from __future__ import annotations

import importlib
import sys
import unittest
from datetime import UTC, datetime, timedelta
from types import SimpleNamespace
from typing import Any, cast
from unittest.mock import patch

from app.config import config
from app.core.logging import clear_log_context, set_log_context
from app.llm import LLMOutputKind, LLMResolvedRoute, LLMStructuredResult, LLMTask, LLMTextResult
from app.llm.routing import DefaultModelSelectionPolicy
from app.llm.schemas import HistoriaClinicaStructured
from app.services.clinical_structured_extraction import ClinicalStructuredExtractionService
from app.services.llm_task_cache import (
    MongoLLMTaskCacheRepository,
    build_llm_task_fingerprint,
)


class _FakeCollection:
    def __init__(self) -> None:
        self.docs: list[dict] = []
        self.indexes: list[tuple] = []
        self._id_counter = 0

    def create_index(self, keys, **kwargs):
        self.indexes.append((tuple(keys), dict(kwargs)))

    def find_one(self, query, sort=None):
        now = datetime.now(UTC)
        for doc in reversed(self.docs):
            if doc.get("tipo_documento") != query.get("tipo_documento"):
                continue
            if doc.get("usuario") != query.get("usuario"):
                continue
            if doc.get("task") != query.get("task"):
                continue
            if doc.get("fingerprint") != query.get("fingerprint"):
                continue
            expires_at = doc.get("expires_at")
            if isinstance(expires_at, datetime) and expires_at <= now:
                continue
            return dict(doc)
        return None

    def update_one(self, query, update, upsert=False):
        if "$inc" in update:
            for doc in self.docs:
                if doc.get("_id") == query.get("_id"):
                    for key, value in update["$inc"].items():
                        doc[key] = int(doc.get(key, 0)) + int(value)
                    return SimpleNamespace(matched_count=1, modified_count=1, upserted_id=None)
            return SimpleNamespace(matched_count=0, modified_count=0, upserted_id=None)

        matched = None
        for doc in self.docs:
            if (
                doc.get("tipo_documento") == query.get("tipo_documento")
                and doc.get("usuario") == query.get("usuario")
                and doc.get("task") == query.get("task")
                and doc.get("fingerprint") == query.get("fingerprint")
            ):
                matched = doc
                break
        if matched is None and upsert:
            self._id_counter += 1
            matched = {"_id": self._id_counter}
            self.docs.append(matched)

        if matched is not None:
            for key, value in update.get("$setOnInsert", {}).items():
                matched.setdefault(key, value)
            matched.update(update.get("$set", {}))
        return SimpleNamespace(matched_count=1 if matched else 0, modified_count=1 if matched else 0, upserted_id=None)


class _MemoryCacheRepository:
    def __init__(self) -> None:
        self.cache_version = "v1"
        self.items: dict[tuple[str, str, str], dict] = {}

    def get(self, *, username: str, task: str, fingerprint: str):
        return self.items.get((username, task, fingerprint))

    def upsert(self, *, username: str, task: str, fingerprint: str, provider: str, model: str, payload, metadata=None):
        self.items[(username, task, fingerprint)] = {
            "usuario": username,
            "task": task,
            "fingerprint": fingerprint,
            "provider": provider,
            "model": model,
            "payload": payload,
            "metadata": dict(metadata or {}),
        }


class _StructuredRouter:
    def __init__(self, *, fail: bool = False) -> None:
        self.fail = fail
        self.calls = 0

    def _resolve_route(self, task):
        return LLMResolvedRoute(
            task=task,
            output_kind=LLMOutputKind.STRUCTURED_OBJECT,
            provider="gemini",
            model="gemini-2.5-flash-lite",
        )

    def generate_structured(self, request):
        self.calls += 1
        if self.fail:
            raise AssertionError("LLM should not be called on cache hit")
        payload = HistoriaClinicaStructured(
            patient_name="PACIENTE CACHE",
            resumen_clinico="Resumen cacheable",
            diagnosticos=[],
            procedimientos=[],
            medicamentos=[],
        ).model_dump(by_alias=True, exclude_none=True, exclude_defaults=True)
        return LLMStructuredResult(
            content=payload,
            provider="gemini",
            model="gemini-2.5-flash-lite",
            output_kind=request.output_kind,
        )


