"""
World Cup Storyteller — Hugging Face Spaces app.py
NLP Homework 4 — ARI 525

Upload this file + requirements.txt + your 3 CSVs to your HF Space.
Set GROQ_API_KEY as a Space secret in Settings.
"""

import os
import time
import warnings
warnings.filterwarnings("ignore")

import pandas as pd
import numpy as np
import gradio as gr
from groq import Groq
from sentence_transformers import SentenceTransformer
import faiss

# -----------------------------------------------
# 1. CONFIG
# -----------------------------------------------

# Loaded from HF Space secret (never hardcode this)
GROQ_API_KEY = os.environ.get("GROQ_API_KEY", "")
MODEL = "llama-3.1-8b-instant"

client = Groq(api_key=GROQ_API_KEY)

# -----------------------------------------------
# 2. LOAD DATA
# CSVs must be uploaded to your HF Space root folder
# -----------------------------------------------

DATA_PATH = "./WC_data/"  # Same folder as app.py on HF Space

try:
    cups    = pd.read_csv(DATA_PATH + "WorldCups.csv")
    matches = pd.read_csv(DATA_PATH + "WorldCupMatches.csv")
    players = pd.read_csv(DATA_PATH + "WorldCupPlayers.csv")
    print("✅ Data loaded successfully.")
except FileNotFoundError as e:
    raise RuntimeError(
        "CSV files not found. Make sure WorldCups.csv, WorldCupMatches.csv, "
        "and WorldCupPlayers.csv are uploaded to your HF Space root folder."
    ) from e


# -----------------------------------------------
# 3. DATA CLEANING
# -----------------------------------------------

def clean_data(cups, matches):
    matches_clean = matches.dropna(
        subset=["Home Team Name", "Away Team Name", "Home Team Goals", "Away Team Goals"]
    ).copy()
    matches_clean["Home Team Name"] = matches_clean["Home Team Name"].str.strip()
    matches_clean["Away Team Name"]  = matches_clean["Away Team Name"].str.strip()
    matches_clean["Home Team Goals"] = matches_clean["Home Team Goals"].astype(int)
    matches_clean["Away Team Goals"] = matches_clean["Away Team Goals"].astype(int)

    cups_clean = cups.dropna(subset=["Year", "Winner"]).copy()
    cups_clean["Year"] = cups_clean["Year"].astype(int)
    return cups_clean, matches_clean


cups_clean, matches_clean = clean_data(cups, matches)

available_years = sorted(cups_clean["Year"].unique().tolist())
available_teams = sorted(set(
    matches_clean["Home Team Name"].tolist() +
    matches_clean["Away Team Name"].tolist()
))


# -----------------------------------------------
# 4. CONTEXT BUILDERS
# -----------------------------------------------

def build_team_context(team, year, cups, matches):
    cup_info = cups[cups["Year"] == year]
    if cup_info.empty:
        return None, f"No data found for year {year}."

    cup = cup_info.iloc[0]
    team_matches = matches[
        (matches["Year"] == year) &
        ((matches["Home Team Name"] == team) | (matches["Away Team Name"] == team))
    ].sort_values("Stage")

    if team_matches.empty:
        return None, f"{team} did not participate in the {year} World Cup."

    match_lines = []
    for _, row in team_matches.iterrows():
        home, away = row["Home Team Name"], row["Away Team Name"]
        hg, ag = int(row["Home Team Goals"]), int(row["Away Team Goals"])
        if home == team:
            result = "WIN" if hg > ag else ("DRAW" if hg == ag else "LOSS")
            line = f"[{row['Stage']}] {team} vs {away}: {hg}-{ag} ({result})"
        else:
            result = "WIN" if ag > hg else ("DRAW" if hg == ag else "LOSS")
            line = f"[{row['Stage']}] {team} vs {home}: {ag}-{hg} ({result})"
        match_lines.append(line)

    final_note = ""
    if cup.get("Winner") == team:
        final_note = f"{team} WON the {year} World Cup! 🏆"
    elif cup.get("Runners-Up") == team:
        final_note = f"{team} were runners-up in {year}."
    elif cup.get("Third") == team:
        final_note = f"{team} finished third in {year}."

    context = (
        f"{team}'s Journey — World Cup {year}\n"
        f"Host: {cup.get('Country', 'Unknown')}\n"
        f"{final_note}\n\n"
        f"Match-by-match results:\n" +
        "\n".join(match_lines)
    )
    return context, None


