import json
import os
import tempfile
import urllib.request
from datetime import date, datetime, timedelta, timezone

import gradio as gr
import pandas as pd
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import seaborn as sns

from simulator import InningOutcomeDistribution, simulate_game

_DATA_DIR  = "/data" if os.path.isdir("/data") else "data"
_HF_REPO_ID = "mc0117/mlb-models-storage"
DATA_PATH  = os.path.join(_DATA_DIR, "full_2026_season.parquet")
_dist_cache: dict = {}


def _ensure_data():
    if not os.path.exists(DATA_PATH):
        from huggingface_hub import download_bucket_files
        os.makedirs(_DATA_DIR, exist_ok=True)
        filename = os.path.basename(DATA_PATH)
        print(f"Downloading {filename} from HF bucket...")
        try:
            download_bucket_files(
                _HF_REPO_ID,
                files=[(filename, DATA_PATH)],
                token=os.environ.get("HF_TOKEN"),
            )
        except Exception:
            print(f"{filename} not found in bucket. Fetching from Baseball Savant...")
            result = update_season_data()
            if result["status"] != "ok":
                raise RuntimeError(f"Auto-fetch failed: {result['message']}")


def _get_dist(team: str) -> InningOutcomeDistribution:
    if team not in _dist_cache:
        _dist_cache[team] = InningOutcomeDistribution(team, DATA_PATH)
    return _dist_cache[team]


def _run_and_aggregate(home_team: str, away_team: str, n_simulations: int):
    _ensure_data()
    home_dist = _get_dist(home_team.upper())
    away_dist = _get_dist(away_team.upper())
    results   = simulate_game(home_dist, away_dist, int(n_simulations))

    rows  = [{"home": h, "away": a, "count": c} for (h, a), c in results.items()]
    df    = pd.DataFrame(rows)
    total = int(df['count'].sum())

    home_wins = int(df[df['home'] > df['away']]['count'].sum())
    away_wins = int(df[df['away'] > df['home']]['count'].sum())
    avg_h = float((df['home'] * df['count']).sum() / total)
    avg_a = float((df['away'] * df['count']).sum() / total)

    return df, total, home_wins, away_wins, avg_h, avg_a


MLB_TEAMS = [
    "ARI", "ATL", "BAL", "BOS", "CHC", "CWS", "CIN", "CLE", "COL", "DET",
    "HOU", "KC",  "LAA", "LAD", "MIA", "MIL", "MIN", "NYM", "NYY", "OAK",
    "PHI", "PIT", "SD",  "SEA", "SF",  "STL", "TB",  "TEX", "TOR", "WSH",
]


# ── Endpoint 1: simulate ─────────────────────────────────────────────────────
# Returns a human-readable summary + heatmap image.
# Gradio UI uses this; Lovable should prefer /simulate_json below.

def simulate(home_team: str, away_team: str, n_simulations: int):
    try:
        df, total, hw, aw, avg_h, avg_a = _run_and_aggregate(home_team, away_team, n_simulations)
    except ValueError as e:
        return str(e), None

    summary = (
        f"{home_team} Win Probability : {hw / total:.1%}\n"
        f"{away_team} Win Probability : {aw / total:.1%}\n"
        f"Expected Score            : {home_team} {avg_h:.1f}  –  {away_team} {avg_a:.1f}\n"
        f"Games Simulated           : {total:,}"
    )

    cap     = 15
    df_plot = df[(df['home'] <= cap) & (df['away'] <= cap)]
    pivot   = df_plot.pivot_table(index="home", columns="away", values="count", aggfunc='sum').fillna(0)

    fig, ax = plt.subplots(figsize=(12, 10))
    sns.heatmap(pivot, annot=True, fmt=".0f", cmap="YlGnBu", ax=ax, annot_kws={"size": 7})
    ax.set_title(f"Score Distribution: {home_team} (Home) vs {away_team} (Away)")
    ax.set_xlabel(f"{away_team} Score")
    ax.set_ylabel(f"{home_team} Score")
    plt.tight_layout()

    return summary, fig


# ── Endpoint 2: simulate_json ────────────────────────────────────────────────
# Returns clean JSON — use this from Lovable / any external client.
# Response shape documented below in the API schema comment.

