from fastapi import FastAPI, UploadFile, File
from fastapi.staticfiles import StaticFiles
from fastapi.responses import HTMLResponse
import cv2
import numpy as np
import torch
import albumentations as A
from albumentations.pytorch import ToTensorV2
from fastapi.middleware.cors import CORSMiddleware
import models as M
from ultralytics import YOLO
import os

app = FastAPI(title="ToothMap AI API")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

# Dynamic checkpoint directory
if os.path.exists("/app/checkpoints"):
    CKPT_DIR = "/app/checkpoints"
else:
    CKPT_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "checkpoints"))
yolo_model = None

@app.on_event("startup")
async def startup_event():
    global yolo_model
    print("Loading internal PyTorch models...")
    M.load_models(CKPT_DIR)

    yolo_path = os.path.join(CKPT_DIR, "yolo_best.pt")
    if os.path.exists(yolo_path):
        yolo_model = YOLO(yolo_path)
        print("✅ YOLO loaded.")

# --- Transforms ---
DET_VAL_TF = A.Compose([
    A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
    ToTensorV2(),
])

CLS_VAL_TF = A.Compose([
    A.Resize(224, 224),
    A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
    ToTensorV2(),
])

SEG_VAL_TF = A.Compose([
    A.Resize(512, 512),
    A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
    ToTensorV2(),
])

def read_image(file_bytes):
    nparr = np.frombuffer(file_bytes, np.uint8)
    img_bgr = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
    return img_bgr

@app.get("/api/health")
def health():
    return {
        "frcnn": "frcnn" in M.MODELS,
        "cls": "cls" in M.MODELS,
        "unet": "unet" in M.MODELS,
        "yolo": yolo_model is not None,
        "device": str(M.DEVICE),
    }

def _classify_crop(img_rgb_crop: np.ndarray) -> dict:
    """Run ResNet18 on a single cropped tooth and return FDI + confidence."""
    cls = M.MODELS.get("cls")
    if cls is None:
        return {"fdi": -1, "confidence": 0.0}
    t = CLS_VAL_TF(image=img_rgb_crop)
    img_t = t["image"].unsqueeze(0).to(M.DEVICE)
    with torch.no_grad():
        pred = cls(img_t)
        probs = torch.softmax(pred, dim=1)
        fdi_idx = pred.argmax(dim=1).item()
        score = probs[0, fdi_idx].item()
    quad = fdi_idx // 8
    num  = fdi_idx % 8
    real_fdi = (quad + 1) * 10 + (num + 1)
    return {"fdi": real_fdi, "confidence": round(score, 4)}

def _segment_crop(img_rgb_crop: np.ndarray) -> str:
    """Run U-Net on a single cropped tooth, return base64 PNG mask with alpha."""
    import base64
    unet = M.MODELS.get("unet")
    if unet is None or img_rgb_crop.size == 0:
        return ""
    t = SEG_VAL_TF(image=img_rgb_crop)
    img_t = t["image"].unsqueeze(0).to(M.DEVICE)
    with torch.no_grad():
        out = unet(img_t)
        mask = (torch.sigmoid(out) > 0.5).squeeze().cpu().numpy()
    mask_cv = (mask * 255).astype(np.uint8)
    rgba = np.zeros((mask_cv.shape[0], mask_cv.shape[1], 4), dtype=np.uint8)
    rgba[:, :, 0] = 255  # B
    rgba[:, :, 1] = 255  # G
    rgba[:, :, 2] = 255  # R
    rgba[:, :, 3] = mask_cv  # A (transparency)
    _, buf = cv2.imencode('.png', rgba)
    return base64.b64encode(buf).decode('utf-8')

@app.post("/api/pipeline/yolo")
async def pipeline_yolo(file: UploadFile = File(...)):
    """Detect teeth (YOLO) then classify each crop (ResNet18) → returns annotated boxes."""
    if yolo_model is None:
        return {"error": "YOLO not loaded"}
    img_bgr = read_image(await file.read())
    img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
    orig_h, orig_w = img_rgb.shape[:2]
    img_yolo = cv2.resize(img_rgb, (640, 640))
    ypred = yolo_model.predict(img_yolo, conf=0.20, verbose=False)[0]
    boxes = ypred.boxes.xyxy.cpu().numpy()
    if len(boxes) > 0:
        boxes[:, [0, 2]] *= (orig_w / 640.0)
        boxes[:, [1, 3]] *= (orig_h / 640.0)
    scores = ypred.boxes.conf.cpu().numpy().tolist()
    results = []
    for i, box in enumerate(boxes):
        x1, y1, x2, y2 = map(int, box)
        x1, y1 = max(0, x1), max(0, y1)
        x2, y2 = min(orig_w, x2), min(orig_h, y2)
        crop = img_rgb[y1:y2, x1:x2]
        cls_result = _classify_crop(crop) if crop.size > 0 else {"fdi": -1, "confidence": 0.0}
        mask_b64   = _segment_crop(crop)
        results.append({"box": list(map(float, box)), "score": scores[i], **cls_result, "mask_base64": mask_b64})
    return {"results": results}