def build_edition_context(year, cups, matches):
    cup_info = cups[cups["Year"] == year]
    if cup_info.empty:
        return None, f"No data found for year {year}."

    cup = cup_info.iloc[0]
    edition_matches = matches[matches["Year"] == year].sort_values("Stage")

    match_lines = [
        f"[{row['Stage']}] {row['Home Team Name']} {int(row['Home Team Goals'])} "
        f"- {int(row['Away Team Goals'])} {row['Away Team Name']}"
        for _, row in edition_matches.iterrows()
    ]

    context = (
        f"World Cup {year} — Host: {cup.get('Country', 'Unknown')}\n"
        f"Winner: {cup.get('Winner', 'Unknown')}\n"
        f"Runners-up: {cup.get('Runners-Up', 'Unknown')}\n"
        f"Third Place: {cup.get('Third', 'Unknown')}\n"
        f"Goals Scored: {cup.get('GoalsScored', 'Unknown')}\n"
        f"Teams: {cup.get('QualifiedTeams', 'Unknown')}\n"
        f"Attendance: {cup.get('Attendance', 'Unknown')}\n\n"
        f"All Matches:\n" +
        "\n".join(match_lines)
    )
    return context, None


# -----------------------------------------------
# 5. PROMPTS
# -----------------------------------------------

SYSTEM_PROMPT = """
You are a passionate, knowledgeable sports journalist and storyteller specializing in
FIFA World Cup history. Your job is to turn raw match data into vivid, engaging,
narrative-driven stories about World Cup tournaments and team journeys.

Guidelines:
- Write in a natural, flowing narrative style (not bullet points)
- Make the story feel alive — build tension, highlight dramatic moments
- Use the match data accurately — never invent scores or results
- Adapt your tone to the user's request (documentary, dramatic, casual, etc.)
- Keep the story between 250-400 words unless asked otherwise
""".strip()

RAG_SYSTEM_PROMPT = """
You are a passionate sports journalist specializing in FIFA World Cup history.
You will be given retrieved historical World Cup data as context, followed by
a user request. Use the retrieved data to write an accurate, vivid, engaging
narrative story. Only use information present in the retrieved context.

Guidelines:
- Write in flowing narrative prose (not bullet points)
- Build tension and highlight drama
- Stick strictly to the facts in the retrieved data
- Adapt tone to user's request
- Keep the story between 250-400 words unless asked otherwise
""".strip()

FEW_SHOT_EXAMPLES = [
    {
        "context": (
            "France's Journey — World Cup 1998\nHost: France\n"
            "France WON the 1998 World Cup! 🏆\n\nMatch-by-match results:\n"
            "[Group Stage] France vs South Africa: 3-0 (WIN)\n"
            "[Group Stage] France vs Saudi Arabia: 4-0 (WIN)\n"
            "[Group Stage] France vs Denmark: 2-1 (WIN)\n"
            "[Round of 16] France vs Paraguay: 1-0 (WIN)\n"
            "[Quarter-finals] France vs Italy: 0-0 (WIN via penalties)\n"
            "[Semi-finals] France vs Croatia: 2-1 (WIN)\n"
            "[Final] France vs Brazil: 3-0 (WIN)"
        ),
        "story": (
            "It was the summer that France found its destiny on home soil. Les Bleus entered "
            "the 1998 World Cup as hosts with immense pressure, but from the very first "
            "whistle, they played with a quiet, relentless authority.\n\n"
            "The group stage was a statement — South Africa swept aside 3-0, Saudi Arabia "
            "dismantled 4-0, and Denmark edged out 2-1. By the knockout rounds, France "
            "carried the weight of a nation's expectations into every match.\n\n"
            "Paraguay made them suffer — a lone goal in extra time was all that separated "
            "the sides. Italy pushed them to penalties, a nerve-shredding duel that France "
            "survived with ice-cold nerves. Then came Croatia — France trailed before turning "
            "the game on its head to win 2-1.\n\n"
            "The final against defending champions Brazil became a coronation. Three goals, "
            "zero reply. The Stade de France erupted, and a generation of French children "
            "discovered what it felt like to be champions of the world."
        )
    },
    {
        "context": (
            "West Germany's Journey — World Cup 1954\nHost: Switzerland\n"
            "West Germany WON the 1954 World Cup! 🏆\n\nMatch-by-match results:\n"
            "[Group Stage] West Germany vs Turkey: 4-1 (WIN)\n"
            "[Group Stage] West Germany vs Hungary: 3-8 (LOSS)\n"
            "[Group Stage Playoff] West Germany vs Turkey: 7-2 (WIN)\n"
            "[Quarter-finals] West Germany vs Yugoslavia: 2-0 (WIN)\n"
            "[Semi-finals] West Germany vs Austria: 6-1 (WIN)\n"
            "[Final] West Germany vs Hungary: 3-2 (WIN)"
        ),
        "story": (
            "They called it the Miracle of Bern, and for good reason. No one believed West "
            "Germany could win the 1954 World Cup — least of all after Hungary handed them "
            "an 8-3 humiliation in the group stage.\n\n"
            "But West Germany, crafty and resilient, rested key players for that match and "
            "quietly plotted their path to the final. They dispatched Turkey twice, squeezed "
            "past Yugoslavia, then demolished Austria 6-1 in a dazzling semi-final.\n\n"
            "The final was a rematch nobody expected. Hungary — the Mighty Magyars, unbeaten "
            "for four years — led 2-0 within eight minutes. The world assumed it was over.\n\n"
            "It was not. West Germany clawed back to 2-2, and then, six minutes from the end, "
            "Helmut Rahn struck. 3-2. A country still rebuilding from the rubble of war had "
            "become world champions."
        )
    }
]


