import os
from pathlib import Path
from typing import Optional, Tuple, List, Dict

import gradio as gr
import pandas as pd
import numpy as np
import plotly.express as px
import joblib

# ZeroGPU + models
import spaces
import torch
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, pipeline
from huggingface_hub import InferenceClient

# ------------------------
# Config & storage
# ------------------------
DATA_DIR = Path("data"); DATA_DIR.mkdir(exist_ok=True)
TS_FMT = "%Y-%m-%d %H:%M:%S"

DT_PATH = "./decision_tree_regressor.joblib"
decision_tree_regressor = joblib.load(DT_PATH)

# Local lightweight model (fallback)
GEN_MODEL = "google/flan-t5-small"
_tokenizer = AutoTokenizer.from_pretrained(GEN_MODEL)
_model = AutoModelForSeq2SeqLM.from_pretrained(GEN_MODEL)
_generate_cpu = pipeline("text2text-generation", model=_model, tokenizer=_tokenizer, device=-1)

# HF Inference API SOTA models
SOTA_MODELS = [
    "Qwen/Qwen2.5-72B-Instruct",            # default: high-quality open model available on HF
    "meta-llama/Meta-Llama-3.1-70B-Instruct",
    "mistralai/Mistral-Nemo-Instruct-2407",
    "Qwen/Qwen2.5-32B-Instruct",
    "Qwen/Qwen2.5-7B-Instruct"
]

def _hf_client(model_id: str) -> InferenceClient:
    token = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACEHUB_API_TOKEN")
    return InferenceClient(model=model_id, token=token, timeout=120)

def generate_with_hf_inference(prompt: str, model_id: str, max_new_tokens: int = 900) -> str:
    """
    Serverless generation via Hugging Face Inference API.
    Works on CPU-only Spaces and with ZeroGPU.
    """
    try:
        client = _hf_client(model_id)
        text = client.text_generation(
            prompt,
            max_new_tokens=max_new_tokens,
            temperature=0.6,
            top_p=0.9,
            repetition_penalty=1.05,
            stop=["</s>"],
            return_full_text=False,
        )
        return text.strip()
    except Exception as e:
        # Fall back to local tiny model inside a GPU window if available
        return f"(HF Inference error: {e})\n" + generate_on_gpu(prompt, max_new_tokens=min(max_new_tokens, 600))

# ------------------------
# ZeroGPU functions (presence at import satisfies ZeroGPU)
# ------------------------
@spaces.GPU
def generate_on_gpu(prompt: str, max_new_tokens: int = 600) -> str:
    """
    Generate with tiny local model. If CUDA is available in the ZeroGPU window,
    bind pipeline to GPU; otherwise use CPU.
    """
    try:
        if torch.cuda.is_available():
            gen = pipeline("text2text-generation", model=_model.to("cuda"), tokenizer=_tokenizer, device=0)
            out = gen(prompt, max_new_tokens=max_new_tokens)
        else:
            out = _generate_cpu(prompt, max_new_tokens=max_new_tokens)
        return out[0]["generated_text"].strip()
    except Exception as e:
        out = _generate_cpu(prompt, max_new_tokens=max_new_tokens)
        return out[0]["generated_text"].strip() + f"\n\n(Note: GPU path failed: {e})"

# ------------------------
# Metrics & helpers
# ------------------------
ACTIVITY = {"Sedentary":1.2,"Lightly active":1.375,"Moderately active":1.55,"Very active":1.725,"Athlete":1.9}
GOAL_CAL_ADJ = {"Fat loss":-0.15,"Recomp/Maintenance":0.0,"Muscle gain":0.10}

def bmi(w,h): return w/((h/100)**2)
def bmr_mifflin(sex,w,h,a): return 10*w+6.25*h-5*a+(5 if sex=="Male" else -161)
def tdee(bmr,act): return bmr*ACTIVITY.get(act,1.2)

def parse_hhmm(hhmm: str) -> Tuple[int,int]:
    h, m = hhmm.split(":")
    h = int(h); m = int(m)
    if not (0 <= h <= 23 and 0 <= m <= 59):
        raise ValueError("Time must be HH:MM (24h).")
    return h, m

