"""Exporterar energidata fran Home Assistant till Woltas uppladdnings-CSV.

Hamtar timvisa matarvarden (sol, batteri, elnat) via HA:s WebSocket-API och
skriver en CSV som kan laddas upp direkt till Wolta (Optimeringsbetyg - folj guiden).

Kraver ett Long-Lived Access Token fran HA (Profil - Sakerhet - Tokens).
Spara token i en fil och peka pa den med --token-file.

Anvandning:
    uv run python tools/export_energy_csv.py \\
        --ha-url wss://hass.local.example.com/api/websocket \\
        --token-file ~/.ha_token \\
        --months 12 \\
        --out wolta-driftdata.csv

Entitet-ID:n ar forifyllda med Sonnen-standardnamn. Justera vid behov:
    --solar sensor.solar_production
    --grid-import sensor.grid_imported_energy
    --grid-export sensor.grid_exported_energy
    --batt-charged sensor.energy_going_into_battery_2
    --batt-discharged sensor.energy_going_out_from_battery

Utelamna en kolumn (om du t.ex. saknar solceller) genom att skicka tomt varde:
    --solar ""

Hogupplost 15-min-export (for optimeringsbetyget - fangar intra-tim-arbitrage):
    --resolution 15min
Detta hamtar ra recorder-history i stallet for timvis statistik. History sparas bara
sa langt 'recorder: purge_keep_days' tillater (ofta ~10-45 dygn) - aldre data saknas da.
Optimeringsbetyget kraver minst 30 dygn, sa hoj purge_keep_days vid behov.

CSV-format (Woltas kanoniska format):
    timestamp,solar_kwh,grid_import_kwh,grid_export_kwh,batt_charged_kwh,batt_discharged_kwh
    En rad per timme, ISO-8601-tidstampel med offset (+00:00).
    Endast kolumner med faktisk data inkluderas.
"""

import argparse
import asyncio
import csv
import json
import os
import sys
from datetime import datetime, timedelta, timezone


try:
    import websockets
except ImportError:
    sys.stderr.write("Paketet 'websockets' saknas. Kor: pip install websockets\n")
    sys.exit(1)

# Tvinga UTF-8 pa stdout/stderr pa Windows (annars kraschar svenska tecken i cp1252)
if hasattr(sys.stdout, "reconfigure"):
    sys.stdout.reconfigure(encoding="utf-8", errors="replace")
if hasattr(sys.stderr, "reconfigure"):
    sys.stderr.reconfigure(encoding="utf-8", errors="replace")


# Månadsvis chunkning för att undvika WebSocket-timeout vid långa perioder
CHUNK_DAYS = 30
# Mindre chunkar för högupplöst history (rå states ger mycket större payload)
CHUNK_DAYS_HISTORY = 7
QUARTER = timedelta(minutes=15)

# Sensor → CSV-kolumn
COLUMN_MAP = {
    "solar": "solar_kwh",
    "grid_import": "grid_import_kwh",
    "grid_export": "grid_export_kwh",
    "batt_charged": "batt_charged_kwh",
    "batt_discharged": "batt_discharged_kwh",
}

# Standardentiteter (Sonnen-installation, vanliga namn)
DEFAULTS = {
    "solar": "sensor.solar_production",
    "grid_import": "sensor.grid_imported_energy",
    "grid_export": "sensor.grid_exported_energy",
    "batt_charged": "sensor.energy_going_into_battery_2",
    "batt_discharged": "sensor.energy_going_out_from_battery",
}


async def ha_stats(
    url: str,
    token: str,
    ids: list[str],
    start: datetime,
    end: datetime,
    period: str,
) -> dict:
    """Hämtar recorder/statistics_during_period för en tidsperiod."""
    async with websockets.connect(url, max_size=2**24) as ws:
        await ws.recv()  # auth_required
        await ws.send(json.dumps({"type": "auth", "access_token": token}))
        assert json.loads(await ws.recv())["type"] == "auth_ok"
        await ws.send(
            json.dumps(
                {
                    "id": 1,
                    "type": "recorder/statistics_during_period",
                    "start_time": start.isoformat(),
                    "end_time": end.isoformat(),
                    "statistic_ids": ids,
                    "period": period,
                    "types": ["sum"],
                }
            )
        )
        resp = json.loads(await ws.recv())
        assert resp["success"], resp
        return resp["result"]


