from __future__ import annotations
import base64
import io
import os
import warnings
from pathlib import Path

import cv2

import numpy as np
from fastapi import FastAPI, File, Form, Header, HTTPException, UploadFile
from fastapi.responses import HTMLResponse
from fastapi.middleware.cors import CORSMiddleware
from huggingface_hub import hf_hub_download

os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "3")
warnings.filterwarnings("ignore")

import tensorflow as tf
tf.get_logger().setLevel("ERROR")

HF_REPO = "jvrfer/medical-ai-model"
IMAGE_SIZE = (384, 384)
THRESHOLD = 0.35
GRADCAM_LAYER = "top_conv"

from model import RandomBrightness
from dataset import load_inference_tensors_from_bytes
from gradcam import make_gradcam_plus_plus_heatmap, overlay_heatmap
from utils import load_json

API_KEY = os.environ.get("MEDICAL_AI_API_KEY")
if not API_KEY:
    raise RuntimeError("MEDICAL_AI_API_KEY no configurada")

print("Descargando modelo...")
model_path = hf_hub_download(repo_id=HF_REPO, filename="best_model_phase2.keras")
class_names_path = hf_hub_download(repo_id=HF_REPO, filename="class_names.json")

CUSTOM_OBJECTS = {"RandomBrightness": RandomBrightness}
print("Cargando modelo...")
MODEL = tf.keras.models.load_model(model_path, custom_objects=CUSTOM_OBJECTS)
CLASS_NAMES = load_json(class_names_path)["class_names"]
print(f"Modelo cargado. Clases: {CLASS_NAMES}")

LEGAL_WARNING = (
    "Advertencia: esta salida es solo apoyo informatico y no constituye diagnostico medico. "
    "Debe ser interpretada por personal de salud calificado."
)

app = FastAPI(title="Medical AI API", version="1.0.0")
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])


def validate_key(api_key: str | None):
    if not api_key:
        raise HTTPException(401, "Falta header X-API-Key")
    if api_key != API_KEY:
        raise HTTPException(403, "API key invalida")


def predict_bytes(filename: str, payload: bytes) -> dict:
    _, processed = load_inference_tensors_from_bytes(filename, payload, image_size=IMAGE_SIZE)
    batch = np.expand_dims(processed, axis=0)
    probs = MODEL.predict(batch, verbose=0)[0]
    results = {cls: float(p) for cls, p in zip(CLASS_NAMES, probs)}
    top_class = max(results, key=results.get)
    return {"diagnosis": top_class, "confidence": results[top_class], "all_probabilities": results}


@app.get("/health")
def health():
    return {"status": "healthy"}


@app.post("/predict")
async def predict(file: UploadFile = File(...), x_api_key: str | None = Header(default=None)):
    validate_key(x_api_key)
    payload = await file.read()
    if not payload:
        raise HTTPException(400, "Archivo vacio")
    try:
        result = predict_bytes(file.filename or "image", payload)
    except Exception as e:
        raise HTTPException(400, f"Error: {e}")
    return {**result, "threshold_used": THRESHOLD, "legal_warning": LEGAL_WARNING}


@app.post("/predict-with-gradcam")
async def predict_gradcam(file: UploadFile = File(...), x_api_key: str | None = Header(default=None)):
    validate_key(x_api_key)
    payload = await file.read()
    if not payload:
        raise HTTPException(400, "Archivo vacio")
    try:
        original, processed = load_inference_tensors_from_bytes(
            file.filename or "image", payload, image_size=IMAGE_SIZE
        )
    except Exception as e:
        raise HTTPException(400, f"Error: {e}")

    batch = np.expand_dims(processed, axis=0)
    probs = MODEL.predict(batch, verbose=0)[0]
    results = {cls: float(p) for cls, p in zip(CLASS_NAMES, probs)}
    top_class = max(results, key=results.get)

    gradcam_b64 = None
    try:
        heatmap = make_gradcam_plus_plus_heatmap(MODEL, batch, layer_name=GRADCAM_LAYER)
        overlay, _ = overlay_heatmap(original, heatmap)
        _, buf = cv2.imencode(".png", cv2.cvtColor(overlay, cv2.COLOR_RGB2BGR))
        gradcam_b64 = base64.b64encode(buf).decode()
    except Exception:
        pass

    return {
        "diagnosis": top_class,
        "confidence": results[top_class],
        "all_probabilities": results,
        "gradcam_b64": gradcam_b64,
        "threshold_used": THRESHOLD,
        "legal_warning": LEGAL_WARNING,
    }


