api/schemas.py

"""
Request/response schemas for PsyAdGenesis API.
Keeps all Pydantic models in one place for consistency and reuse.
"""

from pydantic import BaseModel, Field
from typing import Optional, List, Literal, Any, Dict


# ----- Generate -----
class GenerateRequest(BaseModel):
    """Request schema for ad generation."""
    niche: Literal["home_insurance", "glp1", "auto_insurance"] = Field(
        description="Target niche: home_insurance, glp1, or auto_insurance"
    )
    num_images: int = Field(default=1, ge=1, le=10, description="Number of images to generate (1-10)")
    image_model: Optional[str] = Field(default=None, description="Image generation model to use")
    target_audience: Optional[str] = Field(default=None, description="Optional target audience description")
    offer: Optional[str] = Field(default=None, description="Optional offer to run")
    use_trending: bool = Field(default=False, description="Whether to incorporate current trending topics")
    trending_context: Optional[str] = Field(default=None, description="Specific trending context when use_trending=True")


class GenerateBatchRequest(BaseModel):
    """Request schema for batch ad generation."""
    niche: Literal["home_insurance", "glp1", "auto_insurance"] = Field(description="Target niche")
    count: int = Field(default=5, ge=1, le=100, description="Number of ads to generate (1-100)")
    images_per_ad: int = Field(default=1, ge=1, le=3, description="Images per ad (1-3)")
    image_model: Optional[str] = Field(default=None, description="Image generation model to use")
    method: Optional[Literal["standard", "matrix"]] = Field(default=None, description="Generation method or None for mixed")
    target_audience: Optional[str] = Field(default=None, description="Optional target audience")
    offer: Optional[str] = Field(default=None, description="Optional offer to run")


class ImageResult(BaseModel):
    """Image result schema."""
    filename: Optional[str] = None
    filepath: Optional[str] = None
    image_url: Optional[str] = None
    r2_url: Optional[str] = None
    model_used: Optional[str] = None
    seed: Optional[int] = None
    error: Optional[str] = None


class AdMetadata(BaseModel):
    """Metadata about the generation."""
    strategies_used: List[str]
    creative_direction: str
    visual_mood: str
    framework: Optional[str] = None
    camera_angle: Optional[str] = None
    lighting: Optional[str] = None
    composition: Optional[str] = None
    hooks_inspiration: List[str]
    visual_styles: List[str]


class GenerateResponse(BaseModel):
    """Response schema for ad generation."""
    id: str
    niche: str
    created_at: str
    title: Optional[str] = Field(default=None, description="Short punchy ad title (3-5 words)")
    headline: str
    primary_text: str
    description: str
    body_story: str = Field(description="Compelling 8-12 sentence story that hooks emotionally")
    cta: str
    psychological_angle: str
    why_it_works: Optional[str] = None
    images: List[ImageResult]
    metadata: AdMetadata


class BatchResponse(BaseModel):
    """Response schema for batch generation."""
    count: int
    ads: List[GenerateResponse]


# ----- Matrix -----
class MatrixGenerateRequest(BaseModel):
    """Request for angle × concept matrix generation."""
    niche: Literal["home_insurance", "glp1", "auto_insurance"] = Field(description="Target niche")
    angle_key: Optional[str] = Field(default=None, description="Specific angle key (random if not provided)")
    concept_key: Optional[str] = Field(default=None, description="Specific concept key (random if not provided)")
    custom_angle: Optional[str] = Field(default=None, description="Custom angle text when angle_key is 'custom'")
    custom_concept: Optional[str] = Field(default=None, description="Custom concept text when concept_key is 'custom'")
    num_images: int = Field(default=1, ge=1, le=5, description="Number of images to generate")
    image_model: Optional[str] = Field(default=None, description="Image generation model to use")
    target_audience: Optional[str] = Field(default=None, description="Optional target audience")
    offer: Optional[str] = Field(default=None, description="Optional offer to run")
    core_motivator: Optional[str] = Field(default=None, description="Optional motivator to guide generation")


class RefineCustomRequest(BaseModel):
    """Request to refine custom angle or concept text using AI."""
    text: str = Field(description="The raw custom text from user")
    type: Literal["angle", "concept"] = Field(description="Whether this is an angle or concept")
    niche: Literal["home_insurance", "glp1", "auto_insurance"] = Field(description="Target niche for context")
    goal: Optional[str] = Field(default=None, description="Optional user goal or context")