def hourly_diffs(series: list[dict]) -> dict[datetime, float]:
    """Kumulativ sum-serie → kWh per timme (hanterar nollställningar/omstarter)."""
    out, prev = {}, None
    for row in series:
        ts = datetime.fromtimestamp(row["start"] / 1000, tz=timezone.utc)
        s = row.get("sum")
        if prev is not None and s is not None and s >= prev:
            out[ts] = s - prev
        prev = s if s is not None else prev
    return out


async def ha_history(
    url: str,
    token: str,
    ids: list[str],
    start: datetime,
    end: datetime,
) -> dict:
    """Hämtar rå states via history/history_during_period (full upplösning, ~10-45 dygns
    retention beroende på recorder purge_keep_days). Returnerar {entity_id: [state-rad]}."""
    async with websockets.connect(url, max_size=2**26) as ws:
        await ws.recv()  # auth_required
        await ws.send(json.dumps({"type": "auth", "access_token": token}))
        assert json.loads(await ws.recv())["type"] == "auth_ok"
        await ws.send(
            json.dumps(
                {
                    "id": 1,
                    "type": "history/history_during_period",
                    "start_time": start.isoformat(),
                    "end_time": end.isoformat(),
                    "entity_ids": ids,
                    "minimal_response": True,
                    "no_attributes": True,
                }
            )
        )
        resp = json.loads(await ws.recv())
        assert resp["success"], resp
        return resp["result"]


def _state_points(rows: list[dict]) -> list[tuple[datetime, float]]:
    """Parsar history-rader → [(ts, kumulativt värde)]. Hanterar både fulla rader
    ({state,last_updated}) och minimal_response-kompaktform ({s, lu/lc}). Hoppar över
    unavailable/unknown."""
    pts: list[tuple[datetime, float]] = []
    for r in rows:
        raw = r.get("s", r.get("state"))
        ts_epoch = r.get("lu", r.get("lc"))
        if ts_epoch is None and "last_updated" in r:
            try:
                ts_epoch = datetime.fromisoformat(r["last_updated"]).timestamp()
            except ValueError:
                ts_epoch = None
        if raw is None or ts_epoch is None:
            continue
        try:
            val = float(raw)
        except (TypeError, ValueError):
            continue  # unavailable/unknown
        pts.append((datetime.fromtimestamp(float(ts_epoch), tz=timezone.utc), val))
    pts.sort(key=lambda p: p[0])
    return pts


