src/analyzer/models.py

"""
Pydantic models for Grant Analyst.

Strict schemas for grants, search requests, responses, and QA interactions.
All core entities use validated Pydantic models instead of loose dicts.
"""
from __future__ import annotations

from typing import List, Optional, Dict, Any
from datetime import date, datetime
from pydantic import BaseModel, Field, HttpUrl, field_validator
from enum import Enum


class GrantStatus(str, Enum):
    """Grant/competition status."""
    OPEN = "open"
    UPCOMING = "upcoming"
    CLOSED = "closed"
    UNKNOWN = "unknown"


class FundingRates(BaseModel):
    """Funding rates by organization size."""
    micro_small: Optional[int] = Field(None, description="Funding % for micro/small businesses")
    medium: Optional[int] = Field(None, description="Funding % for medium businesses")
    large: Optional[int] = Field(None, description="Funding % for large businesses")


class FundingInfo(BaseModel):
    """Grant funding information."""
    min: Optional[float] = Field(None, description="Minimum funding amount")
    max: Optional[float] = Field(None, description="Maximum funding amount")
    total_pot: Optional[float] = Field(None, description="Total competition pot")
    rates: Optional[FundingRates] = Field(None, description="Funding rates by org size")


class GrantUrls(BaseModel):
    """URLs associated with a grant."""
    overview: Optional[str] = Field(None, description="Overview page URL")
    apply: Optional[str] = Field(None, description="Application page URL")
    guidance: Optional[str] = Field(None, description="Guidance page URL")
    other: List[str] = Field(default_factory=list, description="Other related URLs")


class Grant(BaseModel):
    """
    Validated grant/competition model.

    Core fields extracted from Innovate UK competition snapshots.
    """
    # Identifiers
    id: str = Field(..., description="Internal grant/competition ID")
    source: str = Field(default="innovate_uk", description="Data source (e.g., 'innovate_uk')")

    # Core content
    title: str = Field(..., description="Grant title")
    programme: Optional[str] = Field(None, description="Programme/scheme name")

    # Descriptive content (raw HTML/text from snapshots)
    summary: Optional[str] = Field(None, description="Summary/description")
    scope: Optional[str] = Field(None, description="Scope and eligibility criteria")
    eligibility: Optional[str] = Field(None, description="Detailed eligibility requirements")

    # Dates
    open_date: Optional[date] = Field(None, description="Competition opening date")
    close_date: Optional[datetime] = Field(None, description="Application deadline (may include time)")
    notify_date: Optional[date] = Field(None, description="Notification date")
    project_start_from: Optional[date] = Field(None, description="Earliest project start date")

    # Funding
    funding: Optional[FundingInfo] = Field(None, description="Funding information")

    # Status (derived from dates)
    status: GrantStatus = Field(default=GrantStatus.UNKNOWN, description="Current status")

    # URLs
    url: Optional[str] = Field(None, description="Primary URL")
    urls: Optional[GrantUrls] = Field(None, description="Structured URLs")

    # Metadata
    created_at: Optional[datetime] = Field(None, description="Record creation timestamp")
    updated_at: Optional[datetime] = Field(None, description="Last update timestamp")

    # Raw sections (for full-text search and citation)
    sections: Optional[Dict[str, Any]] = Field(None, description="Raw HTML sections from source")

    @field_validator('status', mode='before')
    @classmethod
    def derive_status(cls, v: Any, info) -> GrantStatus:
        """Derive status from dates if not explicitly set."""
        if isinstance(v, GrantStatus):
            return v
        if isinstance(v, str):
            try:
                return GrantStatus(v.lower())
            except ValueError:
                pass

        # Derive from dates if available
        data = info.data
        close_date = data.get('close_date')
        open_date = data.get('open_date')

        if close_date:
            now = datetime.now().date()
            close_dt = close_date.date() if isinstance(close_date, datetime) else close_date
            if isinstance(close_dt, date) and close_dt < now:
                return GrantStatus.CLOSED

        if open_date:
            now = datetime.now().date()
            if isinstance(open_date, date) and open_date > now:
                return GrantStatus.UPCOMING

        if open_date and close_date:
            return GrantStatus.OPEN

        return GrantStatus.UNKNOWN

    class Config:
        use_enum_values = True


class SearchFilters(BaseModel):
    """Filters for grant search queries."""
    sources: Optional[List[str]] = Field(None, description="Filter by data sources")
    status: Optional[List[GrantStatus]] = Field(None, description="Filter by status (open/upcoming/closed)")
    min_funding: Optional[float] = Field(None, description="Minimum funding amount")
    max_funding: Optional[float] = Field(None, description="Maximum funding amount")
    location: Optional[List[str]] = Field(None, description="Filter by location/region")
    org_type: Optional[List[str]] = Field(None, description="Filter by organization type")
    theme_keywords: Optional[List[str]] = Field(None, description="Filter by theme keywords")


class SearchHit(BaseModel):
    """A single search result with score."""
    grant: Grant = Field(..., description="The matched grant")
    score: float = Field(..., description="Relevance score (0-1)")

    class Config:
        arbitrary_types_allowed = True


