backend/schemas.py

"""
PromptForge v4.0 — Pydantic schemas
Upgrades: batch generation, stats endpoint, search/filter, duplicate,
          tag management, provider model selection, and theme preference.
"""
from __future__ import annotations
from typing import Any, Dict, List, Optional
from datetime import datetime
from enum import Enum
from pydantic import BaseModel, Field, field_validator


# ── Enumerations ─────────────────────────────────────────────────────────────

class OutputFormat(str, Enum):
    text = "text"
    json = "json"
    both = "both"

class AIProvider(str, Enum):
    none        = "none"
    huggingface = "huggingface"
    google      = "google"

class TargetModel(str, Enum):
    google_ai_studio = "google_ai_studio"
    openai           = "openai"
    anthropic        = "anthropic"
    generic          = "generic"

class PersonaType(str, Enum):
    default        = "default"
    senior_dev     = "senior_dev"
    data_scientist = "data_scientist"
    tech_writer    = "tech_writer"
    product_mgr    = "product_mgr"
    security_eng   = "security_eng"
    devops_eng     = "devops_eng"
    ml_engineer    = "ml_engineer"
    custom         = "custom"

class StyleType(str, Enum):
    professional = "professional"
    concise      = "concise"
    detailed     = "detailed"
    beginner     = "beginner"
    formal       = "formal"
    creative     = "creative"

class PromptStatus(str, Enum):
    pending  = "pending"
    approved = "approved"
    exported = "exported"
    archived = "archived"


# ── InstructionSettings ───────────────────────────────────────────────────────

class InstructionSettings(BaseModel):
    settings_id:   str            = Field(..., description="Unique UUID for this setting.")
    title:         str            = Field(..., min_length=2, max_length=120)
    description:   Optional[str] = Field(None, max_length=1000)
    instruction:   str            = Field(..., min_length=5, max_length=8000)
    extra_context: Optional[str] = Field(None, max_length=2000)
    output_format: OutputFormat   = Field(OutputFormat.both)
    persona:       PersonaType    = Field(PersonaType.default)
    custom_persona: Optional[str] = Field(None, max_length=200)
    style:         StyleType      = Field(StyleType.professional)
    constraints:   List[str]      = Field(default_factory=list)
    tags:          List[str]      = Field(default_factory=list)
    provider:      AIProvider     = Field(AIProvider.none)
    provider_model: Optional[str] = Field(None, max_length=120,
        description="Override the default model for the selected provider.")
    enhance:       bool           = Field(False)
    is_favorite:   bool           = Field(False, description="Starred/pinned setting.")
    created_at:    datetime       = Field(default_factory=datetime.utcnow)
    updated_at:    datetime       = Field(default_factory=datetime.utcnow)
    use_count:     int            = Field(0)
    version:       int            = Field(1, description="Edit version number.")

    @field_validator("tags", mode="before")
    @classmethod
    def normalize_tags(cls, v: List[str]) -> List[str]:
        return [t.strip().lower() for t in v if t.strip()]


class InstructionSettingsCreate(BaseModel):
    title:          str            = Field(..., min_length=2, max_length=120)
    description:    Optional[str] = Field(None, max_length=1000)
    instruction:    str            = Field(..., min_length=5, max_length=8000)
    extra_context:  Optional[str] = Field(None, max_length=2000)
    output_format:  OutputFormat   = Field(OutputFormat.both)
    persona:        PersonaType    = Field(PersonaType.default)
    custom_persona: Optional[str] = Field(None, max_length=200)
    style:          StyleType      = Field(StyleType.professional)
    constraints:    List[str]      = Field(default_factory=list)
    tags:           List[str]      = Field(default_factory=list)
    provider:       AIProvider     = Field(AIProvider.none)
    provider_model: Optional[str] = Field(None, max_length=120)
    enhance:        bool           = Field(False)
    is_favorite:    bool           = Field(False)


class InstructionSettingsUpdate(BaseModel):
    title:          Optional[str]        = Field(None, max_length=120)
    description:    Optional[str]        = None
    instruction:    Optional[str]        = Field(None, min_length=5, max_length=8000)
    extra_context:  Optional[str]        = None
    output_format:  Optional[OutputFormat] = None
    persona:        Optional[PersonaType]  = None
    custom_persona: Optional[str]        = None
    style:          Optional[StyleType]   = None
    constraints:    Optional[List[str]]  = None
    tags:           Optional[List[str]]  = None
    provider:       Optional[AIProvider] = None
    provider_model: Optional[str]        = None
    enhance:        Optional[bool]       = None
    is_favorite:    Optional[bool]       = None


class InstructionSettingsList(BaseModel):
    total:     int
    items:     List[InstructionSettings]
    favorites: int = Field(0, description="Count of starred settings.")


# ── Prompt models ─────────────────────────────────────────────────────────────

