from __future__ import annotations

from enum import Enum
from typing import Literal

from pydantic import Field

from app.models.common import StrictModel
from app.models.internal_retrieve import SourceType


class QueryIntent(str, Enum):
    CODE_LOOKUP = "code_lookup"
    DOCUMENTATION = "documentation"
    HISTORICAL = "historical"
    ARCHITECTURE = "architecture"
    GENERAL = "general"
    MIXED = "mixed"


class IntentConfig(StrictModel):
    source_types: list[SourceType] | None = None
    top_k: int = 20
    needs_graph_expansion: bool = False
    include_commits: bool = False


class IntentResult(StrictModel):
    type: QueryIntent
    confidence: float = Field(ge=0.0, le=1.0)
    source_types_used: list[SourceType] = Field(default_factory=list)
    rewritten_query: str | None = None
    tier: Literal["tier1", "tier2"] = "tier1"
    needs_graph_expansion: bool | None = None

    @property
    def config(self) -> IntentConfig:
        base = INTENT_CONFIGS[self.type]
        source_types = self.source_types_used or base.source_types
        graph_expand = (
            self.needs_graph_expansion
            if self.needs_graph_expansion is not None
            else base.needs_graph_expansion
        )
        return IntentConfig(
            source_types=source_types if source_types else None,
            top_k=base.top_k,
            needs_graph_expansion=graph_expand,
            include_commits=base.include_commits,
        )


INTENT_CONFIGS: dict[QueryIntent, IntentConfig] = {
    QueryIntent.CODE_LOOKUP: IntentConfig(
        source_types=[SourceType.CODE],
        top_k=25,
        needs_graph_expansion=True,
        include_commits=False,
    ),
    QueryIntent.DOCUMENTATION: IntentConfig(
        source_types=[SourceType.DOCS],
        top_k=20,
        needs_graph_expansion=False,
        include_commits=False,
    ),
    QueryIntent.HISTORICAL: IntentConfig(
        source_types=[SourceType.COMMIT, SourceType.CODE],
        top_k=15,
        needs_graph_expansion=False,
        include_commits=True,
    ),
    QueryIntent.ARCHITECTURE: IntentConfig(
        source_types=[SourceType.DOCS, SourceType.CODE],
        top_k=40,
        needs_graph_expansion=True,
        include_commits=False,
    ),
    QueryIntent.GENERAL: IntentConfig(
        source_types=None,
        top_k=20,
        needs_graph_expansion=False,
        include_commits=False,
    ),
    QueryIntent.MIXED: IntentConfig(
        source_types=None,
        top_k=30,
        needs_graph_expansion=True,
        include_commits=True,
    ),
}
