from __future__ import annotations

import json
from dataclasses import dataclass
from typing import Any

from app.config import config
from app.core.logging import get_log_context
from app.llm import LLMOutputKind, LLMStructuredRequest, LLMTask
from app.llm.renderers import render_clinical_document_html
from app.llm.schemas import (
    ClinicalStructuredModel,
    dump_clinical_model,
    get_clinical_document_model,
    get_clinical_schema_name,
)
from app.services.clinical_extraction_policy import decide_clinical_extraction_policy
from app.services.clinical_extraction_quality import assess_clinical_extraction
from app.services.llm_task_cache import build_llm_task_fingerprint


_DOCUMENT_INSTRUCTIONS: dict[str, str] = {
    "historia_clinica": (
        "Extrae metadatos clínicos completos, resumen, diagnósticos, procedimientos, antecedentes y medicamentos. "
        "Devuelve solo datos JSON: ningún campo debe contener HTML, Markdown ni etiquetas. "
        "El prestador de servicio es la institución que emite la historia clínica; prioriza el nombre institucional "
        "cercano al título o NIT. Usa ERP, EPS, pagador o aseguradora solo si no aparece una institución emisora. "
        "Nunca uses una identificación del paciente, una secuencia censurada con X ni el número de caso como prestador. "
        "Fecha de ingreso debe incluir hora si aparece en el documento. "
        "Para cualquier fecha con hora usa ISO `YYYY-MM-DDTHH:mm-05:00`; para fecha simple usa `YYYY-MM-DD`. "
        "Motivo de consulta debe salir de la sección clínica de motivo; no uses estado civil ni datos administrativos. "
        "Diagnósticos debe contener enfermedades o diagnósticos explícitos, no síntomas aislados. "
        "Procedimientos debe contener solo procedimientos actuales/tratados del caso. "
        "Extrae cada antecedente en antecedentes_estructurados indicando categoría, descripción, estado presente o "
        "negado, fecha y evidencia breve cuando estén disponibles. Reconstruye valores partidos por columnas o saltos "
        "de línea. No conviertas cada negación en texto libre ni descartes alergias positivas por una negación posterior. "
        "Marca discordante cuando el documento afirme y niegue el mismo antecedente o niegue de forma general una "
        "categoría con datos positivos. Procedimientos históricos deben ir en procedimientos_antecedentes y como "
        "antecedente quirúrgico estructurado; procedimientos actuales pertenecen solo a procedimientos. "
        "Para cada medicamento separa estrictamente nombre, presentación, dosis, posología (vía, frecuencia y "
        "duración), cantidad numérica y estado. Nunca uses administrado, formulado u ordenado como nombre; nunca "
        "pongas un código CUM/IUM en dosis. Clasifica el estado como administrado solo si existe registro de "
        "administración, formulado si existe fórmula explícita, ordenado si solo existe orden, antecedente u otro. "
        "El resumen debe ser estable, de 3 a 5 frases clínicas, sin datos administrativos redundantes. "
        "No inventes códigos. Si un código o fecha no aparece, omítelo o usa texto mínimo claro."
    ),
    "laboratorio": (
        "Extrae nombre del paciente, fecha del examen, tipo de examen, tabla de resultados, valores alterados, "
        "interpretación y observaciones. Si la fecha incluye hora usa ISO `YYYY-MM-DDTHH:mm-05:00`; "
        "si no incluye hora usa `YYYY-MM-DD`."
    ),
    "radiologia": (
        "Extrae nombre del paciente, fecha del estudio, tipo de estudio, región anatómica, técnica, hallazgos y conclusión. "
        "Si la fecha incluye hora usa ISO `YYYY-MM-DDTHH:mm-05:00`; si no incluye hora usa `YYYY-MM-DD`."
    ),
    "prescripcion": (
        "Extrae fecha, medicamentos formulados y observaciones de la prescripción. Si la fecha incluye hora usa "
        "ISO `YYYY-MM-DDTHH:mm-05:00`; si no incluye hora usa `YYYY-MM-DD`."
    ),
    "quirurgico": (
        "Extrae fecha del procedimiento, diagnósticos, procedimientos realizados, hallazgos, descripción del procedimiento, "
        "complicaciones y conclusión. Si la fecha incluye hora usa ISO `YYYY-MM-DDTHH:mm-05:00`; "
        "si no incluye hora usa `YYYY-MM-DD`."
    ),
    "factura": (
        "Extrae una factura clínica electrónica. Separa estrictamente proveedor, datos de factura, pagador, "
        "paciente, líneas facturadas y resumen financiero. En proveedor usa el nombre institucional y el NIT "
        "correctos. En datos de factura toma el número de factura únicamente del bloque Prefijo/No. Factura, "
        "no de resoluciones DIAN. Devuelve fechas solo en formato ISO: `YYYY-MM-DDTHH:mm-05:00` para fecha-hora y "
        "`YYYY-MM-DD` para fecha simple. Si el vencimiento viene relativo, resuélvelo a ISO usando la fecha "
        "de emisión. Extrae Caso No., Convenio, Fecha Ingreso, Fecha Egreso, Vencimiento y el valor en letras "
        "solo si están explícitos. En las líneas facturadas conserva descripciones reales, "
        "códigos, cantidades, valores y clasifica sin inventar entre procedimientos, laboratorio, imagenología, "
        "hospitalización, honorarios, medicamentos y otros servicios."
    ),
    "generico": "Resume el documento y registra referencias diagnósticas relevantes sin inventar datos.",
}
_CLINICAL_EXTRACT_PROMPT_VERSION = "v1"
_DOCUMENT_PROMPT_VERSIONS: dict[str, str] = {
    "historia_clinica": "v5",
    "factura": "v3",
}