class QARequest(BaseModel):
    """Request model for QA queries."""
    query: str = Field(..., description="User's question about grants", max_length=4000)
    filters: Optional[SearchFilters] = Field(None, description="Optional search filters")
    org_profile: Optional[Dict[str, Any]] = Field(None, description="Optional organization profile for context")
    session_id: Optional[str] = Field(None, description="Optional session ID for tracking")
    use_llm_routing: bool = Field(default=True, description="Use LLM-based query routing")

    @field_validator('query')
    @classmethod
    def validate_query_length(cls, v: str) -> str:
        """Ensure query is not empty and within limits."""
        v = v.strip()
        if not v:
            raise ValueError("Query cannot be empty")
        # Additional check beyond Pydantic's max_length for settings integration
        from analyzer.config import get_settings
        settings = get_settings()
        if len(v) > settings.MAX_QUERY_CHARS:
            raise ValueError(f"Query exceeds maximum length of {settings.MAX_QUERY_CHARS} characters")
        return v


class ChunkType(str, Enum):
    """Types of chunks in streaming responses."""
    METADATA = "metadata"
    INTENT = "intent"
    TOKEN = "token"
    CITATIONS = "citations"
    DONE = "done"
    ERROR = "error"


class QAChunk(BaseModel):
    """
    A chunk in the streaming QA response.

    Supports Server-Sent Events (SSE) and NDJSON streaming.
    """
    type: ChunkType = Field(..., description="Chunk type")

    # Type-specific fields (only one should be populated based on type)
    content: Optional[str] = Field(None, description="Text content (for TOKEN chunks)")
    token: Optional[str] = Field(None, description="Alias for content (for TOKEN chunks)")
    citations: Optional[List[Dict[str, Any]]] = Field(None, description="Citation list (for CITATIONS chunks)")
    error: Optional[str] = Field(None, description="Error message (for ERROR chunks)")

    # Metadata fields
    session_id: Optional[str] = Field(None, description="Session ID (for METADATA chunks)")
    query: Optional[str] = Field(None, description="Original query (for METADATA chunks)")
    intent: Optional[str] = Field(None, description="Detected intent (for INTENT chunks)")
    latency_ms: Optional[int] = Field(None, description="Total latency (for DONE chunks)")

    class Config:
        use_enum_values = True


class CitationInfo(BaseModel):
    """Citation information in response."""
    grant_id: str = Field(..., description="Grant/competition ID")
    title: str = Field(..., description="Grant title")
    url: Optional[str] = Field(None, description="Grant URL")
    confidence: Optional[float] = Field(None, description="Citation confidence score (0-1)")


class QAResponse(BaseModel):
    """Complete QA response (non-streaming)."""
    session_id: str = Field(..., description="Session ID")
    query: str = Field(..., description="Original query")
    answer: str = Field(..., description="Generated answer")
    citations: List[CitationInfo] = Field(default_factory=list, description="Citations used")
    latency_ms: int = Field(..., description="Total latency in milliseconds")
    success: bool = Field(default=True, description="Whether request succeeded")
    error: Optional[str] = Field(None, description="Error message if failed")
    intent: Optional[str] = Field(None, description="Detected query intent")


# Legacy support: simple dict-to-Grant conversion
def dict_to_grant(data: Dict[str, Any]) -> Grant:
    """
    Convert a dictionary (from data loader) to a validated Grant model.

    Handles various field name variations and missing fields gracefully.
    """
    # Extract funding info if present
    funding_data = data.get("funding", {})
    funding = None
    if funding_data and isinstance(funding_data, dict):
        rates_data = funding_data.get("rates", {})
        rates = FundingRates(**rates_data) if rates_data else None
        funding = FundingInfo(
            min=funding_data.get("min"),
            max=funding_data.get("max"),
            total_pot=funding_data.get("total_pot"),
            rates=rates
        )

    # Extract URLs
    url = data.get("url")
    urls = None
    if url:
        urls = GrantUrls(overview=url, apply=None, guidance=None, other=[])

    # Parse dates
    def parse_date(val: Any) -> Optional[date]:
        if isinstance(val, date):
            return val
        if isinstance(val, datetime):
            return val.date()
        if isinstance(val, str):
            try:
                return datetime.fromisoformat(val.replace('Z', '+00:00')).date()
            except:
                return None
        return None

    def parse_datetime(val: Any) -> Optional[datetime]:
        if isinstance(val, datetime):
            return val
        if isinstance(val, str):
            try:
                return datetime.fromisoformat(val.replace('Z', '+00:00'))
            except:
                # Try parsing as date only
                dt = parse_date(val)
                return datetime.combine(dt, datetime.min.time()) if dt else None
        return None

    # Extract summary from various sources
    summary = (
        data.get("summary") or
        data.get("sections", {}).get("summary_raw") or
        data.get("public_description") or
        ""
    )

    # Extract scope and eligibility
    sections = data.get("sections", {})
    scope = sections.get("scope_raw")
    eligibility = sections.get("eligibility_raw")

    return Grant(
        id=data.get("id", "unknown"),
        source=data.get("source", "innovate_uk"),
        title=data.get("title", "Untitled"),
        programme=data.get("programme"),
        summary=summary,
        scope=scope,
        eligibility=eligibility,
        open_date=parse_date(data.get("open_date")),
        close_date=parse_datetime(data.get("close_date")),
        notify_date=parse_date(data.get("notify_date")),
        project_start_from=parse_date(data.get("project_start_from")),
        funding=funding,
        url=url,
        urls=urls,
        created_at=parse_datetime(data.get("created_at")),
        updated_at=parse_datetime(data.get("updated_at")),
        sections=sections if sections else None
    )