"""
TripAlchemy - Gradio Application (Part 5 of the assignment)
================================================================

Flow: pick a city -> set vibe sliders (+ optional filters / free text) ->
get the top-5 matched real experiences (embeddings + slider hybrid, via
scripts/recsys.py) -> generate one AI-written itinerary narrative weaving
together the top-3 of them (via scripts/narrative_generator.py).

Data sourcing (per the assignment's constraints):
  - In production (the deployed HF Space) this reads the dataset straight
    from the HF Dataset repo and the winning embeddings from the HF Model
    repo - see HF_DATASET_REPO / HF_EMBEDDINGS_REPO below. Set those and
    flip USE_HF_HUB = True before deploying.
  - For local development it falls back to the local files in ../data/,
    which is what runs by default right now.

Generation backend:
  - Uses narrative_generator.hf_inference_api_generate_fn, which calls a
    small instruction-tuned HF model through the HF Inference API. Needs an
    HF_TOKEN (set as an HF Space "Secret" in production, or exported
    locally). Falls back to the mock generator with a visible warning if no
    token is present, so the UI still works end-to-end for local testing.
"""

import os
import sys
import json
import html
from pathlib import Path

import gradio as gr

_esc = html.escape

# Local repo layout keeps recsys/narrative_generator in ../scripts; the deployed
# HF Space uploads them next to app.py, so fall back to this file's own directory.
SCRIPT_DIR = Path(__file__).parent.parent / "scripts"
if not (SCRIPT_DIR / "recsys.py").exists():
    SCRIPT_DIR = Path(__file__).parent
sys.path.insert(0, str(SCRIPT_DIR))

from recsys import RecommendationEngine, CATEGORIES  # noqa: E402
from narrative_generator import (  # noqa: E402
    generate_itinerary,
    hf_inference_api_generate_fn,
    local_causal_generate_fn,
    mock_generate_fn,
)

# =====================================================
# Config - flip these before deploying to HF Space
# =====================================================
# Configured via environment variables / HF Space "Variables & secrets" - no code
# editing needed to deploy. Set USE_HF_HUB=true + the two repo ids as Space Variables,
# and HF_TOKEN as a Space Secret, and the app reads dataset + embeddings from the Hub.
USE_HF_HUB = os.environ.get("USE_HF_HUB", "false").lower() == "true"
HF_DATASET_REPO = os.environ.get("HF_DATASET_REPO", "your-username/tripalchemy-experiences")
HF_EMBEDDINGS_REPO = os.environ.get("HF_EMBEDDINGS_REPO", "your-username/tripalchemy-embeddings")
GENERATION_MODEL_ID = os.environ.get("GENERATION_MODEL_ID", "Qwen/Qwen2.5-7B-Instruct")

DATA_DIR = Path(__file__).parent.parent / "data"
CITIES_FILE = DATA_DIR / "cities.json"
EXPERIENCES_FILE = DATA_DIR / "experiences.json"
EMBEDDINGS_FILE = DATA_DIR / "embeddings" / "winning_model.npz"

CATEGORY_LABELS = {
    "culinary": "🍽️ Culinary",
    "historical": "🏛️ Historical",
    "shopping": "🛍️ Shopping",
    "nature": "🌲 Nature",
    "nightlife": "🌃 Nightlife",
    "art_culture": "🎨 Art & Culture",
}

COUNTRY_FLAGS = {
    "Argentina": "🇦🇷", "Australia": "🇦🇺", "Bulgaria": "🇧🇬", "Czech Republic": "🇨🇿",
    "France": "🇫🇷", "Germany": "🇩🇪", "Greece": "🇬🇷", "Hungary": "🇭🇺", "Iceland": "🇮🇸",
    "Indonesia": "🇮🇩", "Italy": "🇮🇹", "Japan": "🇯🇵", "Mexico": "🇲🇽", "Morocco": "🇲🇦",
    "Netherlands": "🇳🇱", "New Zealand": "🇳🇿", "Portugal": "🇵🇹", "Singapore": "🇸🇬",
    "South Africa": "🇿🇦", "South Korea": "🇰🇷", "Spain": "🇪🇸", "Thailand": "🇹🇭",
    "Turkey": "🇹🇷", "UAE": "🇦🇪", "USA": "🇺🇸", "United Kingdom": "🇬🇧", "Vietnam": "🇻🇳",
}