@dataclass(slots=True)
class StructuredClinicalAnalysisResult:
    document_type: str
    analysis_model: ClinicalStructuredModel
    analysis_structured: dict[str, Any]
    rendered_html: str
    provider: str
    model: str
    error_analisis: str | None = None
    metrics: dict[str, Any] | None = None
    route_metadata: dict[str, Any] | None = None
    quality_metadata: dict[str, Any] | None = None


class ClinicalStructuredExtractionService:
    def __init__(self, *, llm_router: Any, cache_repository: Any = None) -> None:
        self._llm_router = llm_router
        self._cache_repository = cache_repository

    def extract(
        self,
        *,
        raw_text: str,
        document_type: str,
        summary_text: str | None = None,
    ) -> StructuredClinicalAnalysisResult:
        normalized_type = str(document_type or "generico").strip() or "generico"
        model_cls = get_clinical_document_model(normalized_type)
        schema_name = get_clinical_schema_name(normalized_type)
        schema_version = str(model_cls.model_fields["schema_version"].default)
        prompt_version = self._get_prompt_version(normalized_type)
        policy_decision = decide_clinical_extraction_policy(
            document_type=normalized_type,
            raw_text=raw_text,
            summary_text=summary_text,
        )
        prompt = self._build_prompt(
            document_type=normalized_type,
            raw_text=raw_text,
            summary_text=summary_text,
            model_cls=model_cls,
            schema_name=schema_name,
        )
        request_metadata = {
            "document_type": normalized_type,
            "schema_name": schema_name,
            "risk_level": policy_decision.risk_level,
            "policy_reason_codes": list(policy_decision.reason_codes),
        }
        cache_context = self._build_cache_context(
            task=LLMTask.CLINICAL_DOCUMENT_EXTRACT,
            document_type=normalized_type,
            schema_name=schema_name,
            schema_version=schema_version,
            source_text=summary_text or raw_text,
            prompt_version=prompt_version,
            request_metadata=request_metadata,
        )
        if cache_context is not None:
            cached = self._cache_repository.get(
                username=cache_context["username"],
                task=LLMTask.CLINICAL_DOCUMENT_EXTRACT.value,
                fingerprint=cache_context["fingerprint"],
            )
            if cached and isinstance(cached.get("payload"), dict):
                return self._hydrate_cached_result(
                    document_type=normalized_type,
                    payload=cached["payload"],
                    provider=str(cached.get("provider") or cache_context["provider"]),
                    model=str(cached.get("model") or cache_context["model"]),
                )
        attempts: list[dict[str, Any]] = []
        response = self._request_structured_analysis(
            prompt=prompt,
            model_cls=model_cls,
            metadata=request_metadata,
        )
        analysis_model = model_cls.model_validate(response.content)
        quality = assess_clinical_extraction(
            document_type=normalized_type,
            raw_text=raw_text,
            analysis_model=analysis_model,
        )
        attempts.append(
            {
                "stage": "initial",
                "provider": response.provider,
                "model": response.model,
                "risk_level": request_metadata["risk_level"],
                "quality_status": quality.status,
                "quality_reasons": list(quality.reason_codes),
            }
        )

        final_response = response
        final_model = analysis_model
        final_quality = quality
        escalated = False
        if self._should_retry_with_high_risk_model(
            current_model=response.model,
            request_metadata=request_metadata,
            quality_payload=quality.to_payload(),
        ):
            escalated = True
            retry_metadata = dict(request_metadata)
            retry_metadata["risk_level"] = "high"
            retry_metadata["retry_reason_codes"] = list(quality.reason_codes)
            retry_metadata["retry_stage"] = "validator_escalation"
            retry_response = self._request_structured_analysis(
                prompt=prompt,
                model_cls=model_cls,
                metadata=retry_metadata,
            )
            retry_model = model_cls.model_validate(retry_response.content)
            retry_quality = assess_clinical_extraction(
                document_type=normalized_type,
                raw_text=raw_text,
                analysis_model=retry_model,
            )
            attempts.append(
                {
                    "stage": "retry",
                    "provider": retry_response.provider,
                    "model": retry_response.model,
                    "risk_level": retry_metadata["risk_level"],
                    "quality_status": retry_quality.status,
                    "quality_reasons": list(retry_quality.reason_codes),
                }
            )
            final_response = retry_response
            final_model = retry_model
            final_quality = retry_quality

        rendered_html = render_clinical_document_html(final_model)
        route_metadata = {
            "document_type": normalized_type,
            "policy_reason_codes": list(policy_decision.reason_codes),
            "initial_risk_level": request_metadata["risk_level"],
            "final_risk_level": "high" if escalated else request_metadata["risk_level"],
            "escalated": escalated,
            "attempts": attempts,
            "selected_provider": final_response.provider,
            "selected_model": final_response.model,
        }
        result = StructuredClinicalAnalysisResult(
            document_type=normalized_type,
            analysis_model=final_model,
            analysis_structured=dump_clinical_model(final_model),
            rendered_html=rendered_html,
            provider=final_response.provider,
            model=final_response.model,
            metrics=final_response.metrics,
            route_metadata=route_metadata,
            quality_metadata=final_quality.to_payload(),
        )
        if cache_context is not None:
            self._cache_repository.upsert(
                username=cache_context["username"],
                task=LLMTask.CLINICAL_DOCUMENT_EXTRACT.value,
                fingerprint=cache_context["fingerprint"],
                provider=cache_context["provider"],
                model=cache_context["model"],
                payload={
                    "analysis_structured": result.analysis_structured,
                    "rendered_html": result.rendered_html,
                    "analysis_schema": final_model.__class__.__name__,
                    "analysis_schema_version": final_model.schema_version,
                    "analysis_render_version": "v1",
                    "route_metadata": dict(result.route_metadata or {}),
                    "quality_metadata": dict(result.quality_metadata or {}),
                },
                metadata={
                    "document_type": normalized_type,
                    "schema_name": schema_name,
                    "schema_version": schema_version,
                    "prompt_version": prompt_version,
                    "response_provider": final_response.provider,
                    "response_model": final_response.model,
                    "risk_level": route_metadata["final_risk_level"],
                    "escalated": escalated,
                },
            )
        return result

    def _build_prompt(
        self,
        *,
        document_type: str,
        raw_text: str,
        summary_text: str | None,
        model_cls: type[ClinicalStructuredModel],
        schema_name: str,
    ) -> str:
        source_text = str(summary_text or raw_text or "").strip()
        schema_json = json.dumps(model_cls.model_json_schema(by_alias=True), ensure_ascii=False, separators=(",", ":"))
        instructions = _DOCUMENT_INSTRUCTIONS.get(document_type, _DOCUMENT_INSTRUCTIONS["generico"])
        return (
            f"Tarea: extraer {document_type}.\n"
            f"Versión de prompt: {self._get_prompt_version(document_type)}.\n"
            f"Contrato: devuelve un único JSON válido para el esquema {schema_name}.\n"
            f"Reglas: {instructions}\n"
            "Omite campos desconocidos, no inventes datos, usa listas vacías si no hay elementos.\n"
            f"Esquema JSON: {schema_json}\n\n"
            f"Documento fuente:\n{source_text}"
        )

    def _build_cache_context(
        self,
        *,
        task: LLMTask,
        document_type: str,
        schema_name: str,
        schema_version: str,
        source_text: str | None,
        prompt_version: str,
        request_metadata: dict[str, Any],
    ) -> dict[str, str] | None:
        username = str(get_log_context().get("username") or "").strip()
        if not username or self._cache_repository is None:
            return None
        route = self._resolve_route_for_cache(task=task, metadata=request_metadata)
        if route is None:
            return None
        return {
            "username": username,
            "provider": route.provider,
            "model": route.model,
            "fingerprint": build_llm_task_fingerprint(
                task=task.value,
                provider=route.provider,
                model=route.model,
                cache_version=getattr(self._cache_repository, "cache_version", "v1"),
                prompt_version=prompt_version,
                schema_version=schema_version,
                payload={
                    "document_type": document_type,
                    "schema_name": schema_name,
                    "source_text": source_text or "",
                },
            ),
        }

    def _get_prompt_version(self, document_type: str) -> str:
        normalized = str(document_type or "").strip().lower() or "generico"
        return _DOCUMENT_PROMPT_VERSIONS.get(normalized, _CLINICAL_EXTRACT_PROMPT_VERSION)

    def _hydrate_cached_result(
        self,
        *,
        document_type: str,
        payload: dict[str, Any],
        provider: str,
        model: str,
    ) -> StructuredClinicalAnalysisResult:
        analysis_structured = dict(payload.get("analysis_structured") or {})
        model_cls = get_clinical_document_model(document_type)
        analysis_model = model_cls.model_validate(analysis_structured)
        rendered_html = str(payload.get("rendered_html") or "") or render_clinical_document_html(analysis_model)
        return StructuredClinicalAnalysisResult(
            document_type=document_type,
            analysis_model=analysis_model,
            analysis_structured=dump_clinical_model(analysis_model),
            rendered_html=rendered_html,
            provider=provider,
            model=model,
            metrics={"source": "cache"},
            route_metadata=dict(payload.get("route_metadata") or {}),
            quality_metadata=dict(payload.get("quality_metadata") or {}),
        )

    def _request_structured_analysis(
        self,
        *,
        prompt: str,
        model_cls: type[ClinicalStructuredModel],
        metadata: dict[str, Any],
    ):
        return self._llm_router.generate_structured(
            LLMStructuredRequest(
                task=LLMTask.CLINICAL_DOCUMENT_EXTRACT,
                prompt=prompt,
                output_model=model_cls,
                output_kind=LLMOutputKind.STRUCTURED_OBJECT,
                system_prompt=(
                    "Eres un extractor clínico. Responde únicamente un objeto JSON válido que cumpla el esquema indicado. "
                    "No uses Markdown, HTML ni texto adicional."
                ),
                metadata=metadata,
            )
        )

    def _resolve_route_for_cache(self, *, task: LLMTask, metadata: dict[str, Any]):
        if not hasattr(self._llm_router, "_resolve_route"):
            return None
        try:
            return self._llm_router._resolve_route(task, metadata=dict(metadata))
        except TypeError:
            return self._llm_router._resolve_route(task)

    def _should_retry_with_high_risk_model(
        self,
        *,
        current_model: str,
        request_metadata: dict[str, Any],
        quality_payload: dict[str, Any],
    ) -> bool:
        if not bool(quality_payload.get("should_retry")):
            return False
        if str(request_metadata.get("risk_level") or "low").strip().lower() == "high":
            return False
        return str(current_model or "").strip() != str(config.GEMINI_MODEL_EXTRACT_HIGH_RISK or "").strip()
