import os, sys
from dotenv import load_dotenv
load_dotenv()
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from download_assets import download_assets
download_assets()

import gradio as gr
from fastapi.middleware.cors import CORSMiddleware
from fastapi import File, UploadFile, HTTPException
from fastapi.responses import JSONResponse
from pydantic import BaseModel
from typing import Optional
from PIL import Image
import io, base64, json

from src.predict import predict_from_bytes
from src.explain import explain_from_bytes
from src.llm_explain import explain_prediction, answer_question
from config import CLASS_LABELS

def run_predict(image_bytes, do_explain, do_llm):
    result = predict_from_bytes(image_bytes)
    heatmap_b64 = None
    if do_explain:
        try:
            pred_idx = CLASS_LABELS[result["predicted_class"]]
            heatmap_b64 = explain_from_bytes(image_bytes, pred_idx)
        except: pass
    explanation = None
    if do_llm:
        try:
            explanation = explain_prediction(
                predicted_class=result["predicted_class"],
                confidence=result["confidence"],
                probabilities=result["probabilities"],
            )
        except Exception as e:
            explanation = str(e)
    return result, heatmap_b64, explanation

def gradio_predict(image: Image.Image):
    buf = io.BytesIO()
    image.save(buf, format="JPEG")
    result, heatmap_b64, explanation = run_predict(buf.getvalue(), True, True)
    heatmap_img = Image.open(io.BytesIO(base64.b64decode(heatmap_b64))) if heatmap_b64 else None
    return heatmap_img, {result["predicted_class"]: result["confidence"]}, explanation or ""

with gr.Blocks(title="DermAI") as demo:
    gr.Markdown("# DermAI — Skin Lesion Classifier\nFrontend: [GitHub Pages](https://ranjithtkm445-blip.github.io/skin-lesion-ai)")
    with gr.Row():
        inp = gr.Image(type="pil", label="Upload Image")
        with gr.Column():
            heatmap = gr.Image(label="Grad-CAM")
            label_out = gr.Label(label="Prediction")
    explanation_out = gr.Textbox(label="Clinical Explanation", lines=6)
    btn = gr.Button("Analyze", variant="primary")
    btn.click(gradio_predict, inputs=inp, outputs=[heatmap, label_out, explanation_out])

# Mount custom API routes on Gradio internal FastAPI app
app = demo.app
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])

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

@app.post("/predict")
async def predict(file: UploadFile = File(...), explain: bool = True, llm_explain: bool = True):
    if not file.content_type.startswith("image/"):
        raise HTTPException(status_code=400, detail="File must be an image.")
    image_bytes = await file.read()
    result, heatmap_b64, explanation = run_predict(image_bytes, explain, llm_explain)
    return {
        "predicted_class": result["predicted_class"],
        "class_name": result["class_name"],
        "confidence": result["confidence"],
        "probabilities": result["probabilities"],
        "heatmap_base64": heatmap_b64,
        "explanation": explanation,
    }

@app.post("/explain")
async def ask(question: str, predicted_class: Optional[str] = None):
    try:
        answer = answer_question(question, predicted_class)
        return {"answer": answer}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)