"""
Epicure Explorer API
Wraps the three Epicure ingredient embedding operations:
  - /neighbors  (substitution + experimental combos)
  - /slerp      (cuisine compass)
  - /modes      (flavour profile)
  - /ingredients (autocomplete vocabulary)
"""

import sys
import os
import importlib.util
from pathlib import Path
from contextlib import asynccontextmanager
from typing import Optional, List

import numpy as np

from fastapi import FastAPI, HTTPException, Query, Body
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse
from huggingface_hub import snapshot_download
from pydantic import BaseModel

MODEL_ID = os.getenv("EPICURE_MODEL", "Kaikaku/epicure-core")

models = {}
ingredients_list = []

def _unit(v: np.ndarray, axis: int = -1, eps: float = 1e-9) -> np.ndarray:
    n = np.linalg.norm(v, axis=axis, keepdims=True)
    return v / np.maximum(n, eps)

def _load_epicure_module():
    snap = snapshot_download(repo_id=MODEL_ID)
    loader_path = Path(snap) / "epicure.py"
    if not loader_path.exists():
        raise RuntimeError(f"epicure.py not found in snapshot: {snap}")
    spec = importlib.util.spec_from_file_location("epicure_lib", loader_path)
    mod = importlib.util.module_from_spec(spec)
    sys.modules["epicure_lib"] = mod
    spec.loader.exec_module(mod)
    return mod

def load_model():
    global models, ingredients_list
    mod = _load_epicure_module()
    for repo_id in ["Kaikaku/epicure-core", "Kaikaku/epicure-cooc", "Kaikaku/epicure-chem"]:
        name = repo_id.split("-")[-1]
        print(f"[epicure] Loading model: {repo_id} ...", flush=True)
        models[name] = mod.Epicure.from_pretrained(repo_id)
    ingredients_list = sorted(list(models["core"].vocab))
    print(f"[epicure] All models ready. Vocab size: {len(ingredients_list)}", flush=True)


@asynccontextmanager
async def lifespan(app: FastAPI):
    load_model()
    yield


app = FastAPI(title="Epicure Explorer API", version="1.0.0", lifespan=lifespan)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)


# ── Health ────────────────────────────────────────────────────────────────────

@app.get("/api/health")
def health():
    return {"status": "ok", "model": MODEL_ID, "vocab_size": len(ingredients_list)}


# ── Vocabulary ────────────────────────────────────────────────────────────────

@app.get("/api/ingredients")
def list_ingredients(q: Optional[str] = Query(None, description="Filter by prefix")):
    """All canonical ingredient names, optionally filtered by prefix."""
    items = ingredients_list
    if q:
        q_lower = q.lower()
        items = [i for i in items if q_lower in i.lower()]
    return {"count": len(items), "ingredients": items}


# ── Neighbors (substitution + experimental combos) ────────────────────────────

@app.get("/api/neighbors")
def neighbors(
    ingredient: str = Query(..., description="Canonical ingredient, e.g. 'chicken'"),
    k: int = Query(10, ge=1, le=30),
    model_variant: str = Query("core", description="core | cooc | chem"),
):
    """
    Top-K nearest neighbours by cosine similarity.
    model_variant controls which embedding space is used:
      - core  = blend (chemistry + recipe context)
      - cooc  = recipe co-occurrence only
      - chem  = chemistry/FlavorDB only
    """
    if ingredient not in ingredients_list:
        raise HTTPException(404, f"'{ingredient}' not in vocabulary. Try /api/ingredients?q=...")
    try:
        results = models["core"].neighbors(ingredient, k=k)
        return {
            "ingredient": ingredient,
            "model": MODEL_ID,
            "neighbors": [{"name": n, "score": round(s, 4)} for n, s in results],
        }
    except Exception as e:
        raise HTTPException(500, str(e))


# ── Cuisine Compass (SLERP) ───────────────────────────────────────────────────

CUISINE_POLES = [
    "South_Asian", "East_Asian", "Southeast_Asian",
    "French", "Italian", "Mediterranean",
    "Mexican", "American", "Middle_Eastern",
    "West_African", "Japanese", "Chinese",
]