def quarter_hour_diffs(rows: list[dict]) -> dict[datetime, float]:
    """Rå kumulativa states → kWh per 15-min-intervall. Resamplar den kumulativa serien
    till 15-min-rutnätet (senaste värde ≤ rutgräns) och diffar; nollställningar (värdet
    minskar, t.ex. mätarbyte) ger 0 för det intervallet. Nyckeln är intervallets START."""
    pts = _state_points(rows)
    if len(pts) < 2:
        return {}
    start = pts[0][0].replace(minute=(pts[0][0].minute // 15) * 15, second=0, microsecond=0)
    end = pts[-1][0]
    # Kumulativt värde vid varje 15-min-gräns = senaste state med ts ≤ gränsen.
    boundary_val: dict[datetime, float] = {}
    i = 0
    b = start
    last = pts[0][1]
    while b <= end:
        while i < len(pts) and pts[i][0] <= b:
            last = pts[i][1]
            i += 1
        boundary_val[b] = last
        b += QUARTER
    bounds = sorted(boundary_val)
    out: dict[datetime, float] = {}
    for j in range(len(bounds) - 1):
        delta = boundary_val[bounds[j + 1]] - boundary_val[bounds[j]]
        out[bounds[j]] = delta if delta >= 0 else 0.0
    return out


def fetch_all_chunks(
    url: str,
    token: str,
    ids: list[str],
    start: datetime,
    end: datetime,
    resolution: str = "hour",
) -> dict[str, list[dict]]:
    """Hämtar data i block och slår ihop till ett sammanhängande dataset.
    resolution='hour' → statistik (timvis); '15min' → rå history (full upplösning)."""
    combined: dict[str, list[dict]] = {}
    highres = resolution == "15min"
    chunk_days = CHUNK_DAYS_HISTORY if highres else CHUNK_DAYS

    chunks: list[tuple[datetime, datetime]] = []
    cur = start
    while cur < end:
        chunk_end = min(cur + timedelta(days=chunk_days), end)
        chunks.append((cur, chunk_end))
        cur = chunk_end

    label = "history-block (15-min)" if highres else "månadsblock"
    print(f"Hämtar {len(chunks)} {label} från HA...", flush=True)

    for i, (cs, ce) in enumerate(chunks, start=1):
        print(f"  Block {i}/{len(chunks)}: {cs.date()} → {ce.date()}", flush=True)
        if highres:
            result = asyncio.run(ha_history(url, token, ids, cs, ce))
        else:
            result = asyncio.run(ha_stats(url, token, ids, cs, ce, "hour"))
        for sid, rows in result.items():
            combined.setdefault(sid, []).extend(rows)

    # Dedup + sortera vid blockgränser (samma tidpunkt kan finnas i två block).
    def _key(row: dict):
        return row.get("start", row.get("lu", row.get("lc")))

    for sid in combined:
        seen: set = set()
        deduped: list[dict] = []
        for row in sorted(combined[sid], key=lambda r: _key(r) or 0):
            k = _key(row)
            if k not in seen:
                seen.add(k)
                deduped.append(row)
        combined[sid] = deduped

    return combined


def main() -> None:
    ap = argparse.ArgumentParser(
        description="Exportera energidata från Home Assistant till Woltas CSV-format.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )
    ap.add_argument(
        "--ha-url",
        required=True,
        metavar="WSS://…",
        help="WebSocket-URL till HA, t.ex. wss://hass.local.example.com/api/websocket",
    )
    ap.add_argument(
        "--token-file",
        required=True,
        metavar="FIL",
        help="Sökväg till fil med Long-Lived Access Token",
    )
    ap.add_argument(
        "--out",
        default="wolta-driftdata.csv",
        metavar="CSV",
        help="Utdatafil (standard: wolta-driftdata.csv)",
    )
    ap.add_argument(
        "--months",
        type=int,
        default=12,
        metavar="N",
        help="Antal månader bakåt att hämta (standard: 12)",
    )
    ap.add_argument(
        "--resolution",
        choices=["hour", "15min"],
        default="hour",
        help="Upplösning: 'hour' (timvis statistik, år bakåt) eller '15min' (rå history, "
             "full upplösning men bara så långt recorder-historiken sparas, ofta ~10-45 dygn). "
             "15min fångar intra-tim-arbitrage i optimeringsbetyget.",
    )

    # Entitets-ID-argument – tomt värde = sensorn finns ej
    ap.add_argument(
        "--solar",
        default=DEFAULTS["solar"],
        metavar="ENTITY_ID",
        help=f"Solproduktion (standard: {DEFAULTS['solar']}). Tom sträng = saknas.",
    )
    ap.add_argument(
        "--grid-import",
        default=DEFAULTS["grid_import"],
        metavar="ENTITY_ID",
        help=f"Nätimport (standard: {DEFAULTS['grid_import']})",
    )
    ap.add_argument(
        "--grid-export",
        default=DEFAULTS["grid_export"],
        metavar="ENTITY_ID",
        help=f"Nätexport (standard: {DEFAULTS['grid_export']})",
    )
    ap.add_argument(
        "--batt-charged",
        default=DEFAULTS["batt_charged"],
        metavar="ENTITY_ID",
        help=f"Batteriladdning (standard: {DEFAULTS['batt_charged']})",
    )
    ap.add_argument(
        "--batt-discharged",
        default=DEFAULTS["batt_discharged"],
        metavar="ENTITY_ID",
        help=f"Batteriurladdning (standard: {DEFAULTS['batt_discharged']})",
    )

    args = ap.parse_args()

    # Läs token
    token_path = os.path.expanduser(args.token_file)
    if not os.path.isfile(token_path):
        print(f"Fel: token-filen hittades inte: {token_path}", file=sys.stderr)
        sys.exit(1)
    token = open(token_path).read().strip()

    # Bygg sensor→nyckel-tabell (hoppa över tomma ID:n)
    sensor_keys: dict[str, str] = {}  # entity_id → logisk nyckel
    for key in COLUMN_MAP:
        entity_id = getattr(args, key.replace("-", "_"), "").strip()
        if entity_id:
            sensor_keys[entity_id] = key

    if not sensor_keys:
        print("Fel: minst en sensor måste vara konfigurerad.", file=sys.stderr)
        sys.exit(1)

    # Tidsintervall
    end = datetime.now(timezone.utc).replace(minute=0, second=0, microsecond=0)
    start = end - timedelta(days=30 * args.months)

    res_label = "15-min (rå history)" if args.resolution == "15min" else "timvis"
    print(f"\nHämtar {args.months} månaders data ({start.date()} → {end.date()}), upplösning: {res_label}")
    if args.resolution == "15min":
        print("OBS: 15-min använder recorder-historiken, som ofta bara sparas ~10-45 dygn "
              "(beror på purge_keep_days). Saknas äldre data exporteras bara det som finns.")
    print(f"Sensorer: {', '.join(sensor_keys)}\n")

    # Hämta data
    raw = fetch_all_chunks(args.ha_url, token, list(sensor_keys), start, end, args.resolution)

    # Auto-detect: vilka sensorer returnerade faktisk data?
    active: dict[str, str] = {}  # logisk nyckel → entity_id
    empty: list[str] = []
    for entity_id, key in sensor_keys.items():
        rows = raw.get(entity_id, [])
        if rows:
            active[key] = entity_id
        else:
            empty.append(entity_id)

    if empty:
        print(f"\nOBS: Ingen data hittades för: {', '.join(empty)} (kolumner utelämnas)")
    if not active:
        print("Fel: ingen sensor returnerade data. Kontrollera entity-ID:n.", file=sys.stderr)
        sys.exit(1)

    print(f"\nAktiva kolumner: {', '.join(COLUMN_MAP[k] for k in active)}")

    # Bygg per-intervall-tabell per aktiv sensor (timvis eller 15-min)
    diff_fn = quarter_hour_diffs if args.resolution == "15min" else hourly_diffs
    hourly_data: dict[str, dict[datetime, float]] = {}
    for key, entity_id in active.items():
        hourly_data[key] = diff_fn(raw[entity_id])

    # Union av alla tidstämplar
    all_timestamps: set[datetime] = set()
    for series in hourly_data.values():
        all_timestamps.update(series.keys())

    if not all_timestamps:
        print("Fel: inga timvärden kunde beräknas.", file=sys.stderr)
        sys.exit(1)

    sorted_ts = sorted(all_timestamps)

    # Skriv CSV
    csv_columns = ["timestamp"] + [COLUMN_MAP[k] for k in active]

    out_path = args.out
    with open(out_path, "w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=csv_columns)
        writer.writeheader()
        for ts in sorted_ts:
            row: dict[str, object] = {
                "timestamp": ts.strftime("%Y-%m-%dT%H:%M:%S+00:00")
            }
            for key in active:
                col = COLUMN_MAP[key]
                val = hourly_data[key].get(ts, 0.0)
                row[col] = f"{val:.4f}"
            writer.writerow(row)

    # Sammanfattning
    date_min = sorted_ts[0].date()
    date_max = sorted_ts[-1].date()
    print(f"\nKlar!")
    print(f"  Fil:      {out_path}")
    print(f"  Rader:    {len(sorted_ts)}")
    print(f"  Period:   {date_min} → {date_max}")
    print(f"  Kolumner: {', '.join(csv_columns[1:])}")
    print(
        "\nGranska filen innan uppladdning till Wolta "
        "(Optimeringsbetyg → följ guiden, steget Driftdata)."
    )


if __name__ == "__main__":
    main()