import re
import unicodedata
from enum import StrEnum
from typing import Annotated

from pydantic import BaseModel, ConfigDict, Field, StringConstraints, field_validator, model_validator

NonEmptyStr = Annotated[str, StringConstraints(strip_whitespace=True, min_length=1)]

RESERVED_RELATIONS = frozenset({"MENTIONS", "HAS_CHUNK"})
FALLBACK_RELATION = "RELATED_TO"


def to_ascii(value: str) -> str:
    return unicodedata.normalize("NFKD", value).encode("ascii", "ignore").decode()


def canonical_key(name: str) -> str:
    return re.sub(r"[^a-z0-9]", "", to_ascii(name).lower()) or name.strip().casefold()


class EntityType(StrEnum):
    COMPONENT = "Component"
    CLASS = "Class"
    FUNCTION = "Function"
    LIBRARY = "Library"
    CONCEPT = "Concept"
    FEATURE = "Feature"
    ERROR = "Error"
    DATABASE = "Database"
    API = "API"


ENTITY_TYPES_BY_NAME = {entity_type.value.lower(): entity_type for entity_type in EntityType}


class RetrievalMode(StrEnum):
    VECTOR = "vector"
    GRAPH = "graphrag"


class Chunk(BaseModel):
    model_config = ConfigDict(frozen=True)

    id: str
    document: str
    position: int
    text: str


class RetrievedChunk(Chunk):
    score: float
    entity_keys: list[str] = Field(default_factory=list)


class Entity(BaseModel):
    name: NonEmptyStr
    type: EntityType = EntityType.CONCEPT

    @field_validator("type", mode="before")
    @classmethod
    def coerce_type(cls, value: object) -> EntityType:
        return ENTITY_TYPES_BY_NAME.get(str(value).strip().lower(), EntityType.CONCEPT)

    @property
    def key(self) -> str:
        return canonical_key(self.name)


class Relation(BaseModel):
    source: NonEmptyStr
    relation: str
    target: NonEmptyStr
    evidence: str = ""

    @field_validator("relation", mode="before")
    @classmethod
    def normalize_relation(cls, value: object) -> str:
        normalized = re.sub(r"[^A-Z0-9]+", "_", to_ascii(str(value)).upper()).strip("_")
        if not normalized or not normalized[0].isalpha() or normalized in RESERVED_RELATIONS:
            return FALLBACK_RELATION
        return normalized

    @field_validator("evidence", mode="before")
    @classmethod
    def normalize_evidence(cls, value: object) -> str:
        return " ".join(str(value or "").split())

    @property
    def source_key(self) -> str:
        return canonical_key(self.source)

    @property
    def target_key(self) -> str:
        return canonical_key(self.target)


class GraphExtraction(BaseModel):
    entities: list[Entity] = Field(default_factory=list)
    relations: list[Relation] = Field(default_factory=list)

    @model_validator(mode="after")
    def reconcile(self) -> "GraphExtraction":
        self.relations = [
            relation for relation in self.relations if relation.source_key != relation.target_key
        ]
        entities = {entity.key: entity for entity in self.entities}
        for relation in self.relations:
            entities.setdefault(relation.source_key, Entity(name=relation.source))
            entities.setdefault(relation.target_key, Entity(name=relation.target))
        self.entities = list(entities.values())
        return self


class GraphFact(BaseModel):
    source: str
    relation: str
    target: str
    evidence: list[str] = Field(default_factory=list)
    documents: list[str] = Field(default_factory=list)
    distance: int = 1


class RetrievalContext(BaseModel):
    chunks: list[RetrievedChunk] = Field(default_factory=list)
    facts: list[GraphFact] = Field(default_factory=list)
    seed_entities: list[str] = Field(default_factory=list)

    @property
    def is_empty(self) -> bool:
        return not self.chunks and not self.facts

    @property
    def sources(self) -> list[str]:
        documents = {chunk.document for chunk in self.chunks}
        documents.update(document for fact in self.facts for document in fact.documents)
        return sorted(documents)


class Answer(BaseModel):
    question: str
    text: str
    mode: RetrievalMode
    context: RetrievalContext