"""Build propagation chain data with time-of-day verdict analysis.

Uses CRSDepTime to bucket flights into morning/afternoon/evening and trace
how delays propagate through aircraft rotations. The output powers the
"Should I book this flight?" verdict panel.

Run AFTER populate_db.py (needs crs_dep_time column).
"""
import json
import sys
from collections import defaultdict

import pandas as pd
from sqlalchemy import text

from core.database import engine
from core.utils import get_output_path
from .stats import MIN_ROUTE_FLIGHTS, MIN_CARRIER_FLIGHTS, classify_risk, generate_reason, compute_verdict
from .builder import build_upstream, build_downstream


def ensure_indexes():
    print("Checking indexes...")
    with engine.connect() as conn:
        conn.execute(text(
            "CREATE INDEX IF NOT EXISTS idx_tail_date ON flights(tail_number, fl_date)"
        ))
        conn.commit()


def load_pairs():
    """Get consecutive flight legs paired by tail number, with time windows."""
    print("Loading consecutive leg pairs (big query, hang tight)...")

    q = """
        WITH ordered_legs AS (
            SELECT origin, dest, tail_number, fl_date, carrier,
                   crs_dep_time, dep_delay, arr_delay, arr_del15,
                   late_aircraft_delay, distance,
                   CASE
                       WHEN crs_dep_time >= 600 AND crs_dep_time < 1200 THEN 'morning'
                       WHEN crs_dep_time >= 1200 AND crs_dep_time < 1800 THEN 'afternoon'
                       ELSE 'evening'
                   END AS time_window,
                   ROW_NUMBER() OVER (
                       PARTITION BY tail_number, fl_date
                       ORDER BY crs_dep_time
                   ) AS leg_num
            FROM flights
            WHERE cancelled = 0 AND diverted = 0
              AND tail_number IS NOT NULL AND tail_number != ''
              AND crs_dep_time IS NOT NULL
        )
        SELECT
            prev.origin     AS prev_origin,
            prev.dest       AS prev_dest,
            prev.time_window AS prev_time,
            prev.arr_delay  AS prev_arr_delay,
            curr.origin     AS curr_origin,
            curr.dest       AS curr_dest,
            curr.carrier    AS carrier,
            curr.time_window AS curr_time,
            curr.dep_delay  AS curr_dep_delay,
            curr.arr_delay  AS curr_arr_delay,
            curr.arr_del15  AS curr_delayed,
            curr.late_aircraft_delay AS late_aircraft,
            curr.distance   AS curr_distance,
            curr.leg_num
        FROM ordered_legs curr
        JOIN ordered_legs prev
          ON curr.tail_number = prev.tail_number
         AND curr.fl_date = prev.fl_date
         AND curr.leg_num = prev.leg_num + 1
    """
    df = pd.read_sql(q, engine)
    print(f"  {len(df):,} consecutive pairs loaded")
    return df


def load_route_stats():
    """Standalone route stats (not chain-dependent)."""
    print("Loading route-level stats...")
    q = """
        SELECT origin, dest,
               COUNT(*) AS flights,
               AVG(dep_delay) AS avg_delay,
               AVG(CASE WHEN arr_del15 = 1 THEN 1.0 ELSE 0.0 END) AS pct_delayed
        FROM flights
        WHERE cancelled = 0
        GROUP BY origin, dest
        HAVING COUNT(*) >= ?
    """
    df = pd.read_sql(q, engine, params=(MIN_ROUTE_FLIGHTS,))
    df["key"] = df["origin"] + "-" + df["dest"]
    return df.set_index("key")


def load_time_stats():
    """Per-route, per-time-window stats from the full flights table."""
    print("Loading time-window stats...")
    q = """
        SELECT origin, dest,
               CASE
                   WHEN crs_dep_time >= 600 AND crs_dep_time < 1200 THEN 'morning'
                   WHEN crs_dep_time >= 1200 AND crs_dep_time < 1800 THEN 'afternoon'
                   ELSE 'evening'
               END AS tw,
               COUNT(*) AS flights,
               AVG(dep_delay) AS avg_delay,
               AVG(CASE WHEN arr_del15 = 1 THEN 1.0 ELSE 0.0 END) AS pct_delayed
        FROM flights
        WHERE cancelled = 0 AND crs_dep_time IS NOT NULL
        GROUP BY origin, dest, tw
        HAVING COUNT(*) >= 50
    """
    df = pd.read_sql(q, engine)
    df["key"] = df["origin"] + "-" + df["dest"]
    return df


def load_carrier_stats():
    """By-carrier breakdown (overall, not by time window)."""
    print("Loading carrier stats...")
    q = """
        SELECT origin, dest, carrier,
               COUNT(*) AS flights,
               AVG(CASE WHEN arr_del15 = 1 THEN 1.0 ELSE 0.0 END) AS pct_delayed
        FROM flights
        WHERE cancelled = 0 AND carrier IS NOT NULL
        GROUP BY origin, dest, carrier
        HAVING COUNT(*) >= ?
    """
    df = pd.read_sql(q, engine, params=(MIN_CARRIER_FLIGHTS,))
    df["key"] = df["origin"] + "-" + df["dest"]
    return df


def _index_time_stats(time_stats_df):
    result = defaultdict(dict)
    for _, row in time_stats_df.iterrows():
        result[row["key"]][row["tw"]] = {
            "flights": int(row["flights"]),
            "avg_delay": round(float(row["avg_delay"]), 1) if pd.notna(row["avg_delay"]) else None,
            "pct_delayed": round(float(row["pct_delayed"]), 3) if pd.notna(row["pct_delayed"]) else None,
        }
    return result