class _TextRouter:
    def __init__(self) -> None:
        self.calls: list[tuple[str, str]] = []

    def _resolve_route(self, task):
        model = "gemini-2.5-flash-lite" if task == LLMTask.HISTORIA_CHUNK_SUMMARY else "gemini-2.5-flash"
        return LLMResolvedRoute(task=task, output_kind=LLMOutputKind.CONTROLLED_PLAIN_TEXT, provider="gemini", model=model)

    def generate_text(self, request):
        self.calls.append((request.task.value, request.model or ""))
        text = f"respuesta-{request.task.value}-{len(self.calls)}"
        return LLMTextResult(content=text, provider="gemini", model=request.model or "unset")


class LLMTaskFingerprintTest(unittest.TestCase):
    def test_fingerprint_is_stable_for_same_semantics(self) -> None:
        first = build_llm_task_fingerprint(
            task="clinical_document_extract",
            provider="gemini",
            model="gemini-2.5-flash-lite",
            cache_version="v1",
            prompt_version="p1",
            schema_version="v1",
            payload={"text": "Paciente   uno"},
            metadata={"b": 2, "a": 1},
        )
        second = build_llm_task_fingerprint(
            task="clinical_document_extract",
            provider="gemini",
            model="gemini-2.5-flash-lite",
            cache_version="v1",
            prompt_version="p1",
            schema_version="v1",
            payload={"text": "Paciente uno"},
            metadata={"a": 1, "b": 2},
        )

        self.assertEqual(first, second)

    def test_fingerprint_changes_when_cache_dimensions_change(self) -> None:
        base = build_llm_task_fingerprint(
            task="clinical_document_extract",
            provider="gemini",
            model="gemini-2.5-flash-lite",
            cache_version="v1",
            prompt_version="p1",
            schema_version="v1",
            payload={"text": "Paciente uno"},
        )

        changed_model = build_llm_task_fingerprint(
            task="clinical_document_extract",
            provider="gemini",
            model="gemini-2.5-pro",
            cache_version="v1",
            prompt_version="p1",
            schema_version="v1",
            payload={"text": "Paciente uno"},
        )
        changed_cache_version = build_llm_task_fingerprint(
            task="clinical_document_extract",
            provider="gemini",
            model="gemini-2.5-flash-lite",
            cache_version="v2",
            prompt_version="p1",
            schema_version="v1",
            payload={"text": "Paciente uno"},
        )

        self.assertNotEqual(base, changed_model)
        self.assertNotEqual(base, changed_cache_version)


class ClinicalExtractRoutingPolicyTest(unittest.TestCase):
    def test_factura_route_uses_high_risk_extract_model(self) -> None:
        route = DefaultModelSelectionPolicy().resolve(
            LLMTask.CLINICAL_DOCUMENT_EXTRACT,
            metadata={"document_type": "factura", "risk_level": "high"},
        )

        self.assertEqual(route.provider, "gemini")
        self.assertEqual(route.model, config.GEMINI_MODEL_EXTRACT_HIGH_RISK)
        self.assertEqual(route.metadata["risk_level"], "high")

    def test_low_risk_route_keeps_low_risk_model(self) -> None:
        route = DefaultModelSelectionPolicy().resolve(
            LLMTask.CLINICAL_DOCUMENT_EXTRACT,
            metadata={"document_type": "laboratorio", "risk_level": "low"},
        )

        self.assertEqual(route.model, config.GEMINI_MODEL_EXTRACT_LOW_RISK)