def simulate_json(home_team: str, away_team: str, n_simulations: int) -> dict:
    """
    Returns:
    {
      "home_team":            str,
      "away_team":            str,
      "n_simulations":        int,
      "home_win_pct":         float,   // 0–1
      "away_win_pct":         float,   // 0–1
      "home_avg_score":       float,
      "away_avg_score":       float,
      "score_distribution": [
        { "home": int, "away": int, "count": int, "probability": float },
        ...                            // sorted by probability desc
      ]
    }
    """
    try:
        df, total, hw, aw, avg_h, avg_a = _run_and_aggregate(home_team, away_team, n_simulations)
    except ValueError as e:
        return {"error": str(e)}

    distribution = (
        df.assign(probability=df['count'] / total)
          .sort_values('probability', ascending=False)
          [['home', 'away', 'count', 'probability']]
          .assign(home=lambda d: d['home'].astype(int),
                  away=lambda d: d['away'].astype(int),
                  count=lambda d: d['count'].astype(int))
          .to_dict(orient='records')
    )

    return {
        "home_team":          home_team.upper(),
        "away_team":          away_team.upper(),
        "n_simulations":      total,
        "home_win_pct":       round(hw / total, 4),
        "away_win_pct":       round(aw / total, 4),
        "home_avg_score":     round(avg_h, 2),
        "away_avg_score":     round(avg_a, 2),
        "score_distribution": distribution,
    }


# MLB Stats API sometimes returns full names instead of abbreviations.
_TEAM_NAME_TO_ABBR = {
    "Arizona Diamondbacks": "ARI", "Atlanta Braves": "ATL",    "Baltimore Orioles": "BAL",
    "Boston Red Sox": "BOS",       "Chicago Cubs": "CHC",       "Chicago White Sox": "CWS",
    "Cincinnati Reds": "CIN",      "Cleveland Guardians": "CLE","Colorado Rockies": "COL",
    "Detroit Tigers": "DET",       "Houston Astros": "HOU",     "Kansas City Royals": "KC",
    "Los Angeles Angels": "LAA",   "Los Angeles Dodgers": "LAD","Miami Marlins": "MIA",
    "Milwaukee Brewers": "MIL",    "Minnesota Twins": "MIN",    "New York Mets": "NYM",
    "New York Yankees": "NYY",     "Oakland Athletics": "OAK",  "Athletics": "OAK",
    "Philadelphia Phillies": "PHI","Pittsburgh Pirates": "PIT", "San Diego Padres": "SD",
    "Seattle Mariners": "SEA",     "San Francisco Giants": "SF","St. Louis Cardinals": "STL",
    "Tampa Bay Rays": "TB",        "Texas Rangers": "TEX",      "Toronto Blue Jays": "TOR",
    "Washington Nationals": "WSH",
    # abbreviation-level overrides (API returns "AZ" for Arizona)
    "AZ": "ARI",
}

_MLB_SCHEDULE_URL  = "https://statsapi.mlb.com/api/v1/schedule?sportId=1&gameType=R&date={date}"

_CACHE_TTL  = timedelta(hours=6)
_mem_cache: dict = {}   # (date_str, n_sims) -> {"cached_at": datetime, "result": dict}


def _cache_hf_path(today: str, n_simulations: int) -> str:
    return f"simulate_today_cache/{today}_{n_simulations}.json"


def _load_cache(today: str, n_simulations: int) -> dict | None:
    key = (today, n_simulations)
    now = datetime.now(timezone.utc)

    entry = _mem_cache.get(key)
    if entry and now - entry["cached_at"] < _CACHE_TTL:
        return entry["result"]

    try:
        from huggingface_hub import download_bucket_files
        remote_path = _cache_hf_path(today, n_simulations)
        with tempfile.TemporaryDirectory() as tmp_dir:
            local_path = os.path.join(tmp_dir, "cache.json")
            download_bucket_files(
                _HF_REPO_ID,
                files=[(remote_path, local_path)],
                token=os.environ.get("HF_TOKEN"),
            )
            with open(local_path) as f:
                stored = json.load(f)
        cached_at = datetime.fromisoformat(stored["cached_at"])
        if cached_at.tzinfo is None:
            cached_at = cached_at.replace(tzinfo=timezone.utc)
        if now - cached_at < _CACHE_TTL:
            _mem_cache[key] = {"cached_at": cached_at, "result": stored["result"]}
            return stored["result"]
    except Exception:
        pass

    return None