WEB_HTML = """<!DOCTYPE html>
<script>const API_KEY = "{api_key}";</script>
<html lang="es">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1">
<title>Medical AI</title>
<style>
  * { box-sizing: border-box; margin: 0; padding: 0; }
  body { font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif; background: #f0f2f5; color: #1a1a1a; }
  .container { max-width: 900px; margin: 0 auto; padding: 2rem 1rem; }
  h1 { font-size: 1.8rem; margin-bottom: .25rem; }
  .sub { color: #555; margin-bottom: 1.5rem; }
  .card { background: #fff; border-radius: 12px; padding: 1.5rem; box-shadow: 0 1px 3px rgba(0,0,0,.1); margin-bottom: 1.5rem; }
  .upload-area { border: 2px dashed #ccc; border-radius: 8px; padding: 2rem; text-align: center; cursor: pointer; transition: .2s; }
  .upload-area:hover, .upload-area.dragover { border-color: #2563eb; background: #f8faff; }
  .upload-area input { display: none; }
  .upload-area img { max-width: 100%; max-height: 300px; margin-top: 1rem; display: none; }
  button { background: #2563eb; color: #fff; border: none; padding: .75rem 2rem; border-radius: 8px; font-size: 1rem; cursor: pointer; margin-top: 1rem; }
  button:hover { background: #1d4ed8; }
  button:disabled { opacity: .6; cursor: not-allowed; }
  .loading { display: none; margin-top: 1rem; }
  .spinner { width: 24px; height: 24px; border: 3px solid #e5e7eb; border-top-color: #2563eb; border-radius: 50%; animation: spin .6s linear infinite; display: inline-block; vertical-align: middle; margin-right: .5rem; }
  @keyframes spin { to { transform: rotate(360deg); } }
  .result { display: none; }
  .result .diagnosis { font-size: 1.2rem; font-weight: 600; margin-bottom: 1rem; }
  .result .diagnosis .conf { font-weight: 400; color: #666; font-size: 1rem; }
  .prob-grid { display: grid; grid-template-columns: 1fr 1fr; gap: .5rem; }
  .prob-item { display: flex; justify-content: space-between; padding: .5rem; background: #f9fafb; border-radius: 6px; }
  .prob-item .bar-wrap { flex: 1; margin: 0 .75rem; background: #e5e7eb; border-radius: 4px; overflow: hidden; }
  .prob-item .bar { height: 100%; background: #2563eb; border-radius: 4px; transition: width .5s; }
  .prob-item .pct { font-weight: 600; min-width: 3.5rem; text-align: right; }
  .top-class { background: #dbeafe; }
  .gradcam-wrap { margin-top: 1.5rem; display: none; }
  .gradcam-wrap img { max-width: 100%; border-radius: 8px; }
  .legal { font-size: .8rem; color: #888; margin-top: 1rem; padding: .75rem; background: #fefce8; border-radius: 8px; }
  @media (max-width: 600px) { .prob-grid { grid-template-columns: 1fr; } }
</style>
</head>
<body>
<div class="container">
  <h1>Medical AI</h1>
  <p class="sub">Analisis de radiografias de torax</p>

  <div class="card">
    <div class="upload-area" id="dropZone">
      <p>Arrastra una imagen aqui o haz clic para seleccionar</p>
      <p style="font-size:.85rem;color:#888;margin-top:.5rem">JPG, PNG, WEBP, BMP, DICOM</p>
      <input type="file" id="fileInput" accept="image/*,.dcm">
      <img id="preview">
    </div>
    <button id="analyzeBtn">Analizar</button>
    <div class="loading" id="loading">
      <span class="spinner"></span> Procesando imagen...
    </div>
  </div>

  <div class="card result" id="resultCard">
    <div class="diagnosis" id="diagnosis"></div>
    <div class="prob-grid" id="probGrid"></div>
    <div class="gradcam-wrap" id="gradcamWrap">
      <h3 style="margin-bottom:.5rem">Mapa de calor GradCAM++</h3>
      <img id="gradcamImg">
    </div>
    <div class="legal" id="legal"></div>
  </div>
</div>

<script>
const dropZone = document.getElementById('dropZone');
const fileInput = document.getElementById('fileInput');
const preview = document.getElementById('preview');
const analyzeBtn = document.getElementById('analyzeBtn');
const loading = document.getElementById('loading');
const resultCard = document.getElementById('resultCard');
const diagnosis = document.getElementById('diagnosis');
const probGrid = document.getElementById('probGrid');
const gradcamWrap = document.getElementById('gradcamWrap');
const gradcamImg = document.getElementById('gradcamImg');
const legal = document.getElementById('legal');

let currentFile = null;

dropZone.addEventListener('click', () => fileInput.click());
dropZone.addEventListener('dragover', e => { e.preventDefault(); dropZone.classList.add('dragover'); });
dropZone.addEventListener('dragleave', () => dropZone.classList.remove('dragover'));
dropZone.addEventListener('drop', e => { e.preventDefault(); dropZone.classList.remove('dragover'); handleFile(e.dataTransfer.files[0]); });
fileInput.addEventListener('change', () => { if (fileInput.files[0]) handleFile(fileInput.files[0]); });

function handleFile(file) {
  currentFile = file;
  const reader = new FileReader();
  reader.onload = e => { preview.src = e.target.result; preview.style.display = 'block'; };
  reader.readAsDataURL(file);
}

analyzeBtn.addEventListener('click', async () => {
  if (!currentFile) return alert('Selecciona una imagen primero');
  loading.style.display = 'block';
  analyzeBtn.disabled = true;
  resultCard.style.display = 'none';
  gradcamWrap.style.display = 'none';

  const form = new FormData();
  form.append('file', currentFile);

  try {
    const res = await fetch('/predict-with-gradcam', {
      method: 'POST',
      headers: { 'X-API-Key': API_KEY },
      body: form,
    });
    if (!res.ok) { const err = await res.json(); throw new Error(err.detail || 'Error'); }
    const data = await res.json();
    showResult(data);
  } catch (err) {
    diagnosis.innerHTML = '<span style="color:#dc2626">Error: ' + err.message + '</span>';
    resultCard.style.display = 'block';
  } finally {
    loading.style.display = 'none';
    analyzeBtn.disabled = false;
  }
});

function showResult(data) {
  const top = data.diagnosis;
  document.getElementById('diagnosis').innerHTML = 'Diagnostico: <strong>' + top + '</strong> <span class="conf">(' + (data.confidence*100).toFixed(1) + '% confianza)</span>';
  if (data.confidence < 0.35) diagnosis.innerHTML += '<br><span style="color:#d97706">Confianza baja. Considere derivar a especialista.</span>';

  const classes = Object.entries(data.all_probabilities).sort((a, b) => b[1] - a[1]);
  probGrid.innerHTML = classes.map(([cls, prob]) => {
    const pct = (prob * 100).toFixed(1);
    const isTop = cls === top;
    return '<div class="prob-item' + (isTop ? ' top-class' : '') + '"><span>' + cls + '</span><div class="bar-wrap"><div class="bar" style="width:' + pct + '%"></div></div><span class="pct">' + pct + '%</span></div>';
  }).join('');

  if (data.gradcam_b64) {
    gradcamImg.src = 'data:image/png;base64,' + data.gradcam_b64;
    gradcamWrap.style.display = 'block';
  }

  legal.textContent = data.legal_warning || '';
  resultCard.style.display = 'block';
}
</script>
</body>
</html>"""


@app.get("/", response_class=HTMLResponse)
def web_ui():
    return HTMLResponse(WEB_HTML.replace("{api_key}", API_KEY if API_KEY else ""))