"""LLM extraction orchestration without giving the model storage authority."""

from __future__ import annotations

from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any, Protocol

from .memory_extraction import MemoryExtractionResult, build_extraction_prompt, parse_extraction_response
from .memory_repository import SqlAlchemyMemoryRepository


class JsonLLMClient(Protocol):
    def chat_json(self, messages: list[dict[str, str]], temperature: float = 0.3, max_tokens: int = 4096) -> dict[str, Any]:
        ...


@dataclass(frozen=True)
class MemoryIngestResult:
    entity_count: int
    edge_count: int
    unresolved_edge_refs: list[str] = field(default_factory=list)


class MemoryExtractionService:
    def __init__(self, client: JsonLLMClient):
        self.client = client

    def extract(
        self,
        *,
        language: str,
        ontology: dict[str, Any],
        episode_text: str,
        context: str = "",
    ) -> MemoryExtractionResult:
        prompt = build_extraction_prompt(
            language=language,
            ontology=ontology,
            episode_text=episode_text,
            context=context,
        )
        raw = self.client.chat_json(
            [{"role": "system", "content": prompt}],
            temperature=0.2,
            max_tokens=8192,
        )
        return parse_extraction_response(raw)

    @staticmethod
    def _ref(value: str) -> str:
        return " ".join(value.casefold().split())

    @staticmethod
    def _timestamp(value: str | None):
        if value is None:
            return None
        if not isinstance(value, str) or not value.strip():
            raise ValueError("invalid_memory_timestamp")
        try:
            parsed = datetime.fromisoformat(value.strip().replace("Z", "+00:00"))
        except ValueError as exc:
            raise ValueError("invalid_memory_timestamp") from exc
        return parsed if parsed.tzinfo is not None else parsed.replace(tzinfo=timezone.utc)

    def persist(
        self,
        repository: SqlAlchemyMemoryRepository,
        result: MemoryExtractionResult,
        *,
        source_type: str,
        source_ref: str,
        episode_text: str,
    ) -> MemoryIngestResult:
        repository.add_episode(
            source_type=source_type,
            source_ref=source_ref,
            normalized_text=episode_text,
            summary=result.episode_summary,
        )

        entity_refs: dict[str, str] = {}
        entity_count = 0
        for entity in result.entities:
            node = repository.upsert_node(
                canonical_name=entity.canonical_name,
                labels=entity.labels,
                aliases=entity.aliases,
                attributes=entity.attributes,
                summary=entity.summary,
                confidence=entity.confidence,
            )
            entity_count += 1
            for reference in [entity.mention, entity.canonical_name, *entity.aliases]:
                entity_refs[self._ref(reference)] = node.id

        unresolved: list[str] = []
        edge_count = 0
        for edge in result.edges:
            source_id = entity_refs.get(self._ref(edge.source_entity_ref))
            target_id = entity_refs.get(self._ref(edge.target_entity_ref))
            if source_id is None or target_id is None:
                unresolved.append(f"{edge.source_entity_ref}->{edge.target_entity_ref}")
                continue
            repository.upsert_edge(
                source_node_id=source_id,
                target_node_id=target_id,
                relation=edge.relation,
                fact=edge.fact,
                attributes=edge.attributes,
                confidence=edge.confidence,
                valid_at=self._timestamp(edge.valid_at),
                invalid_at=self._timestamp(edge.invalid_at),
                expired_at=self._timestamp(edge.expired_at),
            )
            edge_count += 1

        return MemoryIngestResult(
            entity_count=entity_count,
            edge_count=edge_count,
            unresolved_edge_refs=unresolved,
        )