from __future__ import annotations

import json
import os
import urllib.parse
import urllib.request
from dataclasses import dataclass
from datetime import datetime, timedelta
from pathlib import Path
from zoneinfo import ZoneInfo

import gradio as gr
import numpy as np
import pandas as pd
from sklearn.linear_model import Ridge
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler

DATA_PATH = Path(__file__).resolve().parent / "data" / "solar_timeseries_5m.csv"
REGIONS = ["NSW1", "QLD1", "SA1", "TAS1", "VIC1"]
SOLAR_TYPES = ["rooftop", "utility"]
HORIZONS = [5, 15, 30]
HORIZON_STEPS = {5: 1, 15: 3, 30: 6}

BASE_URL = "https://api.openelectricity.org.au/v4/data/network/NEM"
UA = (
    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 "
    "(KHTML, like Gecko) Chrome/125.0.0.0 Safari/537.36"
)
NEM_TZ = ZoneInfo("Australia/Sydney")

_DF_CACHE: pd.DataFrame | None = None
_MODEL_CACHE: dict[tuple[str, str, int], "ModelBundle"] | None = None


@dataclass
class ModelBundle:
    pipeline: Pipeline
    feature_cols: list[str]


def _series_col(region: str, solar_type: str) -> str:
    return f"{region}_gen_solar_{solar_type}_mw"


def _load_data() -> pd.DataFrame:
    global _DF_CACHE
    if _DF_CACHE is None:
        df = pd.read_csv(DATA_PATH)
        df["timestamp"] = pd.to_datetime(df["timestamp"], utc=True)
        _DF_CACHE = df.sort_values("timestamp").reset_index(drop=True)
    return _DF_CACHE.copy()


def _feature_table(series: pd.Series, ts: pd.Series, horizon_steps: int) -> pd.DataFrame:
    out = pd.DataFrame({"y": series})
    out["lag_1"] = series.shift(1)
    out["lag_2"] = series.shift(2)
    out["lag_3"] = series.shift(3)
    out["lag_6"] = series.shift(6)
    out["lag_12"] = series.shift(12)
    out["roll_3"] = series.rolling(3).mean()
    out["roll_12"] = series.rolling(12).mean()

    minute_of_day = ts.dt.hour * 60 + ts.dt.minute
    out["tod_sin"] = np.sin(2 * np.pi * minute_of_day / 1440.0)
    out["tod_cos"] = np.cos(2 * np.pi * minute_of_day / 1440.0)
    dow = ts.dt.day_of_week
    out["dow_sin"] = np.sin(2 * np.pi * dow / 7.0)
    out["dow_cos"] = np.cos(2 * np.pi * dow / 7.0)

    out["target"] = series.shift(-horizon_steps)
    return out


def _fit_model(series: pd.Series, ts: pd.Series, horizon_steps: int) -> ModelBundle:
    table = _feature_table(series, ts, horizon_steps)
    feature_cols = [
        "lag_1",
        "lag_2",
        "lag_3",
        "lag_6",
        "lag_12",
        "roll_3",
        "roll_12",
        "tod_sin",
        "tod_cos",
        "dow_sin",
        "dow_cos",
    ]
    sub = table.dropna(subset=feature_cols + ["target"])
    X = sub[feature_cols].to_numpy(float)
    y = sub["target"].to_numpy(float)

    model = Pipeline(
        [
            ("scaler", StandardScaler()),
            ("ridge", Ridge(alpha=1.5)),
        ]
    )
    model.fit(X, y)
    return ModelBundle(model, feature_cols)


def _train_all(df: pd.DataFrame) -> dict[tuple[str, str, int], ModelBundle]:
    trained: dict[tuple[str, str, int], ModelBundle] = {}
    ts = df["timestamp"]
    for region in REGIONS:
        for solar_type in SOLAR_TYPES:
            col = _series_col(region, solar_type)
            s = df[col].astype(float)
            for horizon in HORIZONS:
                trained[(region, solar_type, horizon)] = _fit_model(s, ts, HORIZON_STEPS[horizon])
    return trained


def _api_key() -> str | None:
    return os.environ.get("OPENELECTRICITY_API_KEY", "").strip() or None


