import io
import json
import os
import asyncio
from contextlib import asynccontextmanager
from typing import Optional, List

from fastapi import FastAPI, File, Form, UploadFile, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse, JSONResponse
from PIL import Image
import uvicorn

from vit_explainer import ViTExplanationService

# Initialize explanation service
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
service = ViTExplanationService(base_dir=BASE_DIR)
simulate_lock = asyncio.Lock()

def run_inference_task(*args, **kwargs):
    return service.run_simulation(*args, **kwargs)



@asynccontextmanager
async def lifespan(app: FastAPI):
    # Preload both ViT models into memory at startup
    print("Preloading ViT models into RAM...")
    try:
        service.get_model("vit_base_patch16_224")
        print("✓ ViT-Base (224) loaded into RAM.")
    except Exception as e:
        print(f"Notice: Failed to preload ViT-Base: {e}")

    try:
        service.get_model("vit_small_patch16_224")
        print("✓ ViT-Small (224) loaded into RAM.")
    except Exception as e:
        print(f"Notice: Failed to preload ViT-Small: {e}")

    print("All models ready in RAM.")
    yield
    print("Shutting down API server...")


app = FastAPI(
    title="ViT Explanation Comparison API",
    description="Interactive evaluation for F2HF, BeyondIntuition, Integrated Gradients, Grad-CAM, and Attention Rollout",
    version="1.0.0",
    lifespan=lifespan
)

# Enable CORS for flexible hosting
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.get("/health")
@app.get("/api/health")
async def health_check():
    """Ultra-lightweight ping endpoint to keep cloud containers awake without running inference."""
    return {"status": "ok", "message": "Service is healthy"}


@app.get("/api/classes")
async def get_classes():
    """Returns the list of all 1000 ImageNet classes for search and selection."""
    classes = service.get_classes_list()
    return {"classes": classes, "total": len(classes)}


@app.post("/api/simulate")
async def simulate(
    file: UploadFile = File(...),
    class_query: str = Form("0"),
    query_type: str = Form("auto"),
    model_name: str = Form("vit_base_patch16_224"),
    methods: Optional[str] = Form(None)
):
    """
    Runs model inference and selected explanation methods on the uploaded image.
    Uses an asyncio lock so simultaneous requests queue cleanly in memory.
    """
    if not file.content_type.startswith("image/"):
        raise HTTPException(status_code=400, detail="Uploaded file is not a valid image.")

    try:
        image_bytes = await file.read()
        pil_image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Cannot decode image file: {str(e)}")

    # Parse requested methods
    selected_methods: Optional[List[str]] = None
    if methods:
        try:
            # Check if JSON array string
            if methods.strip().startswith("["):
                selected_methods = json.loads(methods)
            else:
                selected_methods = [m.strip().lower() for m in methods.split(",") if m.strip()]
        except Exception:
            selected_methods = None

    # Thread-safe sequential processing through the shared in-memory model
    async with simulate_lock:
        try:
            loop = asyncio.get_event_loop()
            result = await loop.run_in_executor(
                None,
                run_inference_task,
                pil_image,
                class_query,
                query_type,
                model_name,
                selected_methods
            )
            return JSONResponse(content={"status": "success", "data": result})
        except Exception as e:
            import traceback
            traceback.print_exc()
            raise HTTPException(status_code=500, detail=f"Explanation calculation failed: {str(e)}")


# Serve frontend static assets
static_dir = os.path.join(BASE_DIR, "static")
os.makedirs(static_dir, exist_ok=True)


@app.get("/example_map.json")
async def serve_example_map():
    map_file = os.path.join(BASE_DIR, "example_map.json")
    if os.path.exists(map_file):
        return FileResponse(map_file)
    alt_file = os.path.join(BASE_DIR, "example", "example_map.json")
    if os.path.exists(alt_file):
        return FileResponse(alt_file)
    raise HTTPException(status_code=404, detail="example_map.json not found")


@app.get("/")
async def serve_index():
    index_file = os.path.join(static_dir, "index.html")
    if os.path.exists(index_file):
        return FileResponse(index_file)
    return JSONResponse(
        content={"message": "ViT Explanation API is running. Frontend static/index.html not yet created."},
        status_code=200
    )


app.mount("/static", StaticFiles(directory=static_dir), name="static")

example_dir = os.path.join(BASE_DIR, "example")
if os.path.exists(example_dir):
    app.mount("/example", StaticFiles(directory=example_dir), name="example")

if __name__ == "__main__":
    uvicorn.run("api:app", host="0.0.0.0", port=8000, reload=False)