from pydantic import BaseModel, Field
from typing import Optional, Any
from datetime import datetime


# ──────────────────────────────────────────────
# Response Models
# ──────────────────────────────────────────────

class EmailStatusEntry(BaseModel):
    """Tracking status and timestamp for a single email."""
    status: str = "pending"
    sent_at: Optional[datetime] = None


class EmailStatusMap(BaseModel):
    """Map of statuses for all emails in a batch."""
    main: EmailStatusEntry = Field(default_factory=EmailStatusEntry)
    follow_up_1: EmailStatusEntry = Field(default_factory=EmailStatusEntry)
    follow_up_2: EmailStatusEntry = Field(default_factory=EmailStatusEntry)
    follow_up_3: EmailStatusEntry = Field(default_factory=EmailStatusEntry)


class EmailBatchResponse(BaseModel):
    """Single email document returned from MongoDB."""
    id: str = Field(..., alias="_id", description="MongoDB document ID")
    run_id: str
    job_id: Optional[str] = None
    place_id: Optional[str] = None
    lead_email: str
    title: str
    generated_emails: dict[str, Any]
    approval_status: str = "pending"
    email_status: EmailStatusMap = Field(default_factory=EmailStatusMap)
    batch_version: int = 1
    metadata: dict[str, Any] = {}
    created_at: Optional[datetime] = None
    updated_at: Optional[datetime] = None

    class Config:
        populate_by_name = True


class EmailBatchListResponse(BaseModel):
    """Paginated list of email documents."""
    total: int
    page: int
    limit: int
    items: list[EmailBatchResponse]


class RunIdSummary(BaseModel):
    """Summary of a single run_id with email count."""
    run_id: str
    email_count: int
    latest_created_at: Optional[datetime] = None


class RunIdListResponse(BaseModel):
    """List of all distinct run_ids with counts."""
    total_runs: int
    runs: list[RunIdSummary]


# ──────────────────────────────────────────────
# Request Models
# ──────────────────────────────────────────────

class EmailPatchRequest(BaseModel):
    """Partial update for an email document. Only provided fields are applied."""
    approval_status: Optional[str] = Field(None, description="e.g. pending, approved, declined")
    generated_emails: Optional[dict[str, Any]] = Field(None, description="Replace the generated email content")
    email_status: Optional[dict[str, Any]] = Field(None, description="Partial update of email tracking statuses")
    metadata: Optional[dict[str, Any]] = Field(None, description="Replace metadata block")


class RegenerateRequest(BaseModel):
    """Optional overrides when regenerating emails for a run."""
    rag_context: Optional[str] = Field(None, description="Custom RAG context to inject into prompts")

class RegenerateResponseAPI(BaseModel):
    """Summary of a regeneration operation."""
    status: str
    run_id: str
    total_leads_queued: int  
    message: str

class RegenerateResponse(BaseModel):
    """Summary of a regeneration operation."""
    run_id: str
    total_leads: int
    succeeded: int
    failed: int
    skipped: int  # leads with no email address


class DeleteResponse(BaseModel):
    """Result of a delete operation."""
    deleted_count: int
    message: str


class BatchActionResponse(BaseModel):
    """Result of a batch operation (approve/decline)."""
    modified_count: int
    message: str


# ──────────────────────────────────────────────
# Lead Models (PostgreSQL)
# ──────────────────────────────────────────────

class LeadResponse(BaseModel):
    """Deeply enriched lead data from PostgreSQL."""
    id: int
    run_id: str
    job_id: Optional[str] = None
    title: Optional[str] = None
    website: Optional[str] = None
    email: Optional[str] = None
    whatsapp_number: Optional[str] = None
    additional_phone_numbers: Optional[Any] = None
    email_generation_status: Optional[str] = None
    is_email_generated: Optional[bool] = None
    social_links: Optional[Any] = None
    campaign_name: Optional[str] = None
    category: Optional[str] = None
    address: Optional[str] = None
    review_count: Optional[int] = None
    review_rating: Optional[float] = None
    latitude: Optional[float] = None
    longitude: Optional[float] = None
    place_id: Optional[str] = None
    business_type: Optional[str] = None
    seo_score: Optional[int] = None
    conversion_score: Optional[int] = None
    lead_score: Optional[int] = None
    tech_stack: Optional[Any] = None
    identified_problems: Optional[Any] = None
    recommended_actions: Optional[Any] = None
    created_at: Optional[datetime] = None
    updated_at: Optional[datetime] = None

    class Config:
        from_attributes = True


class LeadListResponse(BaseModel):
    """Paginated list of enriched leads."""
    total: int
    page: int
    limit: int
    items: list[LeadResponse]


# ──────────────────────────────────────────────
# Prompt Models
# ──────────────────────────────────────────────

class PromptRequest(BaseModel):
    """Model to create or update a custom AI prompt."""
    run_id: Optional[str] = Field(None, description="If provided, prompt is specific to this run")
    system_instruction: str = Field(..., description="The template system instruction")
    user_instruction: Optional[str] = Field(None, description="Custom user/lead instruction (the variable part)")
    rag_context: Optional[str] = Field(None, description="Custom RAG/social proof context")
    is_universal: bool = Field(False, description="If true, becomes the default for the user if no run_id matches")


class PromptResponse(BaseModel):
    """Returned prompt document."""
    id: str = Field(..., alias="_id")
    user_id: str
    run_id: Optional[str] = None
    system_instruction: str
    user_instruction: Optional[str] = None
    rag_context: Optional[str] = None
    is_universal: bool
    created_at: datetime
    updated_at: datetime

    class Config:
        populate_by_name = True


class PromptListResponse(BaseModel):
    """List of all custom prompts for a user."""
    total: int
    items: list[PromptResponse]


class ResolvedPromptResponse(BaseModel):
    """The actually used prompt triplet (resolved from DB or fallback)."""
    system_instruction: str
    user_instruction: str
    rag_context: str