# -----------------------------------------------
# 6. GENERATION FUNCTIONS
# -----------------------------------------------

def generate_zeroshot(user_prompt, context):
    msg = f"Here is the World Cup data:\n---\n{context}\n---\n\nUser request: {user_prompt}"
    start = time.time()
    response = client.chat.completions.create(
        model=MODEL,
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user",   "content": msg}
        ],
        temperature=0.8,
        max_tokens=600
    )
    elapsed = round(time.time() - start, 2)
    return response.choices[0].message.content.strip(), elapsed


def generate_fewshot(user_prompt, context):
    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    for ex in FEW_SHOT_EXAMPLES:
        messages.append({
            "role": "user",
            "content": f"Here is the World Cup data:\n---\n{ex['context']}\n---\n\nUser request: Tell the story of this team's World Cup journey."
        })
        messages.append({"role": "assistant", "content": ex["story"]})
    messages.append({
        "role": "user",
        "content": f"Here is the World Cup data:\n---\n{context}\n---\n\nUser request: {user_prompt}"
    })
    start = time.time()
    response = client.chat.completions.create(
        model=MODEL, messages=messages, temperature=0.8, max_tokens=600
    )
    elapsed = round(time.time() - start, 2)
    return response.choices[0].message.content.strip(), elapsed


# RAG setup
print("⏳ Loading embedding model...")
embedder = SentenceTransformer("all-MiniLM-L6-v2")

def build_knowledge_base(cups, matches):
    chunks, metadata = [], []
    for _, cup_row in cups.iterrows():
        year = int(cup_row["Year"])
        year_matches = matches[matches["Year"] == year]
        teams = set(
            year_matches["Home Team Name"].tolist() +
            year_matches["Away Team Name"].tolist()
        )
        for team in teams:
            ctx, err = build_team_context(team, year, cups, matches)
            if ctx:
                chunks.append(ctx)
                metadata.append({"team": team, "year": year})
        ctx, err = build_edition_context(year, cups, matches)
        if ctx:
            chunks.append(ctx)
            metadata.append({"team": "ALL", "year": year})
    return chunks, metadata

print("⏳ Building knowledge base...")
kb_chunks, kb_metadata = build_knowledge_base(cups_clean, matches_clean)

print("⏳ Embedding knowledge base...")
kb_embeddings = embedder.encode(kb_chunks, show_progress_bar=False, convert_to_numpy=True)
dim = kb_embeddings.shape[1]
faiss_index = faiss.IndexFlatL2(dim)
faiss_index.add(kb_embeddings)
print(f"✅ RAG index ready — {faiss_index.ntotal} vectors")


def retrieve_context(query, top_k=3):
    query_vec = embedder.encode([query], convert_to_numpy=True)
    distances, indices = faiss_index.search(query_vec, top_k)
    retrieved = []
    for idx in indices[0]:
        if idx < len(kb_chunks):
            meta = kb_metadata[idx]
            retrieved.append(f"[Retrieved: {meta['team']} — {meta['year']}]\n{kb_chunks[idx]}")
    return "\n\n---\n\n".join(retrieved)


def generate_rag(user_prompt):
    retrieved = retrieve_context(user_prompt, top_k=3)
    msg = f"Retrieved World Cup data:\n---\n{retrieved}\n---\n\nUser request: {user_prompt}"
    start = time.time()
    response = client.chat.completions.create(
        model=MODEL,
        messages=[
            {"role": "system", "content": RAG_SYSTEM_PROMPT},
            {"role": "user",   "content": msg}
        ],
        temperature=0.8,
        max_tokens=600
    )
    elapsed = round(time.time() - start, 2)
    return response.choices[0].message.content.strip(), elapsed