def _flag(country: str) -> str:
    return COUNTRY_FLAGS.get(country, "🌍")

QUICK_STARTERS = {
    "🌹 Romantic Lisbon Weekend": {
        "city": "lisbon",
        "sliders": {"culinary": 0.8, "historical": 0.3, "shopping": 0.1, "nature": 0.2, "nightlife": 0.5, "art_culture": 0.5},
        "budget": 150, "energy": "low",
    },
    "🍜 Solo Tokyo Foodie Crawl": {
        "city": "tokyo",
        "sliders": {"culinary": 0.95, "historical": 0.1, "shopping": 0.3, "nature": 0.0, "nightlife": 0.6, "art_culture": 0.2},
        "budget": 80, "energy": "moderate",
    },
    "🏛️ Budget Rome History Dive": {
        "city": "rome",
        "sliders": {"culinary": 0.3, "historical": 0.95, "shopping": 0.0, "nature": 0.1, "nightlife": 0.0, "art_culture": 0.6},
        "budget": 40, "energy": "moderate",
    },
}


def load_cities():
    if USE_HF_HUB:
        from huggingface_hub import hf_hub_download
        path = hf_hub_download(repo_id=HF_DATASET_REPO, filename="cities.json", repo_type="dataset")
    else:
        path = CITIES_FILE
    with open(path, encoding="utf-8") as f:
        data = json.load(f)
    return {c["id"]: f"{_flag(c['country'])}  {c['name']}, {c['country']}" for c in data["cities"]}


def build_embed_fn(report_path=None):
    """Loads the SAME embedding model that produced winning_model.npz, so
    free-text / slider-sentence queries land in the same vector space as the
    precomputed document embeddings. Falls back to slider-only search
    (no text embedding) if sentence-transformers or the network isn't
    available, e.g. when testing inside a restricted sandbox."""
    report_path = Path(report_path) if report_path else DATA_DIR / "embeddings" / "comparison_report.json"
    hf_id = "sentence-transformers/all-MiniLM-L6-v2"  # sane default
    query_prefix = ""  # some models (E5) require a "query: " prefix at query time
    if report_path.exists():
        report = json.loads(report_path.read_text(encoding="utf-8"))
        winner = report.get("winner")
        if winner and winner in report.get("models", {}):
            hf_id = report["models"][winner]["hf_id"]
            query_prefix = report["models"][winner].get("prefix_query", "")

    try:
        from sentence_transformers import SentenceTransformer
        model = SentenceTransformer(hf_id)

        def _embed(texts):
            # apply the winner's query prefix so queries match the precomputed
            # document embeddings (which used the model's doc prefix).
            texts = [query_prefix + t for t in texts]
            return model.encode(texts, normalize_embeddings=True, show_progress_bar=False)

        print(f"✅ embed_fn ready ({hf_id}, query_prefix={query_prefix!r})")
        return _embed
    except Exception as e:
        print(f"⚠️ Could not load {hf_id} ({e}). Falling back to slider-only search "
              f"(no free-text embedding search this session).")
        return None


def get_engine():
    if USE_HF_HUB:
        # Production path: pull straight from the HF Hub (dataset repo for experiences,
        # model repo for the winning embeddings + the comparison report that tells the
        # app which embedding model + query prefix to use).
        from huggingface_hub import hf_hub_download
        exp_path = hf_hub_download(repo_id=HF_DATASET_REPO, filename="experiences.json", repo_type="dataset")
        emb_path = hf_hub_download(repo_id=HF_EMBEDDINGS_REPO, filename="winning_model.npz", repo_type="model")
        report_path = hf_hub_download(repo_id=HF_EMBEDDINGS_REPO, filename="comparison_report.json", repo_type="model")
    else:
        exp_path, emb_path = str(EXPERIENCES_FILE), str(EMBEDDINGS_FILE)
        report_path = str(DATA_DIR / "embeddings" / "comparison_report.json")
    return RecommendationEngine.load(exp_path, emb_path, embed_fn=build_embed_fn(report_path))