def _index_carrier_stats(carrier_df):
    result = defaultdict(dict)
    for _, row in carrier_df.iterrows():
        result[row["key"]][row["carrier"]] = {
            "flights": int(row["flights"]),
            "pct_delayed": round(float(row["pct_delayed"]), 3) if pd.notna(row["pct_delayed"]) else None,
        }
    return result


def _build_time_entry(tw, tw_stats, upstream, downstream, route_key):
    pct = tw_stats["pct_delayed"] or 0
    risk = classify_risk(pct)
    up_links = upstream.get(route_key, {}).get(tw, [])
    dn_links = downstream.get(route_key, {}).get(tw, [])

    tw_entry = {
        "flights": tw_stats["flights"],
        "pct_delayed": round(pct * 100, 1),  # percentage for display
        "avg_delay": tw_stats["avg_delay"],
        "risk_level": risk,
        "reason": generate_reason(tw, risk, up_links),
    }
    if up_links:
        tw_entry["upstream"] = up_links
    if dn_links:
        tw_entry["downstream"] = dn_links
    return tw_entry


def _add_verdict(entry, by_time, carriers):
    verdict = compute_verdict(by_time)
    if not verdict:
        return
    if carriers:
        sorted_c = sorted(carriers.items(), key=lambda x: x[1]["pct_delayed"] or 0.5)
        verdict["best_carrier"] = sorted_c[0][0]
        verdict["worst_carrier"] = sorted_c[-1][0]
        verdict["recommendation"] = f"Book {sorted_c[0][0]} {verdict['best_time']} flights"
    entry["verdict"] = verdict


def _safe_round(val, decimals=1):
    return round(float(val), decimals) if pd.notna(val) else None


def assemble(route_stats, time_stats_df, carrier_df, upstream, downstream):
    """Put it all together."""
    print("Assembling final JSON...")

    time_by_route = _index_time_stats(time_stats_df)
    carrier_by_route = _index_carrier_stats(carrier_df)
    output = {}

    for route_key in sorted(route_stats.index):
        rs = route_stats.loc[route_key]

        entry = {
            "overall": {
                "flights": int(rs["flights"]),
                "avg_delay": _safe_round(rs["avg_delay"]),
                "pct_delayed": _safe_round(rs["pct_delayed"], 3),
            },
        }

        # by_time with upstream + downstream embedded
        by_time = {}
        for tw in ["morning", "afternoon", "evening"]:
            tw_stats = time_by_route.get(route_key, {}).get(tw)
            if tw_stats:
                by_time[tw] = _build_time_entry(tw, tw_stats, upstream, downstream, route_key)

        if by_time:
            entry["by_time"] = by_time

        _add_verdict(entry, by_time, carrier_by_route.get(route_key))

        if route_key in carrier_by_route:
            entry["by_carrier"] = carrier_by_route[route_key]

        output[route_key] = entry

    return output


def main():
    with engine.connect() as conn:
        count = conn.execute(text("SELECT COUNT(*) FROM flights")).scalar()
        has_dep = conn.execute(text(
            "SELECT COUNT(*) FROM flights WHERE crs_dep_time IS NOT NULL"
        )).scalar()

    if count == 0:
        print("flights table empty -- run populate_db.py first")
        sys.exit(1)

    if has_dep == 0:
        print("No crs_dep_time data! Re-run populate_db.py with CRSDepTime.")
        sys.exit(1)

    pct_has = has_dep / count * 100
    print(f"Working with {count:,} flights ({pct_has:.0f}% have departure times)\n")

    ensure_indexes()

    pairs = load_pairs()
    route_stats = load_route_stats()
    time_stats = load_time_stats()
    carrier_stats = load_carrier_stats()

    up = build_upstream(pairs)
    down = build_downstream(pairs)

    data = assemble(route_stats, time_stats, carrier_stats, up, down)

    output_path = get_output_path("propagation_chains.json")
    output_path.parent.mkdir(parents=True, exist_ok=True)
    with open(output_path, "w") as f:
        json.dump(data, f, separators=(",", ":"))

    size_mb = output_path.stat().st_size / (1024 * 1024)
    print(f"\nWrote {output_path}")
    print(f"  {len(data)} routes")
    print(f"  {size_mb:.1f} MB")

    # quick sanity check
    has_verdict = sum(1 for r in data.values() if "verdict" in r)
    has_upstream = sum(1 for r in data.values()
                       if "by_time" in r and any("upstream" in r["by_time"].get(tw, {})
                                                  for tw in ["morning", "afternoon", "evening"]))
    print(f"  {has_verdict} routes have verdicts")
    print(f"  {has_upstream} routes have upstream chains")

    # sample a couple routes
    print("\n--- Sample routes ---")
    for key in list(data.keys())[:3]:
        r = data[key]
        v = r.get("verdict", {})
        print(f"\n  {key}: {r['overall']['flights']} flights")
        if v:
            print(f"    Best: {v.get('best_time')} / Worst: {v.get('worst_time')}")
            if "recommendation" in v:
                print(f"    → {v['recommendation']}")
        if "by_time" in r:
            for tw in ["morning", "afternoon", "evening"]:
                tw_data = r["by_time"].get(tw)
                if tw_data:
                    n_up = len(tw_data.get("upstream", []))
                    n_dn = len(tw_data.get("downstream", []))
                    print(f"    {tw}: {tw_data['pct_delayed']}% delayed, "
                          f"{tw_data['risk_level']}, {n_up} upstream, {n_dn} downstream")


if __name__ == "__main__":
    main()