def _save_cache(today: str, n_simulations: int, result: dict) -> None:
    key = (today, n_simulations)
    now = datetime.now(timezone.utc)
    _mem_cache[key] = {"cached_at": now, "result": result}

    payload = {"cached_at": now.isoformat(), "result": result}
    tmp_path = None
    try:
        from huggingface_hub import batch_bucket_files
        with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
            json.dump(payload, f)
            tmp_path = f.name
        batch_bucket_files(
            _HF_REPO_ID,
            add=[(tmp_path, _cache_hf_path(today, n_simulations))],
            token=os.environ.get("HF_TOKEN"),
        )
    except Exception as e:
        print(f"Cache upload failed: {e}")
    finally:
        if tmp_path:
            try:
                os.unlink(tmp_path)
            except Exception:
                pass


# ── Endpoint 3 & 4: simulate_today / simulate_today_refresh ─────────────────
# simulate_today      — returns cached result if within 6 h, else recomputes.
# simulate_today_refresh — always recomputes and overwrites the cache.

def _resolve_abbr(info: dict) -> str:
    raw = info.get("abbreviation") or info.get("name", "")
    return _TEAM_NAME_TO_ABBR.get(raw, raw)


def _compute_today(n_simulations: int) -> dict:
    today = date.today().strftime("%Y-%m-%d")
    url   = _MLB_SCHEDULE_URL.format(date=today)
    try:
        with urllib.request.urlopen(url, timeout=10) as resp:
            data = json.loads(resp.read())
    except Exception as e:
        return {"error": f"Failed to fetch MLB schedule: {e}", "date": today, "n_games": 0, "cached": False, "cached_at": None, "games": []}

    raw_games = data.get("dates", [{}])[0].get("games", []) if data.get("dates") else []
    _ensure_data()

    games_out = []
    for game in raw_games:
        home_abbr = _resolve_abbr(game["teams"]["home"]["team"])
        away_abbr = _resolve_abbr(game["teams"]["away"]["team"])
        status    = game.get("status", {}).get("abstractGameState", "Unknown")

        if home_abbr not in MLB_TEAMS or away_abbr not in MLB_TEAMS:
            games_out.append({
                "game_pk":     game.get("gamePk"),
                "game_status": status,
                "home_team":   home_abbr,
                "away_team":   away_abbr,
                "error":       "Team abbreviation not found in simulation data",
            })
            continue

        result = simulate_json(home_abbr, away_abbr, n_simulations)
        result["game_pk"]     = game.get("gamePk")
        result["game_status"] = status
        games_out.append(result)

    core = {"date": today, "n_games": len(games_out), "games": games_out}
    _save_cache(today, n_simulations, core)
    return {**core, "cached": False, "cached_at": None}


def simulate_today(n_simulations: int = 10000) -> dict:
    today  = date.today().strftime("%Y-%m-%d")
    cached = _load_cache(today, n_simulations)
    if cached is not None:
        entry = _mem_cache.get((today, n_simulations))
        return {**cached, "cached": True, "cached_at": entry["cached_at"].isoformat() if entry else None}
    return _compute_today(n_simulations)


def simulate_today_refresh(n_simulations: int = 10000) -> dict:
    today = date.today().strftime("%Y-%m-%d")
    key   = (today, n_simulations)
    _mem_cache.pop(key, None)
    return _compute_today(n_simulations)


# ── Endpoint 4: update_season_data ──────────────────────────────────────────
# Pulls fresh Statcast data for the given season, saves parquet, uploads to HF.

def update_season_data(season_year: int = 2026) -> dict:
    """
    Fetches all Statcast data for `season_year` from Baseball Savant via
    pybaseball, saves a trimmed parquet locally, and uploads it to HF Hub.
    Clears the distribution cache so the next simulation uses fresh data.

    Returns:
    {
      "status":      "ok" | "error",
      "season_year": int,
      "rows":        int,
      "filename":    str,
      "message":     str
    }
    """
    try:
        import pybaseball
        pybaseball.cache.enable()
    except ImportError:
        return {"status": "error", "message": "pybaseball is not installed. Add it to requirements.txt."}

    start_dt = f"{season_year}-03-01"
    end_dt   = date.today().strftime("%Y-%m-%d")
    filename = f"full_{season_year}_season.parquet"
    local_path = os.path.join(_DATA_DIR, filename)

    print(f"Fetching Statcast data {start_dt} → {end_dt} (this takes a few minutes)...")
    try:
        df = pybaseball.statcast(start_dt=start_dt, end_dt=end_dt, verbose=True)
    except Exception as e:
        return {"status": "error", "message": f"pybaseball fetch failed: {e}"}

    keep = ["events", "home_team", "away_team", "inning_topbot", "game_date"]
    df   = df[keep].copy()
    df["home_team"] = df["home_team"].str.upper()
    df["away_team"] = df["away_team"].str.upper()

    os.makedirs(_DATA_DIR, exist_ok=True)
    df.to_parquet(local_path, index=False)
    print(f"Saved {len(df):,} rows to {local_path}")

    try:
        from huggingface_hub import batch_bucket_files
        batch_bucket_files(
            _HF_REPO_ID,
            add=[(local_path, filename)],
            token=os.environ.get("HF_TOKEN"),
        )
    except Exception as e:
        return {"status": "error", "message": f"HF bucket upload failed: {e}", "rows": len(df), "filename": filename}

    global DATA_PATH
    DATA_PATH = local_path
    _dist_cache.clear()

    return {
        "status":      "ok",
        "season_year": season_year,
        "rows":        len(df),
        "filename":    filename,
        "message":     f"Uploaded {filename} ({len(df):,} rows) to HF Hub and refreshed dist cache.",
    }