def get_generate_fn():
    """Real HF-model generation with a reliable fallback chain:
    1. HF Inference API (fast, needs a token + the model to be served by a provider)
    2. a local HF model loaded in the Space (transformers, CPU) - always works, no token
    Neither is the banned random/formula generation; the mock is a last resort only."""
    token = os.environ.get("HF_TOKEN")
    api_fn = None
    if token:
        try:
            api_fn = hf_inference_api_generate_fn(model_id=GENERATION_MODEL_ID, token=token)
        except Exception as e:
            print(f"⚠️ HF Inference API init failed: {e}")
    _local = {}

    def _generate(prompt: str) -> str:
        if api_fn is not None:
            try:
                return api_fn(prompt)
            except Exception as e:
                print(f"⚠️ HF Inference API call failed ({type(e).__name__}); using local model fallback.")
        if "fn" not in _local:
            print("Loading local generation model (Qwen2.5-0.5B-Instruct)...")
            _local["fn"] = local_causal_generate_fn("Qwen/Qwen2.5-0.5B-Instruct")
        return _local["fn"](prompt)

    return _generate


CITIES = load_cities()
ENGINE = get_engine()
GENERATE_FN = get_generate_fn()


def format_card(exp: dict) -> str:
    scores = exp["category_scores"]
    top_cats = [(c, v) for c, v in sorted(scores.items(), key=lambda kv: -kv[1])[:3] if v > 0]
    chips = "".join(f"<span class='cat-chip'>{CATEGORY_LABELS[c]} {v:.1f}</span>" for c, v in top_cats)
    tags = "".join(f"<span class='vibe-tag'>{_esc(t)}</span>" for t in exp.get("vibe_tags", [])[:4])
    nb = exp.get("neighborhood")
    loc = f"{_esc(nb)}, {_esc(exp['city_name'])}" if nb else _esc(exp["city_name"])
    return (
        f"<div class='exp-card'>"
        f"<div class='exp-title'>{_esc(exp.get('title', 'Untitled'))}</div>"
        f"<div class='exp-meta'>{_flag(exp.get('country',''))} {loc}"
        f" &nbsp;•&nbsp; 💰 ${exp.get('cost_usd','?')} &nbsp;•&nbsp; ⏱️ {exp.get('duration_hours','?')}h"
        f" &nbsp;•&nbsp; {_esc(str(exp.get('time_of_day', 'flexible')))}</div>"
        f"<div class='exp-desc'>{_esc(exp.get('description',''))}</div>"
        f"<div class='chip-row'>{chips}{tags}</div>"
        f"</div>"
    )


def run_search(city_id, culinary, historical, shopping, nature, nightlife, art_culture,
                free_text, max_budget, energy):
    if not city_id:
        return "⚠️ Pick a city first.", "", ""

    weights = {
        "culinary": culinary, "historical": historical, "shopping": shopping,
        "nature": nature, "nightlife": nightlife, "art_culture": art_culture,
    }
    energy_filter = None if energy == "any" else energy
    can_use_text = ENGINE.embed_fn is not None

    results = ENGINE.recommend(
        category_weights=weights,
        free_text=free_text or None,
        city_id=city_id,
        max_budget_usd=max_budget if max_budget else None,
        energy=energy_filter,
        alpha=(0.5 if free_text else 0.2) if can_use_text else 0.0,
        use_text_similarity=can_use_text,
        top_k=5,
    )

    if not results:
        return "No experiences matched those filters — try loosening the budget or sliders.", "", ""

    cards_md = "".join(format_card(e) for e in results)

    # The narrative is secondary: never let a generation failure block the experiences.
    def _template(exps):
        p = [f"<b>{_esc(e['title'])}</b> in {_esc(e.get('neighborhood') or e['city_name'])}" for e in exps]
        s = "Begin your day with " + p[0]
        if len(p) > 1:
            s += ", then head to " + p[1]
        if len(p) > 2:
            s += ", and round it off at " + p[2]
        return s + " — a day woven from the vibe you set."

    try:
        gen = generate_itinerary(results[:3], GENERATE_FN, user_context=free_text or "")
        nar = (gen.get("narrative") or "").strip()
        low = nar.lower()
        # guard against a model that echoes the prompt or returns junk
        echoed = ("travel writer" in low or "real experiences pulled" in low
                  or "curated database" in low or len(nar) < 25)
        narrative_body = _template(results[:3]) if echoed else _esc(nar)
    except Exception:
        narrative_body = _template(results[:3])
    narrative_md = (f"<div class='narr-box'><div class='narr-title'>✨ Your AI-woven itinerary</div>"
                    f"<div class='narr-body'>{narrative_body}</div></div>")

    grounded_md = "<div class='grounded'>🔗 Grounded in: " + ", ".join(_esc(e["title"]) for e in results[:3]) + "</div>"

    return cards_md, narrative_md, grounded_md