class RefinedAngleResponse(BaseModel):
    """Response for refined angle."""
    key: str = Field(default="custom")
    name: str
    trigger: str
    example: str
    category: str = Field(default="Custom")
    original_text: str


class RefinedConceptResponse(BaseModel):
    """Response for refined concept."""
    key: str = Field(default="custom")
    name: str
    structure: str
    visual: str
    category: str = Field(default="Custom")
    original_text: str


class RefineCustomResponse(BaseModel):
    """Response for refined custom angle or concept."""
    status: str
    type: Literal["angle", "concept"]
    refined: Optional[dict] = None
    error: Optional[str] = None


class MotivatorGenerateRequest(BaseModel):
    """Request to generate motivators from niche + angle + concept."""
    niche: Literal["home_insurance", "glp1", "auto_insurance"] = Field(description="Target niche")
    angle: Dict[str, Any] = Field(description="Angle context: name, trigger, example")
    concept: Dict[str, Any] = Field(description="Concept context: name, structure, visual")
    target_audience: Optional[str] = Field(default=None, description="Optional target audience")
    offer: Optional[str] = Field(default=None, description="Optional offer")
    count: int = Field(default=6, ge=3, le=10, description="Number of motivators to generate")


class MotivatorGenerateResponse(BaseModel):
    """Response with generated motivators."""
    motivators: List[str]


class MatrixBatchRequest(BaseModel):
    """Request for batch matrix generation."""
    niche: Literal["home_insurance", "glp1"] = Field(description="Target niche")
    angle_count: int = Field(default=6, ge=1, le=10, description="Number of angles to test")
    concept_count: int = Field(default=5, ge=1, le=10, description="Number of concepts per angle")
    strategy: Literal["balanced", "top_performers", "diverse"] = Field(default="balanced", description="Selection strategy")


class AngleInfo(BaseModel):
    """Angle information."""
    key: str
    name: str
    trigger: str
    category: str


class ConceptInfo(BaseModel):
    """Concept information."""
    key: str
    name: str
    structure: str
    visual: str
    category: str


class MatrixMetadata(BaseModel):
    """Matrix generation metadata."""
    generation_method: str = "angle_concept_matrix"


class MatrixResult(BaseModel):
    """Result from matrix-based generation."""
    angle: AngleInfo
    concept: ConceptInfo


class MatrixGenerateResponse(BaseModel):
    """Response for matrix-based ad generation."""
    id: str
    niche: str
    created_at: str
    title: Optional[str] = Field(default=None, description="Short punchy ad title")
    headline: str
    primary_text: str
    description: str
    body_story: str = Field(description="Compelling 8-12 sentence story that hooks emotionally")
    cta: str
    psychological_angle: str
    why_it_works: Optional[str] = None
    images: List[ImageResult]
    matrix: MatrixResult
    metadata: MatrixMetadata


class CombinationInfo(BaseModel):
    """Info about a single angle × concept combination."""
    combination_id: str
    angle: AngleInfo
    concept: ConceptInfo
    compatibility_score: float
    prompt_guidance: str


class MatrixSummary(BaseModel):
    """Summary of a testing matrix."""
    total_combinations: int
    unique_angles: int
    unique_concepts: int
    average_compatibility: float
    angles_used: List[str]
    concepts_used: List[str]


class TestingMatrixResponse(BaseModel):
    """Response for testing matrix generation."""
    niche: str
    strategy: str
    summary: MatrixSummary
    combinations: List[CombinationInfo]


# ----- Auth -----
class LoginRequest(BaseModel):
    """Login request."""
    username: str = Field(description="Username")
    password: str = Field(description="Password")


class LoginResponse(BaseModel):
    """Login response."""
    token: str
    username: str
    message: str = "Login successful"


# ----- Correction -----
class ImageCorrectRequest(BaseModel):
    """Request schema for image correction."""
    image_id: str = Field(description="ID of existing ad creative or 'temp-id' for images not in DB")
    image_url: Optional[str] = Field(default=None, description="Optional image URL when image_id='temp-id'")
    user_instructions: Optional[str] = Field(default=None, description="User instructions for correction")
    auto_analyze: bool = Field(default=False, description="Auto-analyze image for issues if no instructions")


class SpellingCorrection(BaseModel):
    """Spelling correction entry."""
    detected: str
    corrected: str
    context: Optional[str] = None


class VisualCorrection(BaseModel):
    """Visual correction entry."""
    issue: str
    suggestion: str
    priority: Optional[str] = None