def fmt_hhmm(h: int, m: int) -> str:
    return f"{h:02d}:{m:02d}"

# Meal ideas, workouts, etc.
DIET_STYLES = ["Mediterranean", "Omnivore", "Vegetarian", "Vegan", "Low-carb"]
MEAL_IDEAS = {
    "Mediterranean": [
        "Oats + dates + walnuts + olive oil",
        "Grilled fish, lentil salad, greens",
        "Hummus, wholegrain pita, veggies",
        "Chickpea tomato stew",
        "Feta & olive salad, quinoa",
        "Shakshuka + side salad",
        "Lentils, roasted veg, tahini"
    ],
    "Omnivore": [
        "Yogurt + berries + nuts",
        "Chicken bowl (rice, veg, olive oil)",
        "Eggs, avocado, sourdough",
        "Salmon, quinoa, asparagus",
        "Lean beef, sweet potato, salad",
        "Tuna whole-grain wrap",
        "Cottage cheese + fruit + seeds"
    ],
    "Vegetarian": [
        "Tofu scramble, toast, avocado",
        "Paneer tikka bowl",
        "Bean chili + brown rice",
        "Halloumi, couscous, veg",
        "Greek salad + eggs",
        "Tempeh stir-fry",
        "Yogurt parfait + granola"
    ],
    "Vegan": [
        "Tofu scramble, avocado toast",
        "Lentil curry + basmati",
        "Burrito bowl (beans, corn, salsa)",
        "Seitan, roasted potatoes, veg",
        "Tofu poke bowl",
        "Chickpea pasta + marinara",
        "Overnight oats + banana + PB"
    ],
    "Low-carb": [
        "Eggs, smoked salmon, salad",
        "Chicken Caesar (no croutons)",
        "Beef & greens stir-fry",
        "Omelette + veg + cheese",
        "Zoodles + turkey bolognese",
        "Tofu salad w/ tahini",
        "Yogurt + nuts (moderate)"
    ]
}
WORKOUTS = {
    "Fat loss": [
        "3× LISS cardio 30–40min",
        "2× full-body strength 45min",
        "1× intervals 12–16min",
        "Daily 8–10k steps"
    ],
    "Recomp/Maintenance": [
        "3× full-body strength 45–60min",
        "1–2× LISS cardio 30min",
        "Mobility 10min daily",
        "8–10k steps"
    ],
    "Muscle gain": [
        "4× strength split 45–60min",
        "Optional 1× LISS 20–30min",
        "Mobility 10min",
        "7–9k steps"
    ]
}