def apply_quick_starter(name):
    cfg = QUICK_STARTERS[name]
    s = cfg["sliders"]
    return (
        cfg["city"], s["culinary"], s["historical"], s["shopping"],
        s["nature"], s["nightlife"], s["art_culture"], "",
        cfg["budget"], cfg["energy"],
    )


CUSTOM_CSS = """
/* Force a consistent LIGHT look regardless of the viewer's dark-mode preference,
   by overriding Gradio's theme CSS variables (works in both :root and .dark). */
:root, .dark, gradio-app, .gradio-container {
  --body-background-fill:#f4f2ef !important;
  --background-fill-primary:#ffffff !important;
  --background-fill-secondary:#efedea !important;
  --block-background-fill:#ffffff !important;
  --block-label-background-fill:#ffffff !important;
  --block-title-text-color:#2b3040 !important;
  --block-label-text-color:#3a3f4b !important;
  --body-text-color:#2b3040 !important;
  --body-text-color-subdued:#5a6072 !important;
  --border-color-primary:#e4e2de !important;
  --input-background-fill:#ffffff !important;
  --input-border-color:#e0ded9 !important;
  --panel-background-fill:#ffffff !important;
  --neutral-950:#ffffff !important; --neutral-900:#ffffff !important;
  --neutral-800:#f2f0ed !important; --neutral-700:#e4e2de !important;
}
body, gradio-app, .gradio-container {background:#f4f2ef !important; color:#2b3040 !important;}
.gradio-container {max-width: 1180px !important; margin: auto !important;}
#hero {background: linear-gradient(120deg,#ff6a3d 0%,#f7455d 46%,#8b3dff 100%);
       border-radius: 18px; padding: 30px 34px; margin-bottom: 10px; color:#fff;
       box-shadow: 0 12px 32px rgba(180,60,110,.30);}
#hero h1 {margin:0; font-size:2.35rem; font-weight:800; letter-spacing:-.6px; color:#fff;}
#hero p {margin:8px 0 0; font-size:1.07rem; opacity:.97;}
#hero .pill {display:inline-block; background:rgba(255,255,255,.18); border:1px solid rgba(255,255,255,.4);
       border-radius:999px; padding:4px 13px; font-size:.8rem; margin:12px 8px 0 0;}
.exp-card {background:#ffffff !important; border:1px solid #ececf1; border-left:5px solid #ff6a3d;
       border-radius:14px; padding:16px 18px; margin-bottom:14px;
       box-shadow:0 3px 14px rgba(30,30,60,.10);}
.exp-card, .exp-card * {color:#2b3040 !important;}   /* beat gradio dark-mode white text */
.exp-title {font-size:1.18rem !important; font-weight:800 !important; color:#d63e1f !important; margin-bottom:5px;}
.exp-meta {font-size:.87rem !important; color:#586074 !important; margin-bottom:9px;}
.exp-desc {font-size:.95rem !important; color:#2b3040 !important; line-height:1.55; margin-bottom:11px;}
.chip-row {line-height:2.2;}
.cat-chip {display:inline-block; background:#ffe9df !important; color:#c9451f !important; border-radius:999px;
       padding:3px 11px; font-size:.79rem !important; font-weight:700 !important; margin-right:6px; white-space:nowrap;}
.vibe-tag {display:inline-block; background:#e4f2f1 !important; color:#1c6d69 !important; border-radius:999px;
       padding:3px 10px; font-size:.76rem !important; font-weight:600 !important; margin-right:6px; white-space:nowrap;}
.narr-box {background:linear-gradient(135deg,#fff6ef,#fdeade) !important; border:1px solid #ffd0b5;
       border-radius:14px; padding:18px 22px; margin-top:8px;}
.narr-box, .narr-box * {color:#2b3040 !important;}
.narr-title {font-weight:800 !important; color:#d63e1f !important; font-size:1.14rem !important; margin-bottom:8px;}
.narr-body {font-size:1rem !important; line-height:1.6;}
.grounded {font-size:.8rem !important; color:#7a7f8a !important; margin-top:12px;}
.panel-title {font-weight:800 !important; font-size:1.12rem !important; color:#d63e1f !important; margin:4px 0 8px;}
"""