def _http_get_json(url: str, bearer: str) -> dict:
    req = urllib.request.Request(
        url,
        headers={
            "Authorization": f"Bearer {bearer}",
            "Accept": "application/json",
            "User-Agent": UA,
        },
        method="GET",
    )
    with urllib.request.urlopen(req, timeout=180) as resp:
        return json.load(resp)


def _fetch_latest_live_for_solar_type(bearer: str, solar_type: str) -> dict[str, tuple[pd.Timestamp, float]]:
    """Return region -> (latest_timestamp_utc, value_mw) for one solar type."""
    now_local = datetime.now(NEM_TZ).replace(microsecond=0)
    start_local = now_local - timedelta(hours=12)

    params = [
        ("metrics", "power"),
        ("interval", "5m"),
        ("date_start", start_local.strftime("%Y-%m-%dT%H:%M:%S")),
        ("date_end", now_local.strftime("%Y-%m-%dT%H:%M:%S")),
        ("primary_grouping", "network_region"),
        ("fueltech", f"solar_{solar_type}"),
    ]
    url = BASE_URL + "?" + urllib.parse.urlencode(params)
    payload = _http_get_json(url, bearer)

    if not payload.get("success", True) or payload.get("error"):
        raise RuntimeError(str(payload.get("error") or payload))

    out: dict[str, tuple[pd.Timestamp, float]] = {}
    for block in payload.get("data") or []:
        for series in block.get("results") or []:
            region = (series.get("columns") or {}).get("region")
            if region not in REGIONS:
                continue
            latest_pair: tuple[pd.Timestamp, float] | None = None
            for row in series.get("data") or []:
                if not row or len(row) < 2 or row[1] is None:
                    continue
                ts = pd.to_datetime(str(row[0]), utc=True)
                val = float(row[1])
                if latest_pair is None or ts > latest_pair[0]:
                    latest_pair = (ts, val)
            if latest_pair is not None:
                out[region] = latest_pair
    return out


def _current_snapshot(df: pd.DataFrame) -> tuple[pd.Timestamp, dict[str, dict[str, float]], str]:
    """
    Returns (snapshot_ts_utc, region->solar_type->mw, source_label).
    Uses live API when key is present; falls back to local latest row on failure.
    """
    local_last = df.iloc[-1]
    local_ts = pd.to_datetime(local_last["timestamp"], utc=True)
    local_map: dict[str, dict[str, float]] = {}
    for region in REGIONS:
        local_map[region] = {
            "rooftop": float(local_last[_series_col(region, "rooftop")]),
            "utility": float(local_last[_series_col(region, "utility")]),
        }

    bearer = _api_key()
    if not bearer:
        return local_ts, local_map, "local_fallback"

    try:
        live_roof = _fetch_latest_live_for_solar_type(bearer, "rooftop")
        live_util = _fetch_latest_live_for_solar_type(bearer, "utility")
    except Exception:
        return local_ts, local_map, "local_fallback"

    out_map: dict[str, dict[str, float]] = {}
    ts_candidates: list[pd.Timestamp] = []
    any_live = False

    for region in REGIONS:
        roof_val = local_map[region]["rooftop"]
        util_val = local_map[region]["utility"]

        if region in live_roof:
            ts_candidates.append(live_roof[region][0])
            roof_val = live_roof[region][1]
            any_live = True
        if region in live_util:
            ts_candidates.append(live_util[region][0])
            util_val = live_util[region][1]
            any_live = True

        out_map[region] = {
            "rooftop": roof_val,
            "utility": util_val,
        }

    snapshot_ts = max(ts_candidates) if ts_candidates else local_ts
    return snapshot_ts, out_map, ("openelectricity_api" if any_live else "local_fallback")


def _extend_with_snapshot(
    base_ts: pd.Series,
    base_values: pd.Series,
    snapshot_ts: pd.Timestamp,
    snapshot_value: float,
) -> tuple[pd.Series, pd.Series]:
    """
    Inject snapshot point into the series used for feature construction.
    """
    ts = pd.to_datetime(base_ts, utc=True).reset_index(drop=True)
    vals = base_values.astype(float).reset_index(drop=True)
    last_ts = ts.iloc[-1]

    if snapshot_ts > last_ts:
        ts = pd.concat([ts, pd.Series([snapshot_ts])], ignore_index=True)
        vals = pd.concat([vals, pd.Series([snapshot_value])], ignore_index=True)
    elif snapshot_ts == last_ts:
        vals.iloc[-1] = snapshot_value
    else:
        vals.iloc[-1] = snapshot_value
    return ts, vals


