from pathlib import Path
from typing import Annotated, Literal

from pydantic import BaseModel, Field

from haiku.rag.utils import get_default_data_dir


class ModelConfig(BaseModel):
    """Configuration for a language model.

    Attributes:
        provider: Model provider (ollama, openai, anthropic, etc.)
        name: Model name/identifier
        base_url: Optional base URL for OpenAI-compatible servers (vLLM, LM Studio, etc.)
        enable_thinking: Control reasoning behavior (true/false/None for default)
        temperature: Sampling temperature (0.0 to 1.0+)
        max_tokens: Maximum tokens to generate
        vision: True if the model can interpret images. Default False.
        extra_body: Raw dict forwarded verbatim to the model SDK as
            `ModelSettings.extra_body`. Provider-side escape hatch for
            keys haiku.rag doesn't model explicitly (e.g. vLLM's
            `chat_template_kwargs.enable_thinking: false` for Qwen3).
            Honored by openai/ollama/anthropic/groq; ignored by gemini/bedrock.
    """

    provider: str = "ollama"
    name: str = "gpt-oss"
    base_url: str | None = None

    enable_thinking: bool | None = None
    temperature: float | None = None
    max_tokens: int | None = None
    vision: bool = False
    extra_body: dict | None = None


class EmbeddingModelConfig(BaseModel):
    """Configuration for an embedding model.

    Attributes:
        provider: Model provider (ollama, openai, voyageai, cohere, sentence-transformers)
        name: Model name/identifier
        vector_dim: Vector dimensions produced by the model
        base_url: Optional base URL for OpenAI-compatible servers (vLLM, LM Studio, etc.)
    """

    provider: str = "ollama"
    name: str = "qwen3-embedding:4b"
    vector_dim: int = 2560
    base_url: str | None = None


class StorageConfig(BaseModel):
    data_dir: Path = Field(default_factory=get_default_data_dir)
    auto_vacuum: bool = True
    vacuum_retention_seconds: int = 86400


class LanceDBConfig(BaseModel):
    uri: str = ""
    api_key: str = ""
    region: str = ""
    storage_options: dict[str, str] = Field(default_factory=dict)


class EmbeddingsConfig(BaseModel):
    model: EmbeddingModelConfig = Field(default_factory=EmbeddingModelConfig)
    batch_size: int = 512


class RerankingConfig(BaseModel):
    model: ModelConfig | None = None


class QAConfig(BaseModel):
    model: ModelConfig = Field(
        default_factory=lambda: ModelConfig(
            provider="ollama",
            name="gpt-oss",
            enable_thinking=True,
            temperature=0.3,
        )
    )
    max_searches: int = 5


class AnalysisConfig(BaseModel):
    """Driving model + sandbox limits for the analysis skill.

    ``model`` defaults to ``None``, meaning "no override — use ``qa.model``."
    Consumers resolve via ``config.analysis.model or config.qa.model``. Set
    explicitly when the analysis workload wants a different model from QA
    (e.g. a stronger model for computational tasks)."""

    model: ModelConfig | None = None
    code_timeout: float = 60.0
    max_output_chars: int = 50_000


class PictureDescriptionConfig(BaseModel):
    """How the VLM runs over each picture when it runs at all.

    Activation lives on ``ProcessingConfig.pictures`` — these fields only
    describe *how* the VLM runs once ``pictures == "description"``.
    """

    model: ModelConfig = Field(
        default_factory=lambda: ModelConfig(
            provider="ollama",
            name="ministral-3",
            temperature=0.0,
        )
    )
    timeout: int = 90
    max_tokens: int = 200


class ConversionOptions(BaseModel):
    """Options for document conversion."""

    # OCR options
    do_ocr: bool = True
    force_ocr: bool = False
    ocr_engine: Literal[
        "auto", "easyocr", "ocrmac", "rapidocr", "tesserocr", "tesseract"
    ] = "auto"
    ocr_lang: list[str] = []

    # Table options
    do_table_structure: bool = True
    table_mode: Literal["fast", "accurate"] = "accurate"
    table_cell_matching: bool = True

    # Image options
    images_scale: float = 2.0
    generate_page_images: bool = True

    # Fetch images referenced by URL in HTML and Markdown inputs.
    # docling-local only — docling-serve cannot fetch external images.
    fetch_remote_images: bool = True

    picture_description: PictureDescriptionConfig = Field(
        default_factory=PictureDescriptionConfig
    )