# ── Gradio app ────────────────────────────────────────────────────────────────

with gr.Blocks(title="game-sim-v0 | MLB Game Simulator") as demo:
    gr.Markdown("## game-sim-v0 — MLB Game Simulator\nPlayer-agnostic Monte Carlo simulation using 2025 Statcast data.")

    with gr.Tab("Simulate"):
        with gr.Row():
            home_dd  = gr.Dropdown(choices=MLB_TEAMS, value="NYY", label="Home Team")
            away_dd  = gr.Dropdown(choices=MLB_TEAMS, value="LAD", label="Away Team")
            n_slider = gr.Slider(minimum=1000, maximum=50000, step=1000, value=10000, label="Simulations")
        btn     = gr.Button("Run Simulation")
        summary = gr.Textbox(label="Results")
        heatmap = gr.Plot(label="Score Distribution Heatmap")
        btn.click(fn=simulate, inputs=[home_dd, away_dd, n_slider], outputs=[summary, heatmap],
                  api_name="simulate")

    with gr.Tab("JSON API"):
        gr.Markdown(
            "Use the `/api/simulate_json` endpoint from your app.\n\n"
            "```\nPOST https://mc0117-mlb-models.hf.space/api/simulate_json\n"
            'Content-Type: application/json\n\n{"data": ["NYY", "LAD", 10000]}\n```'
        )
        with gr.Row():
            j_home = gr.Dropdown(choices=MLB_TEAMS, value="NYY", label="Home Team")
            j_away = gr.Dropdown(choices=MLB_TEAMS, value="LAD", label="Away Team")
            j_n    = gr.Slider(minimum=1000, maximum=50000, step=1000, value=10000, label="Simulations")
        j_btn = gr.Button("Run")
        j_out = gr.JSON(label="Response")
        j_btn.click(fn=simulate_json, inputs=[j_home, j_away, j_n], outputs=j_out,
                    api_name="simulate_json")

    with gr.Tab("Today's Games"):
        gr.Markdown("Fetches today's MLB schedule and simulates every game. Results cached 6 h in HF Hub.")
        t_n = gr.Slider(minimum=1000, maximum=50000, step=1000, value=10000, label="Simulations per game")
        with gr.Row():
            t_btn  = gr.Button("Simulate Today's Games")
            tr_btn = gr.Button("Force Refresh", variant="secondary")
        t_out = gr.JSON(label="Results")
        t_btn.click(fn=simulate_today,         inputs=[t_n], outputs=t_out, api_name="simulate_today")
        tr_btn.click(fn=simulate_today_refresh, inputs=[t_n], outputs=t_out, api_name="simulate_today_refresh")

    with gr.Tab("Update Season Data"):
        gr.Markdown(
            "Pulls fresh Statcast data from Baseball Savant via **pybaseball**, "
            "saves a trimmed parquet, and uploads it to HF Hub. "
            "⚠️ Takes several minutes for a full season."
        )
        u_year = gr.Number(value=2026, label="Season year", precision=0)
        u_btn  = gr.Button("Fetch & Upload")
        u_out  = gr.JSON(label="Status")
        u_btn.click(fn=update_season_data, inputs=[u_year], outputs=u_out,
                    api_name="update_season_data")

    gr.Examples(
        examples=[["NYY", "LAD", 10000], ["BOS", "HOU", 5000], ["ATL", "PHI", 10000]],
        inputs=[home_dd, away_dd, n_slider],
    )


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