class GenerateRequest(BaseModel):
    instruction:     str            = Field(..., min_length=5, max_length=8000)
    output_format:   OutputFormat   = Field(OutputFormat.both)
    provider:        AIProvider     = Field(AIProvider.none)
    api_key:         Optional[str] = Field(None)
    provider_model:  Optional[str] = Field(None)
    enhance:         bool           = Field(False)
    extra_context:   Optional[str] = Field(None, max_length=2000)
    persona:         PersonaType    = Field(PersonaType.default)
    custom_persona:  Optional[str] = Field(None, max_length=200)
    style:           StyleType      = Field(StyleType.professional)
    user_constraints: List[str]    = Field(default_factory=list)
    settings_id:     Optional[str] = Field(None)
    target_model:    TargetModel    = Field(TargetModel.google_ai_studio,
        description="Target model for output formatting (does not affect inference).")


class GenerateFromSettingsRequest(BaseModel):
    settings_id: str
    api_key:     Optional[str] = Field(None)


class BatchGenerateRequest(BaseModel):
    """Generate multiple prompts in one call."""
    requests: List[GenerateRequest] = Field(..., min_length=1, max_length=10)


class ApproveRequest(BaseModel):
    prompt_id: str
    edits:     Optional[Dict[str, Any]] = None


class ExportRequest(BaseModel):
    prompt_id:     str
    export_format: OutputFormat = Field(OutputFormat.json)


class RefineRequest(BaseModel):
    prompt_id:  str
    feedback:   str           = Field(..., min_length=3, max_length=2000)
    provider:   AIProvider    = Field(AIProvider.none)
    api_key:    Optional[str] = Field(None)


class SearchRequest(BaseModel):
    query:       str           = Field(..., min_length=1, max_length=500)
    search_in:   List[str]     = Field(default_factory=lambda: ["instruction", "role", "task"])
    status_filter: Optional[PromptStatus] = None


# ── Structured prompt ─────────────────────────────────────────────────────────

class StructuredPrompt(BaseModel):
    role:          str
    task:          str
    input_format:  str
    output_format: str
    constraints:   List[str]                    = Field(default_factory=list)
    style:         str
    safety:        List[str]                    = Field(default_factory=list)
    examples:      Optional[List[Dict[str, str]]] = None
    raw_prompt_text: str
    word_count:    int = Field(0, description="Approximate word count of raw_prompt_text.")


class PromptManifest(BaseModel):
    prompt_id:         str
    version:           int      = 1
    created_at:        datetime = Field(default_factory=datetime.utcnow)
    updated_at:        datetime = Field(default_factory=datetime.utcnow)
    instruction:       str
    status:            str      = "pending"
    structured_prompt: StructuredPrompt
    enhancement_notes: Optional[str] = None
    explanation:       Optional[str] = None
    settings_id:       Optional[str] = None
    persona_used:      PersonaType   = Field(PersonaType.default)
    style_used:        StyleType     = Field(StyleType.professional)
    tags:              List[str]     = Field(default_factory=list)
    is_favorite:       bool          = Field(False)


# ── Response models ───────────────────────────────────────────────────────────

class GenerateResponse(BaseModel):
    success:   bool
    prompt_id: str
    manifest:  PromptManifest
    message:   str = "Manifest generated — awaiting approval."


class BatchGenerateResponse(BaseModel):
    success:   bool
    total:     int
    results:   List[GenerateResponse]
    failed:    int = 0


class ApproveResponse(BaseModel):
    success:          bool
    prompt_id:        str
    message:          str
    finalized_prompt: StructuredPrompt


class ExportResponse(BaseModel):
    success:   bool
    prompt_id: str
    data:      Any


class ExplainResponse(BaseModel):
    prompt_id:     str
    explanation:   str
    key_decisions: List[str] = Field(default_factory=list)


class HistoryEntry(BaseModel):
    prompt_id:   str
    version:     int
    created_at:  datetime
    updated_at:  datetime
    instruction: str
    status:      str
    settings_id: Optional[str] = None
    explanation: Optional[str] = None
    tags:        List[str]     = Field(default_factory=list)
    is_favorite: bool          = False


class HistoryResponse(BaseModel):
    total:   int
    entries: List[HistoryEntry]


class SearchResult(BaseModel):
    prompt_id:   str
    instruction: str
    status:      str
    score:       float
    snippet:     str


class SearchResponse(BaseModel):
    total:   int
    results: List[SearchResult]


class StatsResponse(BaseModel):
    total_prompts:       int
    total_settings:      int
    pending_count:       int
    approved_count:      int
    exported_count:      int
    archived_count:      int
    favorite_prompts:    int
    favorite_settings:   int
    top_personas:        Dict[str, int]
    top_styles:          Dict[str, int]
    total_refinements:   int
    avg_constraints:     float
    uptime_since:        datetime


# ── Env config ────────────────────────────────────────────────────────────────

class EnvConfigStatus(BaseModel):
    hf_key_set:     bool = Field(..., description="True if HF_API_KEY is set.")
    google_key_set: bool = Field(..., description="True if GOOGLE_API_KEY is set.")
    port:           str  = "7860"
    version:        str  = "4.0"
    hf_model:       str  = "mistralai/Mistral-7B-Instruct-v0.2"
    google_model:   str  = "gemini-1.5-flash"