PicturesMode = Literal["none", "description", "image"]


class ProcessingConfig(BaseModel):
    chunk_size: int = 256
    converter: str = "docling-local"
    chunker: str = "docling-local"
    chunker_type: str = "hybrid"
    chunking_tokenizer: str = "Qwen/Qwen3-Embedding-0.6B"
    chunking_merge_peers: bool = True
    chunking_use_markdown_tables: bool = False
    conversion_options: ConversionOptions = Field(default_factory=ConversionOptions)
    split_pages: int = Field(
        default=0,
        ge=0,
        description=(
            "If >0, PDFs are split into N-page slices, each converted "
            "independently, and merged via DoclingDocument.concatenate. 0 "
            "disables splitting (single-pass conversion). Recommended: 10 "
            "for memory-bound or large (>100 page) PDFs."
        ),
    )
    pictures: PicturesMode = "image"
    """How embedded pictures are handled at ingest.

    - ``"none"``: docling skips picture-image generation; ``label="picture"``
      rows still exist as structure but carry no bytes or description. Use
      this when you don't need picture content and want to keep RAM and DB
      size low on large documents.
    - ``"description"``: docling generates picture images, the configured
      VLM produces text descriptions woven into chunk text, AND the bytes
      are retained in ``document_items.picture_data`` so a vision-capable
      QA model or multimodal embedder can be enabled later without
      reingesting.
    - ``"image"``: docling generates picture images and stores them in
      ``document_items.picture_data``; no VLM runs at ingest.
    """
    auto_title: bool = False
    title_model: ModelConfig = Field(
        default_factory=lambda: ModelConfig(
            provider="ollama",
            name="gpt-oss",
            enable_thinking=False,
            temperature=0.3,
            max_tokens=100,
        )
    )


class SearchConfig(BaseModel):
    limit: int = 5
    max_context_chars: int = 10000
    vector_index_metric: Literal["cosine", "l2", "dot"] = "cosine"
    vector_refine_factor: int = 30


class OllamaConfig(BaseModel):
    base_url: str = Field(
        default_factory=lambda: __import__("os").environ.get(
            "OLLAMA_BASE_URL", "http://localhost:11434"
        )
    )


class DoclingServeConfig(BaseModel):
    """docling-serve endpoints. Accepts a single URL or a list — when a list
    is given, the client round-robins jobs across the URLs. Each job's
    submit/poll/result trio stays on the same instance (task IDs are
    instance-local). The round-robin counter is per-process; for true load
    balancing or failover, put an LB in front."""

    base_url: str | list[str] = "http://localhost:5001"
    api_key: str = ""

    @property
    def base_urls(self) -> list[str]:
        """Always-a-list view of base_url. Empty input falls back to localhost."""
        if isinstance(self.base_url, str):
            return [self.base_url]
        return list(self.base_url) or ["http://localhost:5001"]


class ProvidersConfig(BaseModel):
    ollama: OllamaConfig = Field(default_factory=OllamaConfig)
    docling_serve: DoclingServeConfig = Field(default_factory=DoclingServeConfig)


class PromptsConfig(BaseModel):
    domain_preamble: str = ""
    picture_description: str = (
        "Describe this image for a blind user. "
        "State the image type (screenshot, chart, photo, etc.), "
        "what it depicts, any visible text, and key visual details. "
        "Be concise and accurate."
    )


class EvaluationsConfig(BaseModel):
    """Settings consumed only by the `evaluations` package."""

    judge: ModelConfig | None = Field(
        default=None,
        description=(
            "Judge model for `evaluations run`'s LLM-as-judge step. "
            "ModelConfig's base_url lets the judge point at any "
            "OpenAI-compatible endpoint."
        ),
    )


class QueueConfig(BaseModel):
    """SQLite queue for the production ingester."""

    path: Path = Field(
        default_factory=lambda: get_default_data_dir() / "ingester.db",
        description="Location of the ingester's SQLite queue file.",
    )


class RetryPolicyConfig(BaseModel):
    """Per-job retry policy. Per-source override is allowed under
    SourceConfig.retry so a flaky source doesn't drag the rest of the queue."""

    max_attempts: int = 5
    base_delay_s: float = 2.0
    max_delay_s: float = 300.0
    jitter: float = Field(default=0.25, ge=0.0, le=1.0)


