import os
import tempfile
import shutil
import base64
from pathlib import Path
from typing import Optional
from fastapi import FastAPI, UploadFile, File, HTTPException, Query
from fastapi.middleware.cors import CORSMiddleware
import cv2
import numpy as np
from ultralytics import YOLO
from huggingface_hub import hf_hub_download

app = FastAPI(title="WeldSight YOLO Model API Space")

# Enable CORS so the local app can connect
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Your Hugging Face model repositoryy
HF_MODEL_REPO = "chakib2f2sdf/weldsight-yolo-models"

# In-memory dictionary to hold loaded models
_models = {
    "radio": {"binary": None, "4cls": None, "7cls": None},
    "visual": {"binary": None, "4cls": None, "7cls": None}
}

MODEL_VERSIONS = {
    "4cls": "WeldSight-Space-4CLS (P:84.3% R:75.6% mAP50:78.5%)",
    "binary": "WeldSight-Space-Binary (P:93.0% R:79.7% mAP50:88.0%)",
    "7cls": "WeldSight-Space-7CLS-Elite (P:79.7% R:78.1% mAP50:79.5%)"
}

def download_and_load_model(inspection_type: str, model_type: str) -> YOLO:
    global _models
    
    filenames = {
        "radio": {
            "binary": "RT_binary.pt",
            "4cls": "RT_4classe.pt",
            "7cls": "RT_7classes.pt"
        },
        "visual": {
            "binary": "VT_binary.pt",
            "4cls": "VT_6classes.pt",
            "7cls": "VT_6classes.pt"
        }
    }
    
    filename = filenames[inspection_type][model_type]
    
    if _models[inspection_type][model_type] is None:
        print(f"[Loading] Fetching {filename} from Hub repo: {HF_MODEL_REPO}...")
        try:
            model_path = hf_hub_download(
                repo_id=HF_MODEL_REPO,
                filename=filename,
                token=os.getenv("HF_TOKEN")
            )
            device = "cuda" if cv2.cuda.getCudaEnabledDeviceCount() > 0 else "cpu"
            _models[inspection_type][model_type] = YOLO(model_path).to(device)
            print(f"[Success] Loaded model [{inspection_type} -> {model_type}] to {device}")
        except Exception as e:
            print(f"[Error] Failed to load model {filename}: {e}")
            raise RuntimeError(f"Failed to load model {filename}: {e}")
            
    return _models[inspection_type][model_type]


@app.on_event("startup")
def startup_event():
    print(f"[Startup] Pre-loading models from: {HF_MODEL_REPO}")
    for insp_type in ["radio", "visual"]:
        for model_type in ["binary", "4cls"]:
            try:
                download_and_load_model(insp_type, model_type)
            except Exception as e:
                print(f"[Startup Warn] Pre-loading failed for [{insp_type} -> {model_type}]: {e}")


def classify_image_type(image_path: str) -> str:
    try:
        img = cv2.imread(image_path)
        if img is not None and len(img.shape) == 3:
            b, g, r = cv2.split(img)
            if not (np.allclose(b, g) and np.allclose(g, r)):
                return "visual"
    except Exception as ex:
        print(f"[Classifier] Error: {ex}. Defaulting to radio.")
    return "radio"


def preprocess_radio_image(image_path: str):
    try:
        img_array = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
        if img_array is not None:
            denoised = cv2.fastNlMeansDenoising(img_array, None, h=10, templateWindowSize=7, searchWindowSize=21)
            clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8))
            enhanced = clahe.apply(denoised)
            cv2.imwrite(image_path, enhanced)
    except Exception as e:
        print(f"[Preprocessing] Preprocessing failed: {e}")


@app.get("/")
def read_root():
    return {
        "status": "online",
        "service": "WeldSight YOLO Model API Space",
        "model_repo": HF_MODEL_REPO
    }


@app.post("/analyze")
async def analyze(
    file: UploadFile = File(...),
    model_type: str = Query("4cls"),
    inspection_type: str = Query("auto")
):
    if model_type not in ["4cls", "binary", "7cls"]:
        model_type = "4cls"

    suffix = Path(file.filename).suffix or ".jpg"
    with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
        shutil.copyfileobj(file.file, tmp)
        tmp_path = tmp.name

    try:
        resolved_type = inspection_type
        if resolved_type == "auto":
            resolved_type = classify_image_type(tmp_path)

        # Download and load the model on-demand
        model = download_and_load_model(resolved_type, model_type)

        if resolved_type == "radio":
            preprocess_radio_image(tmp_path)

        with open(tmp_path, "rb") as f:
            b64_data = base64.b64encode(f.read()).decode("utf-8")
        preprocessed_image_url = f"data:image/jpeg;base64,{b64_data}"

        if model_type == "4cls":
            imgsz = 1280
        elif model_type == "7cls":
            imgsz = 640
        else:
            imgsz = 1024

        device = "cuda" if cv2.cuda.getCudaEnabledDeviceCount() > 0 else "cpu"
        results = model(tmp_path, imgsz=imgsz, conf=0.10, verbose=False, device=device)
        
        detections = []
        class_names = getattr(model, "names", {})

        for result in results:
            boxes = result.boxes
            masks = getattr(result, "masks", None)
            
            if boxes is None:
                continue

            for i, box in enumerate(boxes):
                cls_id = int(box.cls[0].item())
                conf   = float(box.conf[0].item())
                x1, y1, x2, y2 = [float(v) for v in box.xyxy[0].tolist()]
                label  = class_names.get(cls_id, f"class_{cls_id}")

                detections.append({
                    "type":       "box",
                    "label":      label,
                    "confidence": conf,
                    "xyxy":       [x1, y1, x2, y2],
                })

                if masks is not None and i < len(masks.xy):
                    poly = masks.xy[i]
                    if len(poly) >= 3:
                        points = [[float(p[0]), float(p[1])] for p in poly]
                        detections.append({
                            "type":       "mask",
                            "label":      label,
                            "confidence": conf,
                            "points":     points,
                            "xyxy":       [x1, y1, x2, y2],
                        })

        model_version = f"WeldSight-VT-Visual" if resolved_type == "visual" else MODEL_VERSIONS.get(model_type, model_type)

        return {
            "detections": detections,
            "model_used": model_version,
            "preprocessed_image": preprocessed_image_url
        }

    except Exception as e:
        print(f"[Error] Inference failed: {e}")
        raise HTTPException(status_code=500, detail=str(e))
        
    finally:
        if os.path.exists(tmp_path):
            os.unlink(tmp_path)