def _latest_features(series: pd.Series, ts: pd.Series, feature_cols: list[str]) -> np.ndarray:
    table = _feature_table(series, ts, horizon_steps=1)
    row = table.iloc[-1]
    return row[feature_cols].to_numpy(float).reshape(1, -1)


def _build_dashboard() -> tuple[str, pd.DataFrame, dict]:
    df = _load_data()
    global _MODEL_CACHE
    if _MODEL_CACHE is None:
        _MODEL_CACHE = _train_all(df)
    models = _MODEL_CACHE

    snapshot_ts, current_map, source_kind = _current_snapshot(df)

    rows = []
    payload_regions = []
    pred_generated = pd.Timestamp.utcnow().isoformat().replace("+00:00", "Z")
    snapshot_iso = snapshot_ts.isoformat().replace("+00:00", "Z")

    for region in REGIONS:
        row: dict[str, object] = {
            "Region": region,
            "Prediction Generated At": pred_generated,
            "Source Snapshot At": snapshot_iso,
            "Current Rooftop Solar (MW)": round(current_map[region]["rooftop"], 3),
            "Current Utility Solar (MW)": round(current_map[region]["utility"], 3),
        }

        region_payload = {
            "prediction_generated_at": pred_generated,
            "source_snapshot_at": snapshot_iso,
            "region": region,
            "source": source_kind,
            "current_values": {
                "solar_rooftop_mw": current_map[region]["rooftop"],
                "solar_utility_mw": current_map[region]["utility"],
            },
            "forecasts": {},
        }

        for horizon in HORIZONS:
            horizon_key = f"{horizon}m"
            region_payload["forecasts"][horizon_key] = {}

            for solar_type in SOLAR_TYPES:
                col = _series_col(region, solar_type)
                base_vals = df[col]
                ext_ts, ext_vals = _extend_with_snapshot(
                    df["timestamp"],
                    base_vals,
                    snapshot_ts,
                    current_map[region][solar_type],
                )
                bundle = models[(region, solar_type, horizon)]
                x = _latest_features(ext_vals, ext_ts, bundle.feature_cols)
                yhat = float(bundle.pipeline.predict(x)[0])
                yhat = max(0.0, yhat)

                label_prefix = "Rooftop" if solar_type == "rooftop" else "Utility"
                row[f"{label_prefix} In {horizon}m (MW)"] = round(yhat, 3)
                region_payload["forecasts"][horizon_key][f"solar_{solar_type}_mw"] = yhat

        rows.append(row)
        payload_regions.append(region_payload)

    out_df = pd.DataFrame(rows)
    payload = {
        "prediction_generated_at": pred_generated,
        "source_snapshot_at": snapshot_iso,
        "source": source_kind,
        "prediction_horizons_minutes": HORIZONS,
        "regions": payload_regions,
    }
    summary = (
        f"Prediction generated at: {pred_generated} | "
        f"Source snapshot at: {snapshot_iso} | "
        f"Source: {source_kind} | Horizons: [5, 15, 30]"
    )
    return summary, out_df, payload


def refresh_solar_dashboard() -> tuple[str, pd.DataFrame, dict]:
    return _build_dashboard()


with gr.Blocks(theme=gr.themes.Soft(), title="NEM Solar Predictor") as demo:
    gr.Markdown("## NEM Solar Predictor")
    gr.Markdown("Current regional solar baselines (rooftop + utility) with 5m, 15m, and 30m forecasts.")

    summary = gr.Textbox(label="Summary", interactive=False)
    table = gr.Dataframe(label="Regional Predictions", interactive=False)
    payload = gr.JSON(label="Raw Prediction Payload")
    refresh = gr.Button("Refresh Predictions", variant="primary")

    refresh.click(
        fn=refresh_solar_dashboard,
        inputs=[],
        outputs=[summary, table, payload],
        api_name="/refresh_solar_dashboard",
    )

    demo.load(
        fn=refresh_solar_dashboard,
        inputs=[],
        outputs=[summary, table, payload],
        api_name="/refresh_solar_dashboard_1",
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", "7860")))