import argparse
import json
import os
from time import perf_counter
from typing import Any

from dotenv import load_dotenv
from groq import Groq
from langchain_community.vectorstores import FAISS

from app.config import config, require_env
from app.core.logging import InstrumentedGroqClient, get_audit_logger, get_logger
from app.llm import LLMTask
from app.llm.task_text import generate_task_text
from modules.processing.embeddings import get_hf_embeddings


load_dotenv()
GROQ_API_KEY = require_env(os.getenv("GROQ_API_KEY") or config.GROQ_API_KEY, "GROQ_API_KEY")
EMBEDDINGS_MODEL = os.getenv('EMBEDDINGS_MODEL', 'intfloat/e5-small-v2')
logger = get_logger(__name__)
audit_logger = get_audit_logger()

class SOATRetriever:
    def __init__(
        self,
        faiss_index_path: str,
        groq_api_key: str,
        groq_model: str = None,
        llm_router: Any | None = None,
    ):
        self.embeddings = get_hf_embeddings(EMBEDDINGS_MODEL)  # "google/embeddinggemma-300m"
        self.db = FAISS.load_local(faiss_index_path, self.embeddings, allow_dangerous_deserialization=True)
        self.llm_router = llm_router
        self.groq = None
        if self.llm_router is None:
            self.groq = InstrumentedGroqClient(
                Groq(api_key=groq_api_key),
                provider="groq",
                service="soat_retriever",
            )
        self.groq_model = groq_model or config.GROQ_MODEL_SOAT
        self.index_path = faiss_index_path

    def buscar(self, consulta: str, k: int = 20):
        started_at = perf_counter()
        try:
            docs = self.db.max_marginal_relevance_search(consulta, k=k, lambda_param=0.7)
            audit_logger.rag_event(
                action="soat.buscar",
                outcome="success",
                service="soat_retriever",
                resource={"index_path": self.index_path, "query_type": "soat", "index_name": "soat"},
                metrics={
                    "duration_ms": round((perf_counter() - started_at) * 1000, 3),
                    "k": k,
                    "result_count": len(docs),
                },
            )
            return docs
        except Exception as exc:
            audit_logger.rag_event(
                action="soat.buscar",
                outcome="error",
                service="soat_retriever",
                resource={"index_path": self.index_path, "query_type": "soat", "index_name": "soat"},
                metrics={"duration_ms": round((perf_counter() - started_at) * 1000, 3), "k": k},
                error={"class": exc.__class__.__name__, "message": str(exc)},
            )
            raise

    def proponer_codigos_soat(self, descripcion: str, docs) -> str:
        contexto = "\n\n".join(
            f"(p.{d.metadata.get('page')}) {d.page_content[:1200]}" for d in docs
        )
        prompt = f"""
Eres un especialista en auditoría médica en Colombia. Basado EXCLUSIVAMENTE en el Manual SOAT 2024 siguiente (contexto), propone códigos SOAT válidos para el procedimiento:

Descripción del procedimiento: {descripcion}

Contexto (Manual SOAT 2024 - extractos relevantes):
{contexto}

Reglas:
- Devuelve SOLO códigos SOAT que aparezcan explícitamente en el contexto con su descripción exacta.
- Incluye la página entre paréntesis de donde lo tomaste.
- Si no hay coincidencia exacta, devuelve vacío.

Formato por línea: <codigo_soat> - <descripcion> (p.<pagina>)
"""
        if self.llm_router is not None:
            response = generate_task_text(
                self.llm_router,
                task=LLMTask.SOAT_MANUAL_CODE_LOOKUP,
                prompt=prompt,
                system_prompt="Eres estricto: no inventes códigos. Usa solo lo presente en el contexto.",
                metadata={"operation": "soat_manual_code_lookup", "descripcion": descripcion[:200]},
                generation_config={"temperature": 0.0},
            )
            return response.content.strip()

        resp = self.groq.chat.completions.create(
            model=self.groq_model,
            temperature=0.0,
            messages=[
                {"role": "system", "content": "Eres estricto: no inventes códigos. Usa solo lo presente en el contexto."},
                {"role": "user", "content": prompt},
            ],
        )
        return resp.choices[0].message.content.strip()

    def asignar_codigos(self, procedimientos: list[str]) -> list[dict[str, str]]:
        resultados = []
        for p in procedimientos:
            docs = self.buscar(p, k=12)
            propuesta = self.proponer_codigos_soat(p, docs)
            if propuesta:
                for linea in propuesta.splitlines():
                    linea = linea.strip()
                    if not linea:
                        continue
                    resultados.append({"procedimiento": p, "codigo_soat": linea})
        return resultados


def main():
    parser = argparse.ArgumentParser(description="Proponer códigos SOAT vía RAG desde un índice FAISS")
    parser.add_argument("--index", dest="index_dir", default="soat_faiss", help="Directorio del índice FAISS (por defecto: soat_faiss)")
    parser.add_argument("--groq_key", dest="groq_key", default=GROQ_API_KEY, help="API key de Groq (por defecto, toma GROQ_API_KEY del entorno)")
    parser.add_argument("--groq_model", dest="groq_model", default=os.getenv("GROQ_MODEL_SOAT") or config.GROQ_MODEL_SOAT, help="Modelo Groq (por defecto: GROQ_MODEL_SOAT)")
    parser.add_argument("--proc", dest="procedimientos", action="append", help="Descripción de procedimiento (repetible)")
    parser.add_argument("--proc_file", dest="proc_file", help="Archivo con procedimientos, uno por línea")
    args = parser.parse_args()

    if not args.groq_key:
        raise SystemExit("Falta GROQ_API_KEY (usa --groq_key o exporta variable de entorno)")

    procedimientos: list[str] = []
    if args.procedimientos:
        procedimientos.extend([p for p in args.procedimientos if p and p.strip()])
    if args.proc_file and os.path.exists(args.proc_file):
        with open(args.proc_file, encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if line:
                    procedimientos.append(line)
    if not procedimientos:
        procedimientos = [
           "REDUCCIÓN ABIERTA DE FRACTURA CON FIJACIÓN INTERNA  (DISPOSITIVOS DE FIJACIÓN U OSTEOSÍNTESIS) DE CLAVÍCULA",
           # "Desbridamiento de tejido blando de pie izquierdo",
            #"Reducción cerrada de fractura de falange proximal del pie"
        ]

    retriever = SOATRetriever(faiss_index_path=args.index_dir, groq_api_key=args.groq_key, groq_model=args.groq_model)
    resultados = retriever.asignar_codigos(procedimientos)
    print(json.dumps({"procedimientos": procedimientos, "resultados": resultados}, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()