class CorrectionData(BaseModel):
    """Correction data structure."""
    spelling_corrections: List[SpellingCorrection]
    visual_corrections: List[VisualCorrection]
    corrected_prompt: str


class CorrectedImageResult(BaseModel):
    """Corrected image result."""
    filename: Optional[str] = None
    filepath: Optional[str] = None
    image_url: Optional[str] = None
    r2_url: Optional[str] = None
    model_used: Optional[str] = None
    corrected_prompt: Optional[str] = None


class ImageCorrectResponse(BaseModel):
    """Response schema for image correction."""
    status: str
    analysis: Optional[str] = None
    corrections: Optional[CorrectionData] = None
    corrected_image: Optional[CorrectedImageResult] = None
    error: Optional[str] = None


class ImageRegenerateRequest(BaseModel):
    """Request schema for image regeneration."""
    image_id: str = Field(description="ID of existing ad creative in database")
    image_model: Optional[str] = Field(default=None, description="Image model to use (or original if not provided)")
    preview_only: bool = Field(default=True, description="If True, preview only; user confirms selection later")


class RegeneratedImageResult(BaseModel):
    """Regenerated image result."""
    filename: Optional[str] = None
    filepath: Optional[str] = None
    image_url: Optional[str] = None
    r2_url: Optional[str] = None
    model_used: Optional[str] = None
    prompt_used: Optional[str] = None
    seed_used: Optional[int] = None


class ImageRegenerateResponse(BaseModel):
    """Response schema for image regeneration."""
    status: str
    regenerated_image: Optional[RegeneratedImageResult] = None
    original_image_url: Optional[str] = None
    original_preserved: bool = Field(default=True, description="Whether original image info was preserved")
    is_preview: bool = Field(default=False, description="Whether this is a preview (not yet saved)")
    error: Optional[str] = None


class ImageSelectionRequest(BaseModel):
    """Request schema for confirming image selection after regeneration."""
    image_id: str = Field(description="ID of existing ad creative in database")
    selection: str = Field(description="Which image to keep: 'new' or 'original'")
    new_image_url: Optional[str] = Field(default=None, description="URL of new image (required if selection='new')")
    new_r2_url: Optional[str] = Field(default=None, description="R2 URL of the new image")
    new_filename: Optional[str] = Field(default=None, description="Filename of the new image")
    new_model: Optional[str] = Field(default=None, description="Model used for the new image")
    new_seed: Optional[int] = Field(default=None, description="Seed used for the new image")


# ----- Extensive -----
class ExtensiveGenerateRequest(BaseModel):
    """Request for extensive generation."""
    niche: str = Field(description="Target niche or 'others' with custom_niche")
    custom_niche: Optional[str] = Field(default=None, description="Custom niche when 'others' is selected")
    target_audience: Optional[str] = Field(default=None, description="Optional target audience")
    offer: Optional[str] = Field(default=None, description="Optional offer to run")
    num_images: int = Field(default=1, ge=1, le=3, description="Number of images per strategy (1-3)")
    image_model: Optional[str] = Field(default=None, description="Image generation model to use")
    num_strategies: int = Field(default=5, ge=1, le=10, description="Number of creative strategies (1-10)")
    use_creative_inventor: bool = Field(
        default=True,
        description="If True, invent new angles/concepts/visuals/triggers; if False, use researcher",
    )
    trend_context: Optional[str] = Field(default=None, description="Optional trend or occasion (e.g. New Year, Valentine's)")


class ExtensiveJobResponse(BaseModel):
    """Response when extensive generation is started (202 Accepted)."""
    job_id: str
    message: str = "Extensive generation started. Poll /extensive/status/{job_id} for progress."


class InventOnlyRequest(BaseModel):
    """Request for invent-only (no ad generation)."""
    niche: str = Field(description="Niche (e.g. GLP-1, Home Insurance)")
    target_audience: Optional[str] = Field(default=None, description="Target audience")
    offer: Optional[str] = Field(default=None, description="Offer to run")
    n: int = Field(default=5, ge=1, le=15, description="Number of invented essentials")
    trend_context: Optional[str] = Field(default=None, description="Optional trend or occasion")
    export_as_text: bool = Field(default=False, description="If True, return human-readable text; else JSON")


class InventedEssentialSchema(BaseModel):
    """One invented creative essential (for API response)."""
    psychology_trigger: str
    angles: List[str]
    concepts: List[str]
    visual_directions: List[str]
    hooks: List[str] = []
    visual_styles: List[str] = []
    target_audience: str = ""  # Hyper-specific audience for this essential (AI-decided)