class CircuitBreakerConfig(BaseModel):
    """Per-source breaker over discover() failures. Stops the ingester from
    hammering a source that's persistently failing."""

    failure_threshold: int = Field(
        default=5, description="Consecutive failures before the breaker opens."
    )
    cooldown_s: float = Field(
        default=600.0,
        description="How long the breaker stays open before allowing a probe.",
    )


class WorkerConfig(BaseModel):
    worker_count: int = 4
    max_concurrent: int = 4
    poll_idle_interval_s: float = 1.0
    claim_timeout_s: int = 1800
    reaper_interval_s: int = 60
    retry: RetryPolicyConfig = Field(default_factory=RetryPolicyConfig)
    shutdown_grace_s: float = Field(
        default=60.0,
        description="On SIGINT/SIGTERM, how long to wait for in-flight jobs to "
        "finish before forcing cancellation. Cancelled jobs stay 'claimed' in "
        "the queue; the reaper resets them after claim_timeout_s.",
    )


class APIConfig(BaseModel):
    """HTTP control plane settings for the ingester."""

    enabled: bool = True
    host: str = "127.0.0.1"
    port: int = 8765
    auth_token: str | None = None


class _SourceBase(BaseModel):
    """Fields common to every source. `id` is optional; if omitted the source
    derives a deterministic id from its target (root path / bucket+prefix /
    user-supplied tag)."""

    id: str | None = None
    delete_orphans: bool = True
    poll_interval_s: float = Field(
        default=300.0,
        description="How often discover() runs. FS additionally uses watchfiles "
        "for push events between sweeps.",
    )
    retry: RetryPolicyConfig | None = Field(
        default=None,
        description="Override the worker's default retry policy for jobs from "
        "this source. None = inherit from WorkerConfig.retry.",
    )
    circuit_breaker: CircuitBreakerConfig = Field(default_factory=CircuitBreakerConfig)


class FSSourceConfig(_SourceBase):
    type: Literal["fs"]
    root: Path
    ignore_patterns: list[str] = []
    include_patterns: list[str] = []


class HTTPSourceConfig(_SourceBase):
    type: Literal["http"]
    urls: list[str] = []
    headers: dict[str, str] = Field(default_factory=dict)


class S3SourceConfig(_SourceBase):
    type: Literal["s3"]
    uri: str
    storage_options: dict[str, str] = Field(default_factory=dict)
    ignore_patterns: list[str] = []
    include_patterns: list[str] = []


class WebDAVSourceConfig(_SourceBase):
    """A WebDAV collection (Nextcloud, ownCloud, Apache mod_dav, etc.). Files
    are discovered via PROPFIND on `base_url`; fetch is plain HTTP GET."""

    type: Literal["webdav"]
    base_url: str
    username: str | None = None
    password: str | None = None
    headers: dict[str, str] = Field(default_factory=dict)
    ignore_patterns: list[str] = []
    include_patterns: list[str] = []


SourceConfig = Annotated[
    FSSourceConfig | HTTPSourceConfig | S3SourceConfig | WebDAVSourceConfig,
    Field(discriminator="type"),
]


class IngesterConfig(BaseModel):
    """Production ingester settings."""

    sources: list[SourceConfig] = []
    queue: QueueConfig = Field(default_factory=QueueConfig)
    workers: WorkerConfig = Field(default_factory=WorkerConfig)
    api: APIConfig = Field(default_factory=APIConfig)


class AppConfig(BaseModel):
    environment: str = "production"
    storage: StorageConfig = Field(default_factory=StorageConfig)
    lancedb: LanceDBConfig = Field(default_factory=LanceDBConfig)
    embeddings: EmbeddingsConfig = Field(default_factory=EmbeddingsConfig)
    reranking: RerankingConfig = Field(default_factory=RerankingConfig)
    qa: QAConfig = Field(default_factory=QAConfig)
    analysis: AnalysisConfig = Field(default_factory=AnalysisConfig)
    processing: ProcessingConfig = Field(default_factory=ProcessingConfig)
    search: SearchConfig = Field(default_factory=SearchConfig)
    providers: ProvidersConfig = Field(default_factory=ProvidersConfig)
    prompts: PromptsConfig = Field(default_factory=PromptsConfig)
    ingester: IngesterConfig = Field(default_factory=IngesterConfig)
    evaluations: "EvaluationsConfig" = Field(
        default_factory=lambda: EvaluationsConfig()
    )