@app.post("/api/pipeline/frcnn")
async def pipeline_frcnn(file: UploadFile = File(...)):
    """Detect teeth (FRCNN) then classify each crop (ResNet18) → returns annotated boxes."""
    frcnn = M.MODELS.get("frcnn")
    if frcnn is None:
        return {"error": "FRCNN not loaded"}
    img_bgr = read_image(await file.read())
    img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
    orig_h, orig_w = img_rgb.shape[:2]
    IMG_W, IMG_H = 1000, 500
    img_det = cv2.resize(img_rgb, (IMG_W, IMG_H))
    t = DET_VAL_TF(image=img_det)
    img_t = t["image"].unsqueeze(0).to(M.DEVICE)
    with torch.no_grad():
        pred = frcnn(img_t)[0]
    keep = pred["scores"] >= 0.20
    boxes = pred["boxes"][keep].cpu().numpy()
    scores = pred["scores"][keep].cpu().numpy().tolist()
    if len(boxes) > 0:
        boxes[:, [0, 2]] *= (orig_w / float(IMG_W))
        boxes[:, [1, 3]] *= (orig_h / float(IMG_H))
    results = []
    for i, box in enumerate(boxes):
        x1, y1, x2, y2 = map(int, box)
        x1, y1 = max(0, x1), max(0, y1)
        x2, y2 = min(orig_w, x2), min(orig_h, y2)
        crop = img_rgb[y1:y2, x1:x2]
        cls_result = _classify_crop(crop) if crop.size > 0 else {"fdi": -1, "confidence": 0.0}
        mask_b64   = _segment_crop(crop)
        results.append({"box": list(map(float, box)), "score": scores[i], **cls_result, "mask_base64": mask_b64})
    return {"results": results}

@app.post("/api/detect/yolo")
async def detect_yolo(file: UploadFile = File(...)):
    if yolo_model is None:
        return {"error": "YOLO not loaded"}
    img_bgr = read_image(await file.read())
    img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
    orig_h, orig_w = img_rgb.shape[:2]
    img_yolo = cv2.resize(img_rgb, (640, 640))
    ypred = yolo_model.predict(img_yolo, conf=0.20, verbose=False)[0]
    boxes = ypred.boxes.xyxy.cpu().numpy()
    if len(boxes) > 0:
        boxes[:, [0, 2]] *= (orig_w / 640.0)
        boxes[:, [1, 3]] *= (orig_h / 640.0)
    scores = ypred.boxes.conf.cpu().numpy().tolist()
    labels = ypred.boxes.cls.cpu().numpy().astype(int).tolist()
    return {"boxes": boxes.tolist(), "scores": scores, "labels": labels}

@app.post("/api/detect/frcnn")
async def detect_frcnn(file: UploadFile = File(...)):
    frcnn = M.MODELS.get("frcnn")
    if frcnn is None:
        return {"error": "FRCNN not loaded"}
    img_bgr = read_image(await file.read())
    img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
    orig_h, orig_w = img_rgb.shape[:2]
    IMG_W, IMG_H = 1000, 500
    img_det = cv2.resize(img_rgb, (IMG_W, IMG_H))
    t = DET_VAL_TF(image=img_det)
    img_t = t["image"].unsqueeze(0).to(M.DEVICE)
    with torch.no_grad():
        pred = frcnn(img_t)[0]
    keep = pred["scores"] >= 0.20
    boxes = pred["boxes"][keep].cpu().numpy()
    if len(boxes) > 0:
        boxes[:, [0, 2]] *= (orig_w / float(IMG_W))
        boxes[:, [1, 3]] *= (orig_h / float(IMG_H))
    return {
        "boxes": boxes.tolist(),
        "scores": pred["scores"][keep].cpu().numpy().tolist(),
        "labels": pred["labels"][keep].cpu().numpy().tolist()
    }

@app.post("/api/classify")
async def classify_crops(file: UploadFile = File(...)):
    cls = M.MODELS.get("cls")
    if cls is None:
        return {"error": "Classifier not loaded"}
    img_bgr = read_image(await file.read())
    img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
    t = CLS_VAL_TF(image=img_rgb)
    img_t = t["image"].unsqueeze(0).to(M.DEVICE)
    with torch.no_grad():
        pred = cls(img_t)
        confidences = torch.softmax(pred, dim=1)
        fdi_idx = pred.argmax(dim=1).item()
        score = confidences[0, fdi_idx].item()
    quad = fdi_idx // 8
    num = fdi_idx % 8
    real_fdi = (quad + 1) * 10 + (num + 1)
    return {"fdi": real_fdi, "confidence": score}

@app.post("/api/segment")
async def segment_unet(file: UploadFile = File(...)):
    unet = M.MODELS.get("unet")
    if unet is None:
        return {"error": "U-Net not loaded"}
    import base64
    img_bgr = read_image(await file.read())
    img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
    t = SEG_VAL_TF(image=img_rgb)
    img_t = t["image"].unsqueeze(0).to(M.DEVICE)
    with torch.no_grad():
        out = unet(img_t)
        mask = (torch.sigmoid(out) > 0.5).squeeze().cpu().numpy()
    mask_img = (mask * 255).astype(np.uint8)
    _, buffer = cv2.imencode('.png', mask_img)
    encoded = base64.b64encode(buffer).decode('utf-8')
    return {"mask_base64": encoded}

# --- Mount React Frontend ---
# This MUST be the last route registered so it doesn't shadow API routes
app.mount("/assets", StaticFiles(directory="static/assets"), name="assets")
app.mount("/", StaticFiles(directory="static", html=True), name="static")