class MongoLLMTaskCacheRepositoryTest(unittest.TestCase):
    def test_upsert_and_get_return_active_hit(self) -> None:
        collection = _FakeCollection()
        repository = MongoLLMTaskCacheRepository(SimpleNamespace(collection=collection))
        repository.ensure_indexes()
        repository.upsert(
            username="auditor",
            task="historia_chunk_summary",
            fingerprint="abc",
            provider="gemini",
            model="gemini-2.5-flash-lite",
            payload={"text": "Resumen"},
            metadata={"prompt_version": "v1"},
        )

        cached = repository.get(username="auditor", task="historia_chunk_summary", fingerprint="abc")

        self.assertIsNotNone(cached)
        cached_doc = cast(dict[str, Any], cached)
        self.assertEqual(cached_doc["payload"]["text"], "Resumen")
        self.assertEqual(collection.docs[0]["hit_count"], 1)
        self.assertEqual(len(collection.indexes), 3)
        self.assertIs(repository.collection, collection)

    def test_get_skips_expired_documents(self) -> None:
        collection = _FakeCollection()
        collection.docs.append(
            {
                "_id": 1,
                "tipo_documento": "llm_task_cache",
                "usuario": "auditor",
                "task": "historia_chunk_summary",
                "fingerprint": "abc",
                "payload": {"text": "viejo"},
                "expires_at": datetime.now(UTC) - timedelta(days=1),
            }
        )
        repository = MongoLLMTaskCacheRepository(SimpleNamespace(collection=collection))

        cached = repository.get(username="auditor", task="historia_chunk_summary", fingerprint="abc")

        self.assertIsNone(cached)

    def test_repository_accepts_direct_collection_backend(self) -> None:
        collection = _FakeCollection()
        repository = MongoLLMTaskCacheRepository(collection)

        repository.upsert(
            username="auditor",
            task="clinical_document_extract",
            fingerprint="fp-001",
            provider="gemini",
            model="gemini-2.5-flash-lite",
            payload={"text": "ok"},
        )

        cached = repository.get(username="auditor", task="clinical_document_extract", fingerprint="fp-001")

        self.assertIsNotNone(cached)
        cached_doc = cast(dict[str, Any], cached)
        self.assertEqual(cached_doc["payload"]["text"], "ok")


class ClinicalStructuredExtractionCacheTest(unittest.TestCase):
    def tearDown(self) -> None:
        clear_log_context()

    def test_extract_reuses_cached_structured_result(self) -> None:
        cache_repository = _MemoryCacheRepository()
        set_log_context({"username": "auditor"})

        first_service = ClinicalStructuredExtractionService(
            llm_router=_StructuredRouter(),
            cache_repository=cache_repository,
        )
        first = first_service.extract(raw_text="Historia clínica ejemplo", document_type="historia_clinica")

        second_service = ClinicalStructuredExtractionService(
            llm_router=_StructuredRouter(fail=True),
            cache_repository=cache_repository,
        )
        second = second_service.extract(raw_text="Historia clínica ejemplo", document_type="historia_clinica")

        self.assertEqual(first.analysis_structured, second.analysis_structured)
        self.assertEqual(second.metrics, {"source": "cache"})
        self.assertEqual(second.analysis_model.patient_name, "PACIENTE CACHE")


class HistoriaSummaryCacheTest(unittest.TestCase):
    def test_prepare_historia_prompt_reuses_chunk_and_final_summary_cache(self) -> None:
        cache_repository = _MemoryCacheRepository()
        router = _TextRouter()
        with patch.object(config, "GROQ_API_KEY", "fake-key"):
            module = sys.modules.get("modules.processing.resumen_google")
            if module is None:
                historia_module = importlib.import_module("modules.processing.resumen_google")
            else:
                historia_module = module

        descripcion = "\n".join(
            f"Evolucion dia {idx}: dolor abdominal persistente, signos vitales estables y plan terapeutico {idx}."
            for idx in range(24)
        )

        with patch.object(config, "HISTORIA_COMPACT_MAX_CHARS", 6000):
            first = historia_module._prepare_historia_prompt_text(
                descripcion,
                summary_threshold=300,
                summary_chunk=250,
                max_chars=500,
                llm_router=router,
                cache_repository=cache_repository,
                username="auditor",
            )
            first_call_count = len(router.calls)
            second = historia_module._prepare_historia_prompt_text(
                descripcion,
                summary_threshold=300,
                summary_chunk=250,
                max_chars=500,
                llm_router=router,
                cache_repository=cache_repository,
                username="auditor",
            )

        self.assertIn("version=v7", first)
        self.assertIn("[RESUMEN PREVIO DEL DOCUMENTO COMPLETO]", first)
        self.assertEqual(first, second)
        self.assertGreater(first_call_count, 1)
        self.assertEqual(len(router.calls), first_call_count)
