"""Domain models for structured data extraction."""

from __future__ import annotations

from enum import StrEnum

from pydantic import BaseModel, Field


class FieldType(StrEnum):
    """Supported extraction field types."""

    TEXT = "text"
    NUMBER = "number"
    DATE = "date"
    BOOLEAN = "boolean"
    CURRENCY = "currency"
    ENTITY = "entity"  # person, org, location
    LIST = "list"


_FIELD_TYPE_JSON: dict[FieldType, dict] = {
    FieldType.TEXT: {"type": "string"},
    FieldType.NUMBER: {"type": "number"},
    FieldType.DATE: {"type": "string", "format": "date"},
    FieldType.BOOLEAN: {"type": "boolean"},
    FieldType.CURRENCY: {"type": "string", "description": "Monetary amount with currency"},
    FieldType.ENTITY: {"type": "string", "description": "Named entity (person, org, location)"},
    FieldType.LIST: {"type": "array", "items": {"type": "string"}},
}


class ExtractionField(BaseModel):
    """A single field to extract from text."""

    name: str
    field_type: FieldType
    description: str = ""
    required: bool = True
    examples: list[str] = Field(default_factory=list)


class ExtractionSchema(BaseModel):
    """Schema defining what to extract -- domain-agnostic."""

    name: str
    description: str = ""
    fields: list[ExtractionField]

    def to_json_schema(self) -> dict:
        """Convert to JSON schema for LLM structured output.

        Returns:
            A JSON Schema ``object`` with properties derived from each
            :class:`ExtractionField` in ``self.fields``.
        """
        properties: dict = {}
        required_fields: list[str] = []

        for field in self.fields:
            prop = dict(_FIELD_TYPE_JSON.get(field.field_type, {"type": "string"}))
            if field.description:
                prop["description"] = field.description
            if field.examples:
                prop["examples"] = field.examples
            properties[field.name] = prop
            if field.required:
                required_fields.append(field.name)

        schema: dict = {
            "type": "object",
            "properties": properties,
        }
        if required_fields:
            schema["required"] = required_fields
        if self.description:
            schema["description"] = self.description

        return schema


class ExtractionResult(BaseModel):
    """Result of a structured extraction."""

    schema_name: str
    data: dict  # extracted key-value pairs matching the schema
    confidence: float = Field(default=1.0, ge=0, le=1)
    model_used: str = ""
    tokens_used: int = 0
    cost_usd: float = 0.0
    processing_time_ms: int = 0
    warnings: list[str] = Field(default_factory=list)


class ExtractionConfig(BaseModel):
    """Configuration for the extraction service."""

    default_model: str = "claude-haiku-3-20250414"  # cheap model for extraction
    temperature: float = 0.1
    max_tokens: int = 4096
    max_retries: int = 2