def feeding_schedule(first_meal_hhmm: str, fasting_hours: float) -> List[Tuple[str, str]]:
    h, m = parse_hhmm(first_meal_hhmm)
    window = max(0.0, 24 - float(fasting_hours))
    start_minutes = h * 60 + m
    end_minutes = int((start_minutes + window * 60) % (24 * 60))
    sched = []
    for _ in range(7):
        start = fmt_hhmm(h, m)
        end = fmt_hhmm(end_minutes // 60, end_minutes % 60)
        sched.append((start, end))
    return sched

def weekly_plan(diet: str, sched: List[Tuple[str, str]], kcal: int, protein_g: int) -> pd.DataFrame:
    ideas = MEAL_IDEAS[diet]
    rows = []
    for i in range(7):
        day = ["Mon","Tue","Wed","Thu","Fri","Sat","Sun"][i]
        start, end = sched[i]
        meal1 = ideas[i % len(ideas)]
        meal2 = ideas[(i+3) % len(ideas)]
        snack = "Fruit or nuts (optional)"
        rows.append({
            "Day": day,
            "Feeding window": f"{start}–{end}",
            "Meal 1": meal1,
            "Meal 2": meal2,
            "Protein target": f"≥ {protein_g} g",
            "Daily kcal": kcal,
            "Snack": snack,
        })
    return pd.DataFrame(rows)

def shopping_list(diet: str) -> List[str]:
    core = [
        "Leafy greens, mixed veg, berries",
        "Olive oil, nuts/seeds, herbs & spices",
        "Coffee/tea, mineral water, electrolytes",
    ]
    extras = {
        "Omnivore": ["Chicken, fish, eggs, yogurt, cottage cheese", "Rice/quinoa/sourdough", "Beans/lentils"],
        "Mediterranean": ["Fish, feta, olives", "Whole grains (bulgur, farro)", "Chickpeas/lentils"],
        "Vegetarian": ["Eggs, dairy, paneer", "Legumes", "Tofu/tempeh"],
        "Vegan": ["Tofu/tempeh/seitan", "Beans/lentils", "Plant yogurt/milk"],
        "Low-carb": ["Eggs, fish, meat", "Green veg", "Greek yogurt, cheese"],
    }
    return core + extras[diet]

# ------------------------
# Plan builder (with SOTA + local fallback)
# ------------------------
def predict_and_plan(
    fasting_duration, meal_timing, weight, age, gender, height,
    activity, goal, diet, lang, use_sota_model, sota_model_id
) -> Tuple[Optional[float], str, str, pd.DataFrame, object, str]:
    try:
        if fasting_duration < 0 or fasting_duration > 72: raise ValueError("Fasting must be 0–72h.")
        h, m = parse_hhmm(meal_timing)
        if weight <= 0 or height <= 0 or age < 0: raise ValueError("Invalid weight/height/age.")

        # Predict score
        df = pd.DataFrame({
            "Fasting Duration (hours)": [float(fasting_duration)],
            "Meal Timing (hour:minute)": [h + m/60],
            "Body Weight (kg)": [float(weight)],
            "Age (years)": [float(age)],
            "Height (cm)": [float(height)],
            "Gender_Male": [1 if gender == "Male" else 0],
            "Gender_Other": [1 if gender == "Other" else 0],
        })
        score = float(decision_tree_regressor.predict(df)[0])

        # Metrics
        bmr = bmr_mifflin(gender, weight, height, age)
        tdee_kcal = tdee(bmr, activity)
        target_kcal = int(round(tdee_kcal * (1 + GOAL_CAL_ADJ[goal])))
        protein_g = int(round(max(1.6 * weight, 80)))
        bmi_val = round(bmi(weight, height), 1)

        # Schedule, plan table, chart
        sched = feeding_schedule(meal_timing, float(fasting_duration))
        plan_df = weekly_plan(diet, sched, target_kcal, protein_g)

        chart_df = pd.DataFrame({
            "Day": ["Mon","Tue","Wed","Thu","Fri","Sat","Sun"],
            "start": [int(s.split(":")[0])*60 + int(s.split(":")[1]) for s,_ in sched],
            "length": [max(0, int((24 - float(fasting_duration))*60))]*7,
        })
        fig = px.bar(chart_df, y="Day", x="length", base="start", orientation="h",
                     title="Feeding window each day (minutes)")
        fig.update_layout(
            xaxis=dict(range=[0,1440], tickvals=[0,360,720,1080,1440],
                       ticktext=["00:00","06:00","12:00","18:00","24:00"]),
            height=300, margin=dict(l=10,r=10,t=40,b=10)
        )

        # Base markdown (deterministic, structured). We’ll optionally enhance with SOTA.
        kpis = (
            f"**Score:** {score:.1f}  •  **BMI:** {bmi_val}  •  **BMR:** {int(bmr)} kcal  •  "
            f"**TDEE:** {int(tdee_kcal)} kcal  •  **Target:** {target_kcal} kcal  •  **Protein:** ≥ {protein_g} g  •  "
            f"**Diet:** {diet}"
        )
        sched_md = "\n".join([f"- **{d}**: {s} – {e}" for d,(s,e) in zip(["Mon","Tue","Wed","Thu","Fri","Sat","Sun"], sched)])
        workouts_md = "\n".join([f"- {w}" for w in WORKOUTS[goal]])
        shop_md = "\n".join([f"- {x}" for x in shopping_list(diet)])

        base_plan_md = f"""
## Your 7-day intermittent fasting plan

{kpis}

### Feeding window (daily)
{sched_md}

### Weekly training
{workouts_md}

### Daily meals (example week)
(See the table below.)

### Shopping list
{shop_md}

> Hydration & electrolytes during the fast, protein at each meal, whole foods, and 7–9 hours sleep.
""".strip()

        # Enhance/format with chosen generator
        if use_sota_model:
            plan_md = generate_with_hf_inference(
                prompt=(
                    "You are an expert health coach. Refine the following intermittent fasting plan. "
                    "Keep markdown headings and bullets; be concise and specific; keep the meaning. "
                    f"Language: '{lang}'.\n\n{base_plan_md}"
                ),
                model_id=sota_model_id,
                max_new_tokens=900,
            )
        else:
            # Use local tiny model inside ZeroGPU window (or CPU fallback)
            plan_md = generate_on_gpu(
                "Rewrite in a friendly coaching tone; keep markdown structure; do not remove tables or metrics.\n\n" + base_plan_md,
                max_new_tokens=700
            )

        # Save for download
        md_path = DATA_DIR / "plan.md"
        md_path.write_text(plan_md, encoding="utf-8")

        return score, kpis, plan_md, plan_df, fig, str(md_path)

    except Exception as e:
        return None, "", f"⚠️ {e}", pd.DataFrame(), None, ""

# ------------------------
# Tracker logic
# ------------------------
active_fasts: Dict[str, pd.Timestamp] = {}

def _csv(user: str) -> Path:
    safe = "".join(ch for ch in (user or "default") if ch.isalnum() or ch in ("_","-"))
    return DATA_DIR / f"{safe}.csv"

def hist_load(user: str) -> pd.DataFrame:
    p = _csv(user)
    if p.exists():
        d = pd.read_csv(p)
        for c in ["start_time","end_time"]:
            if c in d: d[c] = pd.to_datetime(d[c], errors="coerce")
        return d
    return pd.DataFrame(columns=["start_time","end_time","duration_hours","note"])

def hist_save(user: str, d: pd.DataFrame):
    d.to_csv(_csv(user), index=False)

def make_hist_chart(df: pd.DataFrame):
    if df.empty: return None
    d = df.dropna(subset=["end_time"]).copy()
    if d.empty: return None
    d["date"] = pd.to_datetime(d["end_time"]).dt.date
    fig = px.bar(d, x="date", y="duration_hours", title="Fasting duration by day (h)")
    fig.update_layout(height=300, margin=dict(l=10,r=10,t=40,b=10))
    return fig

def compute_streak(df: pd.DataFrame) -> int:
    d = df.dropna(subset=["end_time"]).copy()
    if d.empty: return 0
    days = set(pd.to_datetime(d["end_time"]).dt.date)
    cur = pd.Timestamp.now().date(); streak = 0
    while cur in days:
        streak += 1; cur = cur - pd.Timedelta(days=1)
    return streak

def hist_stats(df: pd.DataFrame) -> str:
    if df.empty: return "No history yet."
    last7 = df.tail(7)
    avg = last7["duration_hours"].mean()
    streak = compute_streak(df)
    return f"Total fasts: {len(df)}\nAvg (last 7): {avg:.2f} h\nCurrent streak: {streak} day(s)"

def start_fast(user: str, note: str):
    if not user: return "Enter username in Tracker.", None
    if user in active_fasts: return f"Already fasting since {active_fasts[user].strftime(TS_FMT)}.", None
    active_fasts[user] = pd.Timestamp.now()
    return f"✅ Fast started at {active_fasts[user].strftime(TS_FMT)}.", None

def end_fast(user: str):
    if not user: return "Enter username.", None, None, None
    if user not in active_fasts: return "No active fast.", None, None, None
    end = pd.Timestamp.now(); start = active_fasts.pop(user)
    dur = round((end - start).total_seconds()/3600, 2)
    df = hist_load(user)
    df.loc[len(df)] = [start, end, dur, ""]
    hist_save(user, df)
    return f"✅ Fast ended at {end.strftime(TS_FMT)} • {dur} h", df.tail(12), make_hist_chart(df), hist_stats(df)

def refresh_hist(user: str):
    df = hist_load(user)
    return df.tail(12), make_hist_chart(df), hist_stats(df)

# ------------------------
# UI
# ------------------------
with gr.Blocks(
    title="Intermittent Fasting Coach — Pro (SOTA)",
    theme=gr.themes.Soft(primary_hue=gr.themes.colors.orange, neutral_hue=gr.themes.colors.gray),
) as demo:
    gr.Markdown("""
# 🥣 Intermittent Fasting — Pro (SOTA)
Detailed coaching plans + tracker. ZeroGPU-ready (with CPU fallback). Data stored locally in this Space.
""")

    with gr.Tabs():
        # --- Coach tab
        with gr.TabItem("Coach"):
            with gr.Row():
                with gr.Column():
                    fasting_duration = gr.Number(label="Fasting duration (h)", value=16, minimum=0, maximum=72, step=0.5)
                    meal_timing = gr.Textbox(label="First meal time (HH:MM)", value="12:30")
                    weight = gr.Number(label="Body weight (kg)", value=70, step=0.5)
                with gr.Column():
                    age = gr.Slider(label="Age (years)", minimum=18, maximum=100, value=35)
                    gender = gr.Radio(["Male","Female","Other"], label="Gender", value="Male")
                    height = gr.Number(label="Height (cm)", value=175)

            with gr.Row():
                activity = gr.Dropdown(choices=list(ACTIVITY.keys()), value="Lightly active", label="Activity")
                goal = gr.Dropdown(choices=list(GOAL_CAL_ADJ.keys()), value="Recomp/Maintenance", label="Goal")
                diet = gr.Dropdown(choices=DIET_STYLES, value="Mediterranean", label="Diet style")
                lang = gr.Radio(["en","es"], value="en", label="Language")
                use_sota_model = gr.Checkbox(value=True, label="Use SOTA model (HF Inference)")
                sota_model_id = gr.Dropdown(choices=SOTA_MODELS, value=SOTA_MODELS[0], label="HF model")

            btn = gr.Button("Predict & Build Plan", variant="primary")

            score_out = gr.Number(label="Predicted score")
            kpi_out = gr.Markdown()
            plan_md = gr.Markdown()
            plan_tbl = gr.Dataframe(headers=["Day","Feeding window","Meal 1","Meal 2","Protein target","Daily kcal","Snack"], interactive=False)
            fig = gr.Plot()
            dl = gr.DownloadButton(label="Download plan (.md)")

            btn.click(
                predict_and_plan,
                inputs=[fasting_duration, meal_timing, weight, age, gender, height, activity, goal, diet, lang, use_sota_model, sota_model_id],
                outputs=[score_out, kpi_out, plan_md, plan_tbl, fig, dl],
                api_name="coach_plan"
            )

        # --- Tracker tab
        with gr.TabItem("Tracker"):
            with gr.Row():
                user = gr.Textbox(label="Username", value="")
                note = gr.Textbox(label="Note (optional)")
            with gr.Row():
                b1 = gr.Button("Start fast", variant="primary")
                b2 = gr.Button("End fast")
                b3 = gr.Button("Reload history")
            status = gr.Markdown("Not fasting.")
            hist = gr.Dataframe(interactive=False)
            hist_fig = gr.Plot()
            stats = gr.Markdown()

            b1.click(start_fast, inputs=[user, note], outputs=[status, note])
            b2.click(end_fast, inputs=[user], outputs=[status, hist, hist_fig, stats])  # <-- FIXED: no None
            b3.click(refresh_hist, inputs=[user], outputs=[hist, hist_fig, stats])
            demo.load(refresh_hist, inputs=[user], outputs=[hist, hist_fig, stats])

        # --- About tab
        with gr.TabItem("About"):
            gr.Markdown("""
**How it works**  
• The predictor estimates a health score from inputs.  
• The coach builds a 7-day schedule matching your fasting window, goal, activity and diet style.  
• SOTA option uses Hugging Face Inference API; fallback uses a tiny local model in the ZeroGPU window.  
• Tracker stores CSVs under `/data/` and never sends data elsewhere.
""")

if __name__ == "__main__":
    demo.queue().launch()