class InventOnlyResponse(BaseModel):
    """Response from invent-only endpoint."""
    essentials: List[InventedEssentialSchema]
    export_text: Optional[str] = None


# ----- Creative (upload / analyze / modify) -----
class CreativeAnalysisData(BaseModel):
    """Structured analysis of a creative."""
    visual_style: str
    color_palette: List[str]
    mood: str
    composition: str
    subject_matter: str
    text_content: Optional[str] = None
    current_angle: Optional[str] = None
    current_concept: Optional[str] = None
    target_audience: Optional[str] = None
    strengths: List[str]
    areas_for_improvement: List[str]


class CreativeAnalyzeRequest(BaseModel):
    """Request for creative analysis."""
    image_url: Optional[str] = Field(default=None, description="URL of the image to analyze (alternative to file upload)")


class CreativeAnalysisResponse(BaseModel):
    """Response for creative analysis."""
    status: str
    analysis: Optional[CreativeAnalysisData] = None
    suggested_angles: Optional[List[str]] = None
    suggested_concepts: Optional[List[str]] = None
    error: Optional[str] = None


class CreativeModifyRequest(BaseModel):
    """Request for creative modification."""
    image_url: str = Field(description="URL of the original image")
    analysis: Optional[Dict[str, Any]] = Field(default=None, description="Previous analysis data (optional)")
    angle: Optional[str] = Field(default=None, description="Angle to apply to the creative")
    concept: Optional[str] = Field(default=None, description="Concept to apply to the creative")
    mode: Literal["modify", "inspired"] = Field(default="modify", description="modify = image-to-image, inspired = new generation")
    image_model: Optional[str] = Field(default=None, description="Image generation model to use")
    user_prompt: Optional[str] = Field(default=None, description="Optional custom user prompt for modification")


class ModifiedImageResult(BaseModel):
    """Result of creative modification."""
    filename: Optional[str] = None
    filepath: Optional[str] = None
    image_url: Optional[str] = None
    r2_url: Optional[str] = None
    model_used: Optional[str] = None
    mode: Optional[str] = None
    applied_angle: Optional[str] = None
    applied_concept: Optional[str] = None


class CreativeModifyResponse(BaseModel):
    """Response for creative modification."""
    status: str
    prompt: Optional[str] = None
    image: Optional[ModifiedImageResult] = None
    error: Optional[str] = None


class FileUploadResponse(BaseModel):
    """Response for file upload."""
    status: str
    image_url: Optional[str] = None
    filename: Optional[str] = None
    error: Optional[str] = None


# ----- Database -----
class AdCreativeDB(BaseModel):
    """Ad creative from database."""
    id: str
    niche: str
    title: Optional[str] = None
    headline: str
    primary_text: Optional[str] = None
    description: Optional[str] = None
    body_story: Optional[str] = None
    cta: Optional[str] = None
    psychological_angle: Optional[str] = None
    why_it_works: Optional[str] = None
    image_url: Optional[str] = None
    image_filename: Optional[str] = None
    image_model: Optional[str] = None
    image_seed: Optional[int] = None
    angle_key: Optional[str] = None
    angle_name: Optional[str] = None
    concept_key: Optional[str] = None
    concept_name: Optional[str] = None
    generation_method: Optional[str] = None
    created_at: Optional[str] = None


class DbStatsResponse(BaseModel):
    """Database statistics response."""
    connected: bool
    total_ads: Optional[int] = None
    by_niche: Optional[Dict[str, int]] = None
    by_method: Optional[Dict[str, int]] = None
    error: Optional[str] = None


class EditAdCopyRequest(BaseModel):
    """Request for editing ad copy."""
    ad_id: str = Field(description="ID of the ad to edit")
    field: Literal["title", "headline", "primary_text", "description", "body_story", "cta"] = Field(description="Field to edit")
    value: str = Field(description="New value (manual) or current value (AI edit)")
    mode: Literal["manual", "ai"] = Field(description="Edit mode: manual or ai")
    user_suggestion: Optional[str] = Field(default=None, description="User suggestion for AI editing (optional)")


# ----- Export -----
class BulkExportRequest(BaseModel):
    """Request schema for bulk export."""
    ad_ids: List[str] = Field(description="List of ad IDs to export", min_length=1, max_length=50)


class BulkExportResponse(BaseModel):
    """Response schema for bulk export (actual response is FileResponse with ZIP)."""
    status: str
    message: str
    filename: str