#!/usr/bin/env python3
"""
servir_ensemble.py — API do ensemble (média das probs calibradas de N decisors).

Mesmos endpoints do servir_decisor (GET /saude, POST /decidir); carrega 2+
checkpoints, roda cada pergunta em cada modelo e devolve a probabilidade
MÉDIA (calibrada por modelo antes da média — média de distribuições
calibradas continua calibrada). Roteamento: mínima confiança da média.

Uso (host .144, GPU 1, porta de produção ou teste):
  python3 scripts/servir_ensemble.py \
      --checkpoints models/decisor-v10 models/decisor-v11_cons \
      --device cuda --porta 8101
"""
import argparse
import json
import sys
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
from chamados_para_decisor import CATEGORIAS, DECISOES, URGENCIA  # noqa: E402
from inferencia_decisor import DecisorCPU                         # noqa: E402

PERGUNTAS = {
    "categoria": {"tipo": "choice",
                  "instrucao": "Qual a categoria principal do chamado?",
                  "opcoes": CATEGORIAS},
    "decisao":   {"tipo": "choice",
                  "instrucao": "Qual decisão tomar para este caso?",
                  "opcoes": DECISOES},
    "urgencia":  {"tipo": "score",
                  "instrucao": "Qual a urgência? (0=baixa, 3=crítica)",
                  "opcoes": URGENCIA},
}

DECISORES = []


def decidir(estado: str, limiar: float) -> dict:
    respostas, confs = {}, []
    for nome, pergunta in PERGUNTAS.items():
        ops = pergunta["opcoes"]
        media = None
        votos = []
        for d in DECISORES:
            r = d["modelo"].prever(estado, pergunta)
            probs = [r["distribuicao"].get(op, 0.0) for op in ops]
            votos.append(probs)
            votos.append(None)  # placeholder p/ manter par modelo/(peso)
            votos.pop()
        media = [sum(v[k] for v in votos) / len(votos) for k in range(len(ops))]
        i = max(range(len(ops)), key=lambda k: media[k])
        respostas[nome] = {
            "resposta": ops[i],
            "confianca": round(media[i], 4),
            "distribuicao": {op: round(media[k], 4)
                             for op, k in zip(ops, range(len(ops)))},
        }
        confs.append(media[i])
    conf = min(confs)
    return {"respostas": respostas,
            "confianca_roteamento": round(conf, 4),
            "limiar": limiar,
            "destino": "automatico" if conf >= limiar else "system2"}


class Handler(BaseHTTPRequestHandler):
    def _json(self, obj, status=200):
        corpo = json.dumps(obj, ensure_ascii=False).encode()
        self.send_response(status)
        self.send_header("Content-Type", "application/json; charset=utf-8")
        self.send_header("Content-Length", str(len(corpo)))
        self.end_headers()
        self.wfile.write(corpo)

    def log_message(self, *a):
        pass

    def do_GET(self):
        if self.path == "/saude":
            self._json({
                "status": "ok",
                "ensemble": [d["nome"] for d in DECISORES],
                "n_modelos": len(DECISORES),
                "calibrado": all(bool(d["modelo"].temps) for d in DECISORES),
                "categorias": len(CATEGORIAS),
            })
            return
        self._json({"erro": "use GET /saude ou POST /decidir"}, 404)

    def do_POST(self):
        if self.path != "/decidir":
            self._json({"erro": "rota desconhecida"}, 404)
            return
        try:
            corpo = json.loads(self.rfile.read(int(self.headers["Content-Length"])))
            estado = str(corpo["estado"]).strip()
            if not estado:
                raise ValueError("estado vazio")
            limiar = float(corpo.get("limiar", 0.8))
            if not 0.0 <= limiar <= 1.0:
                raise ValueError("limiar fora de [0,1]")
        except (KeyError, ValueError, TypeError) as e:
            self._json({"erro": f"payload inválido: {e}",
                        "formato": {"estado": "texto", "limiar": 0.8}}, 400)
            return
        self._json(decidir(estado, limiar))


def _device_auto() -> str:
    """Escolhe o melhor device disponível: cuda → mps → cpu."""
    import torch
    if torch.cuda.is_available():
        return "cuda"
    if getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
        return "mps"
    return "cpu"


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--checkpoints", nargs="+", required=True)
    ap.add_argument("--device", default="auto",
                    choices=["auto", "cuda", "mps", "xpu", "cpu"],
                    help="default: auto (cuda → mps → cpu)")
    ap.add_argument("--porta", type=int, default=8101)
    args = ap.parse_args()
    device = _device_auto() if args.device == "auto" else args.device
    print(f"device: {device}")
    for ck in args.checkpoints:
        DECISORES.append({"nome": ck.rstrip("/").split("/")[-1],
                          "modelo": DecisorCPU(ck, device=device)})
        print(f"✓ carregado: {ck}")
    print(f"ensemble de {len(DECISORES)} modelos na porta {args.porta}")
    ThreadingHTTPServer(("0.0.0.0", args.porta), Handler).serve_forever()


if __name__ == "__main__":
    main()