import json
from pathlib import Path
from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.staticfiles import StaticFiles
from fastapi.responses import HTMLResponse
import torch
from transformers import AutoImageProcessor, AutoModelForImageClassification
from PIL import Image
import io

# ── Paths ──────────────────────────────────────────────────────────────────────
BASE_DIR   = Path(__file__).parent
MODEL_DIR  = BASE_DIR / "models" / "aishrica_food_predictor"
STATIC_DIR = BASE_DIR / "static"
CALORIE_DB = BASE_DIR / "calorie_db.json"

# ── Load calorie database ──────────────────────────────────────────────────────
with open(CALORIE_DB, "r", encoding="utf-8") as f:
    calorie_data = json.load(f)

# ── Load model ─────────────────────────────────────────────────────────────────
def load_model():
    """
    Load the bundled model from ./models/aishrica_food_predictor/ — local only,
    no internet or HuggingFace download required.
    """
    if not (MODEL_DIR.exists() and any(MODEL_DIR.iterdir())):
        raise FileNotFoundError(
            f"Model files not found in {MODEL_DIR}. "
            f"Expected config.json, model.safetensors and preprocessor_config.json."
        )

    print(f"[load] Loading model from: {MODEL_DIR}")
    processor = AutoImageProcessor.from_pretrained(
        str(MODEL_DIR),
        local_files_only=True,
    )
    model = AutoModelForImageClassification.from_pretrained(
        str(MODEL_DIR),
        local_files_only=True,
    )
    model.eval()
    print("[load] Model ready!")
    return processor, model

processor, model = load_model()

# ── FastAPI app ────────────────────────────────────────────────────────────────
app = FastAPI(title="Smart Plate - Indian Food Classification and Calorie Estimation")

STATIC_DIR.mkdir(exist_ok=True)
app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static")


PAGES_DIR = STATIC_DIR / "pages"
VALID_PAGES = {"profile", "home", "scan", "dashboard", "history", "model"}


def _serve_page(name: str) -> HTMLResponse:
    path = PAGES_DIR / f"{name}.html"
    if not path.exists():
        raise HTTPException(status_code=404, detail="Page not found")
    return HTMLResponse(path.read_text(encoding="utf-8"))


@app.get("/", response_class=HTMLResponse)
async def root():
    return _serve_page("welcome")


@app.post("/predict")
async def predict(file: UploadFile = File(...)):
    # Validate file type
    if file.content_type not in ("image/jpeg", "image/png", "image/webp", "image/jpg"):
        raise HTTPException(status_code=400, detail="Only JPEG/PNG/WEBP images are supported.")

    # Read and preprocess image
    contents = await file.read()
    try:
        image = Image.open(io.BytesIO(contents)).convert("RGB")
    except Exception:
        raise HTTPException(status_code=400, detail="Could not read image file.")

    inputs = processor(images=image, return_tensors="pt")

    # Run inference
    with torch.no_grad():
        outputs = model(**inputs)
        logits  = outputs.logits
        probs   = torch.nn.functional.softmax(logits, dim=1)[0]

    # Top-3 predictions
    top3_indices = probs.topk(3).indices.tolist()
    top3_scores  = probs.topk(3).values.tolist()

    results = []
    for idx, score in zip(top3_indices, top3_scores):
        label = model.config.id2label[idx]          # e.g. "biryani"
        info  = calorie_data.get(label, {
            "calories_per_100g": None,
            "serving_g": None,
            "calories_per_serving": None,
            "category": "Unknown"
        })
        results.append({
            "label":                label.replace("_", " ").title(),
            "raw_label":            label,
            "confidence":           round(score * 100, 1),
            "category":             info["category"],
            "calories_per_100g":    info["calories_per_100g"],
            "serving_g":            info["serving_g"],
            "calories_per_serving": info["calories_per_serving"],
        })

    return {
        "predictions": results,
        "top_prediction": results[0],
    }


@app.get("/foods")
async def list_foods():
    """Return all supported food items with their calorie info."""
    return {
        k: v for k, v in sorted(calorie_data.items())
    }


# Page routes — declared last so /predict and /foods take precedence.
@app.get("/{page}", response_class=HTMLResponse)
async def page(page: str):
    if page not in VALID_PAGES:
        raise HTTPException(status_code=404, detail="Page not found")
    return _serve_page(page)


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