with gr.Blocks(title="TripAlchemy") as demo:
    gr.HTML(
        f"<style>{CUSTOM_CSS}</style>"
        "<div id='hero'>"
        "<h1>🧪 TripAlchemy</h1>"
        "<p>Set your travel vibe — get real matching experiences + one AI-woven itinerary.</p>"
        "<span class='pill'>🌍 30 cities</span>"
        "<span class='pill'>✨ 10,000+ experiences</span>"
        "<span class='pill'>🧠 embedding + slider hybrid search</span>"
        "</div>"
    )

    with gr.Row():
        with gr.Column(scale=1):
            city_dd = gr.Dropdown(choices=[(v, k) for k, v in CITIES.items()], label="🏙️ Pick a city", value="lisbon")
            gr.Markdown("**Set your vibe** — how strongly do you want each?")
            culinary_s = gr.Slider(0, 1, 0.5, step=0.05, label=CATEGORY_LABELS["culinary"])
            historical_s = gr.Slider(0, 1, 0.5, step=0.05, label=CATEGORY_LABELS["historical"])
            shopping_s = gr.Slider(0, 1, 0.2, step=0.05, label=CATEGORY_LABELS["shopping"])
            nature_s = gr.Slider(0, 1, 0.2, step=0.05, label=CATEGORY_LABELS["nature"])
            nightlife_s = gr.Slider(0, 1, 0.3, step=0.05, label=CATEGORY_LABELS["nightlife"])
            art_s = gr.Slider(0, 1, 0.3, step=0.05, label=CATEGORY_LABELS["art_culture"])
            free_text_tb = gr.Textbox(label="💬 Describe the vibe (optional)",
                                       placeholder="e.g. cozy jazz bar with cocktails")
            budget_s = gr.Slider(0, 500, 200, step=10, label="💰 Max budget per experience (USD)")
            energy_dd = gr.Dropdown(choices=["any", "low", "moderate", "high"], value="any", label="⚡ Energy level")
            search_btn = gr.Button("🔮 Find my experiences", variant="primary", size="lg")

            gr.Markdown("**⚡ Or try a quick starter:**")
            with gr.Row():
                starter_buttons = [gr.Button(name, size="sm") for name in QUICK_STARTERS]

        with gr.Column(scale=2):
            gr.HTML("<div class='panel-title'>🎯 Your matched experiences</div>")
            results_md = gr.HTML()
            narrative_md = gr.HTML()
            grounded_md = gr.HTML()

    inputs = [city_dd, culinary_s, historical_s, shopping_s, nature_s, nightlife_s, art_s,
              free_text_tb, budget_s, energy_dd]
    search_btn.click(fn=run_search, inputs=inputs, outputs=[results_md, narrative_md, grounded_md])

    for name, btn in zip(QUICK_STARTERS, starter_buttons):
        btn.click(fn=(lambda n=name: apply_quick_starter(n)), outputs=inputs).then(
            fn=run_search, inputs=inputs, outputs=[results_md, narrative_md, grounded_md]
        )

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