@app.get("/api/cuisines")
def list_cuisines():
    return {"cuisines": CUISINE_POLES}


@app.get("/api/slerp")
def slerp(
    ingredient: str = Query(..., description="Starting ingredient, e.g. 'rice'"),
    direction: str = Query(..., description="Cuisine name (e.g. 'South_Asian') or ingredient"),
    theta: float = Query(30.0, ge=5.0, le=85.0, description="Degrees to rotate (5–85)"),
    k: int = Query(10, ge=1, le=30),
):
    """
    SLERP direction arithmetic.
    Rotate 'ingredient' toward a cuisine pole or another ingredient by theta degrees.
    Returns the k nearest neighbours of the resulting vector.
    """
    if ingredient not in ingredients_list:
        raise HTTPException(404, f"'{ingredient}' not in vocabulary.")
    try:
        # cuisine poles are prefixed with "cuisine:" in the model
        if direction in CUISINE_POLES:
            target = f"cuisine:{direction}"
        else:
            target = direction  # ingredient-to-ingredient
        results = models["core"].slerp(ingredient, target, theta_deg=theta, k=k)
        return {
            "ingredient": ingredient,
            "direction": direction,
            "theta_deg": theta,
            "suggestions": [{"name": n, "score": round(s, 4)} for n, s in results],
        }
    except Exception as e:
        raise HTTPException(400, f"SLERP failed: {e}")


# ── Flavour Modes ─────────────────────────────────────────────────────────────

@app.get("/api/modes")
def modes(
    ingredient: str = Query(..., description="Ingredient to profile"),
    kind: str = Query("factor", description="'factor' or 'supervised'"),
    k: int = Query(5, ge=1, le=15),
):
    """Find the closest flavour/aroma mode clusters for an ingredient."""
    if ingredient not in ingredients_list:
        raise HTTPException(404, f"'{ingredient}' not in vocabulary.")
    try:
        results = models["core"].closest_mode(ingredient, kind=kind, k=k)
        # results is list of (mode_id, description, score)
        return {
            "ingredient": ingredient,
            "kind": kind,
            "modes": [
                {"id": r[0], "description": r[1], "score": round(r[2], 4)}
                for r in results
            ],
        }
    except Exception as e:
        raise HTTPException(400, f"Mode lookup failed: {e}")


# ── Graph (force-directed neighbourhood) ──────────────────────────────────────

@app.get("/api/graph")
def graph(
    ingredient: str = Query(..., description="Center ingredient"),
    k: int = Query(15, ge=1, le=30),
    model_variant: str = Query("core", description="core | cooc | chem"),
    cross_edge_threshold: float = Query(0.45, ge=0.0, le=1.0),
):
    """
    Return a force-directed graph of an ingredient's neighbourhood.
    Includes center->neighbour edges and cross-edges between neighbours
    whose cosine similarity exceeds the threshold.
    """
    if ingredient not in ingredients_list:
        raise HTTPException(404, f"'{ingredient}' not in vocabulary.")
    m = models.get(model_variant, models["core"])
    try:
        neighbors = m.neighbors(ingredient, k=k)
        node_names = [ingredient] + [n for n, _ in neighbors]
        name_to_idx = {n: i for i, n in enumerate(node_names)}
        # Get vectors for all nodes
        vectors = np.stack([m.vec(n) for n in node_names])
        # Compute pairwise similarities
        sims = vectors @ vectors.T  # (k+1, k+1)
        nodes = [
            {"id": name, "score": 1.0 if i == 0 else float(sims[0, i]), "is_center": i == 0}
            for i, name in enumerate(node_names)
        ]
        edges = []
        # Center edges
        for i in range(1, len(node_names)):
            edges.append({"source": ingredient, "target": node_names[i], "weight": float(sims[0, i])})
        # Cross-edges between neighbours
        for i in range(1, len(node_names)):
            for j in range(i + 1, len(node_names)):
                w = float(sims[i, j])
                if w >= cross_edge_threshold:
                    edges.append({"source": node_names[i], "target": node_names[j], "weight": w})
        return {"center": ingredient, "model": model_variant, "nodes": nodes, "edges": edges}
    except Exception as e:
        raise HTTPException(500, str(e))