# -----------------------------------------------
# 7. GRADIO UI
# -----------------------------------------------

def generate_story_ui(mode, year, team, user_prompt, approach):
    year = int(year)

    if mode == "Team Journey":
        if not team:
            return "⚠️ Please select a team.", ""
        context, err = build_team_context(team, year, cups_clean, matches_clean)
    else:
        context, err = build_edition_context(year, cups_clean, matches_clean)

    if err:
        return f"⚠️ {err}", ""

    if not user_prompt.strip():
        if mode == "Team Journey":
            user_prompt = f"Tell me {team}'s {year} World Cup story in a dramatic, engaging way."
        else:
            user_prompt = f"Tell the full story of the {year} World Cup — the drama, the upsets, the champion."

    try:
        if approach == "Zero-Shot":
            story, elapsed = generate_zeroshot(user_prompt, context)
            info = f"⚡ Zero-Shot | ⏱️ {elapsed}s | 📝 {len(story.split())} words"
        elif approach == "Few-Shot":
            story, elapsed = generate_fewshot(user_prompt, context)
            info = f"📖 Few-Shot | ⏱️ {elapsed}s | 📝 {len(story.split())} words"
        else:
            story, elapsed = generate_rag(user_prompt)
            info = f"🔍 RAG | ⏱️ {elapsed}s | 📝 {len(story.split())} words"
        return story, info
    except Exception as e:
        return f"❌ Error: {str(e)}", ""


year_choices = [str(y) for y in available_years]

with gr.Blocks(
    title="⚽ World Cup Storyteller",
    theme=gr.themes.Base(),
    css="""
        #header { text-align: center; padding: 1.5em 0 0.5em 0; }
        #header h1 { font-size: 2.2em; margin-bottom: 0.1em; }
        #header p  { color: #888; font-size: 1.05em; }
        #story-box textarea { font-size: 1.05em; line-height: 1.8; }
        #info-bar { font-size: 0.9em; color: #555; margin-top: 0.3em; }
        .approach-note { font-size: 0.85em; color: #777; margin-top: 0.4em; }
    """
) as demo:

    with gr.Column(elem_id="header"):
        gr.Markdown("# ⚽ World Cup Storyteller")
        gr.Markdown("Generate vivid, narrative-driven stories about any World Cup edition or team journey.")

    with gr.Row():

        # --- Left panel: controls ---
        with gr.Column(scale=1, min_width=280):
            gr.Markdown("### ⚙️ Settings")

            mode = gr.Radio(
                choices=["Team Journey", "Full Edition"],
                value="Team Journey",
                label="Storytelling Mode"
            )
            year = gr.Dropdown(
                choices=year_choices,
                value="2002",
                label="World Cup Year"
            )
            team = gr.Dropdown(
                choices=available_teams,
                value="Brazil",
                label="Team (Team Journey only)"
            )
            approach = gr.Radio(
                choices=["Zero-Shot", "Few-Shot", "RAG"],
                value="Few-Shot",
                label="NLP Approach"
            )
            gr.Markdown(
                "- **Zero-Shot** — No examples, direct generation\n"
                "- **Few-Shot** — Guided by hand-crafted story examples\n"
                "- **RAG** — Retrieves context from full knowledge base",
                elem_classes="approach-note"
            )

        # --- Right panel: prompt + output ---
        with gr.Column(scale=2):
            gr.Markdown("### ✍️ Your Prompt")
            user_prompt = gr.Textbox(
                placeholder='e.g. "Tell Brazil\'s 2002 story like a sports documentary" — or leave blank for a default story.',
                label="Free-form prompt (optional)",
                lines=3
            )
            btn = gr.Button("🎙️ Generate Story", variant="primary", size="lg")

            gr.Markdown("### 📖 Story")
            story_out = gr.Textbox(
                label="",
                lines=16,
                interactive=False,
                elem_id="story-box"
            )
            info_out = gr.Markdown("", elem_id="info-bar")

    btn.click(
        fn=generate_story_ui,
        inputs=[mode, year, team, user_prompt, approach],
        outputs=[story_out, info_out]
    )

    gr.Markdown(
        "---\n*Data: FIFA World Cup dataset (1930–2014) · Model: Llama 3 8B via Groq · "
        "Embeddings: all-MiniLM-L6-v2 · Built for NLP HW4 — ARI 525*"
    )

demo.launch()