import io
import base64
import numpy as np
from pathlib import Path
from PIL import Image, ImageDraw, ImageFont
from .disease_knowledge import get_disease_info

MODEL_LABEL = "best.onnx (YOLOv12)"

CLASS_NAMES = [
    "Darier_s-Disease", "Epidermolysis-Bullosa-Pruriginosa", "Hailey-Hailey-Disease",
    "Hemangiome", "Impetigo", "Leishmanios", "Lupus-Erythematosus-Chronicus-Discoides",
    "Melanoma", "Molluscum-Contagiosum", "Porokeratosis", "Psoriasis", "Tinea-Corporis",
    "Tungiasis", "acne", "basal-cell-carcinoma", "eczema", "lichen", "nevus", "normal skin",
]

DISPLAY_NAMES = {
    "Darier_s-Disease": "Darier's Disease",
    "Epidermolysis-Bullosa-Pruriginosa": "Epidermolysis Bullosa Pruriginosa",
    "Hailey-Hailey-Disease": "Hailey-Hailey Disease",
    "Hemangiome": "Hemangioma",
    "Impetigo": "Impetigo",
    "Leishmanios": "Leishmaniasis",
    "Lupus-Erythematosus-Chronicus-Discoides": "Discoid Lupus Erythematosus",
    "Melanoma": "Melanoma",
    "Molluscum-Contagiosum": "Molluscum Contagiosum",
    "Porokeratosis": "Porokeratosis",
    "Psoriasis": "Psoriasis",
    "Tinea-Corporis": "Tinea Corporis (Ringworm)",
    "Tungiasis": "Tungiasis",
    "acne": "Acne",
    "basal-cell-carcinoma": "Basal Cell Carcinoma",
    "eczema": "Eczema",
    "lichen": "Lichen Planus",
    "nevus": "Nevus (Mole)",
    "normal skin": "Normal Skin",
}

_ort_session = None
MODEL_PATH = Path(__file__).parent.parent / "best.onnx"

def get_ort_session():
    global _ort_session
    if _ort_session is None:
        try:
            import onnxruntime as ort
            providers = ["CPUExecutionProvider"]
            _ort_session = ort.InferenceSession(str(MODEL_PATH), providers=providers)
        except Exception as e:
            raise Exception(f"Failed to load CV model: {str(e)}. Ensure best.onnx exists in the backend folder.")
    return _ort_session

def preprocess_image(image_bytes: bytes, input_size: int = 640) -> np.ndarray:
    img = Image.open(io.BytesIO(image_bytes)).convert("RGB")
    w, h = img.size
    scale = input_size / max(w, h)
    new_w, new_h = int(w * scale), int(h * scale)
    img = img.resize((new_w, new_h), Image.BILINEAR)

    canvas = Image.new("RGB", (input_size, input_size), (114, 114, 114))
    pad_x = (input_size - new_w) // 2
    pad_y = (input_size - new_h) // 2
    canvas.paste(img, (pad_x, pad_y))

    arr = np.array(canvas, dtype=np.float32) / 255.0
    arr = arr.transpose(2, 0, 1)
    arr = np.expand_dims(arr, 0)
    return arr

def parse_yolo_output(output: np.ndarray) -> tuple[str, float]:
    num_classes = len(CLASS_NAMES)
    out = output[0]  
    
    if out.ndim == 2 and out.shape[0] == (4 + num_classes):
        class_scores = out[4:, :]
        per_anchor_max = class_scores.max(axis=0)
        best_anchor = int(np.argmax(per_anchor_max))
        anchor_scores = class_scores[:, best_anchor]
    elif out.ndim == 2 and out.shape[1] == (4 + num_classes):
        class_scores = out[:, 4:]
        per_anchor_max = class_scores.max(axis=1)
        best_anchor = int(np.argmax(per_anchor_max))
        anchor_scores = class_scores[best_anchor, :]
    elif out.ndim == 1:
        anchor_scores = out[:num_classes]
    elif out.ndim == 2 and out.shape[1] == num_classes:
        anchor_scores = out[0]
    else:
        flat = out.flatten()
        anchor_scores = flat[:num_classes] if len(flat) >= num_classes else flat

    exp_s = np.exp(anchor_scores - anchor_scores.max())
    probs = exp_s / exp_s.sum()

    class_idx = int(np.argmax(probs))
    confidence = float(probs[class_idx])
    class_idx = min(class_idx, num_classes - 1)

    raw_name = CLASS_NAMES[class_idx]
    display_name = DISPLAY_NAMES.get(raw_name, raw_name)
    return display_name, round(confidence, 4)

CRITICAL_DISEASES = {"Melanoma", "Basal Cell Carcinoma"}
SEVERE_DISEASES = {"Discoid Lupus Erythematosus", "Darier's Disease", "Epidermolysis Bullosa Pruriginosa", "Hailey-Hailey Disease"}

def determine_severity_str(confidence: float, disease_name: str) -> str:
    if disease_name in CRITICAL_DISEASES: return "critical"
    if disease_name in SEVERE_DISEASES: return "severe" if confidence >= 0.5 else "moderate"
    if disease_name == "Normal Skin": return "mild"
    if confidence >= 0.85: return "severe"
    elif confidence >= 0.65: return "moderate"
    else: return "mild"

_PALETTE = [
    (255, 80,  80),  (80, 200, 80),  (80, 120, 255), (255, 180,  0),
    (200,  0, 200), (0, 210, 210),  (255, 100,  0),  (0, 150, 255),
    (180, 255,  0), (255,  0, 120), (120,  0, 255),  (0, 255, 150),
    (255, 220, 50), (50, 255, 220), (200, 100, 50),  (100, 200, 50),
    (50, 100, 200), (200, 50, 100), (150, 150, 150),
]

def analyze_skin_image_core(image_bytes: bytes):
    session = get_ort_session()
    input_name = session.get_inputs()[0].name
    input_tensor = preprocess_image(image_bytes)
    outputs = session.run(None, {input_name: input_tensor})
    
    disease_name, confidence = parse_yolo_output(outputs[0])
    severity = determine_severity_str(confidence, disease_name)
    disease_info = get_disease_info(disease_name)
    
    recommendations = []
    if "lifestyle" in disease_info.get("recommendations", {}):
        recommendations.extend(disease_info["recommendations"]["lifestyle"])
        
    suggested_medicines = [m["name"] for m in disease_info.get("medicines", [])]
    precautions = disease_info.get("recommendations", {}).get("precautions", [])
    
    return {
        "condition": disease_name,
        "confidence": confidence,
        "severity": severity,
        "description": disease_info["description"],
        "recommendations": recommendations,
        "suggestedMedicines": suggested_medicines,
        "precautions": precautions,
        "doctorConsultationUrgent": severity in ["severe", "critical"],
        "boundingBox": None # Optional
    }