# ── Recipe Lab ────────────────────────────────────────────────────────────────

class RecipeRequest(BaseModel):
    ingredients: List[str]
    action: str  # score | suggest | affinity | surprise


@app.post("/api/recipe")
def recipe(req: RecipeRequest):
    """
    Recipe lab operations:
      - score:      mean pairwise coherence of the basket
      - suggest:    best ingredient to add next
      - affinity:   cuisine-pole affinity of the basket centroid
      - surprise:   unexpected but chemistry-compatible ingredient
    """
    basket = [i for i in req.ingredients if i in ingredients_list]
    if not basket:
        raise HTTPException(400, "No valid ingredients in basket.")
    m = models["core"]
    try:
        if req.action == "score":
            vecs = np.stack([m.vec(i) for i in basket])
            sims = vecs @ vecs.T
            pairs = []
            total = 0.0
            count = 0
            for i in range(len(basket)):
                for j in range(i + 1, len(basket)):
                    s = float(sims[i, j])
                    pairs.append({"a": basket[i], "b": basket[j], "score": round(s, 4)})
                    total += s
                    count += 1
            mean_score = total / count if count else 0.0
            pairs.sort(key=lambda x: -x["score"])
            label = "coherent" if mean_score >= 0.45 else "mixed" if mean_score >= 0.30 else "adventurous"
            return {"score": round(mean_score, 4), "label": label, "pairwise": pairs}

        elif req.action == "suggest":
            basket_vecs = np.stack([m.vec(i) for i in basket])
            best_name = None
            best_score = -1.0
            for cand in ingredients_list:
                if cand in basket:
                    continue
                cand_vec = m.vec(cand)
                avg_sim = float(np.mean(basket_vecs @ cand_vec))
                if avg_sim > best_score:
                    best_score = avg_sim
                    best_name = cand
            reason = f"complements {', '.join(basket[:3])}"
            if len(basket) > 3:
                reason += f" and {len(basket) - 3} more"
            return {"suggestion": best_name, "score": round(best_score, 4), "reason": reason}

        elif req.action == "affinity":
            basket_vec = _unit(np.mean(np.stack([m.vec(i) for i in basket]), axis=0))
            affinities = []
            for pole_name in CUISINE_POLES:
                pole_key = f"cuisine:{pole_name}"
                if pole_key in m.supervised_poles:
                    pole_vec = _unit(m.supervised_poles[pole_key])
                    affinities.append({"cuisine": pole_name, "score": round(float(basket_vec @ pole_vec), 4)})
            affinities.sort(key=lambda x: -x["score"])
            return {"affinities": affinities[:6]}

        elif req.action == "surprise":
            chem = models.get("chem")
            if chem is None:
                raise HTTPException(503, "Chem model not loaded")
            basket_vec = _unit(np.mean(np.stack([m.vec(i) for i in basket]), axis=0))
            best_name = None
            best_chem_score = -1.0
            best_core_score = 0.0
            for cand in ingredients_list:
                if cand in basket:
                    continue
                core_score = float(basket_vec @ m.vec(cand))
                chem_score = float(basket_vec @ chem.vec(cand))
                # Want low core similarity but high chem similarity
                if core_score < 0.35 and chem_score > best_chem_score and chem_score > 0.40:
                    best_chem_score = chem_score
                    best_core_score = core_score
                    best_name = cand
            if best_name is None:
                # Fallback: medium core similarity ingredient not in basket
                for cand in ingredients_list:
                    if cand in basket:
                        continue
                    core_score = float(basket_vec @ m.vec(cand))
                    if 0.25 <= core_score <= 0.40 and core_score > best_core_score:
                        best_core_score = core_score
                        best_name = cand
                best_chem_score = float(basket_vec @ chem.vec(best_name)) if best_name else 0.0
            return {"suggestion": best_name, "core_score": round(best_core_score, 4), "chemistry_score": round(best_chem_score, 4)}

        else:
            raise HTTPException(400, f"Unknown action: {req.action}")
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(500, str(e))


# ── Static frontend ───────────────────────────────────────────────────────────

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

@app.get("/")
def root():
    return FileResponse("static/index.html")