#!/usr/bin/env python3
"""
fetch_3a.py - fetch stations, sensors and weather data from the 3A API
(https://api.3acloud.it, spec in 3a_api_dump/spec__serve_docs_base.json).

    python fetch_3a.py overview                      # stations count, sensors, columns, date ranges
    python fetch_3a.py stations                      # list of stations (+ count)
    python fetch_3a.py sensors                       # hourly/daily sensors per station
    python fetch_3a.py catalogue                     # ask again (3 requests), rewrite 3a_stations_sensors.json
    python fetch_3a.py columns                       # all hourly/daily column names
    python fetch_3a.py daily  --from 2024-01-01 --to 2024-12-31
    python fetch_3a.py hourly --from 2024-01-01 --to 2024-01-31 --stations 4,46 \
                              --columns CodiceStazione,Stazione,Timestamp,TemperaturaAriaMedia,Pioggia

    options: --stations 1,2,3   (default: all stations)
             --columns a,b,c    (default: all available columns)
             --chunk-days N     (split long periods into N-day requests)
             --refresh          (ignore cached metadata and ask the server again)

How auth works (from the spec):
    POST /login with headers name, password, totp  -> token
    every other call sends header  token: <token>

Credentials are not stored in this file:
    THREEA_USER         username (default: mohsen)
    THREEA_PASSWORD     password (else asked with a hidden prompt)
    THREEA_TOTP_SECRET  optional base32 TOTP secret -> codes generated automatically
                        (else you are asked to paste the 6-digit code when a login is needed)

The token is cached in 3a_data/.token, so you only need a new TOTP when it expires.
The server allows only 5 requests per 15 minutes; the script waits automatically.
Results are written to 3a_data/ as .json and .csv.
"""

import argparse
import getpass
import json
import os
import sys
import time
from datetime import datetime, timedelta
from pathlib import Path

import pandas as pd
import requests

BASE_URL = "https://api.3acloud.it"
OUT = Path(__file__).resolve().parent / "3a_data"
OUT.mkdir(exist_ok=True)
TOKEN_FILE = OUT / ".token"
TIMEOUT = 120

sys.stdout.reconfigure(encoding="utf-8", errors="replace")
session = requests.Session()
session.headers["output-fmt"] = "json"
_password = None


# ─────────────────────────── rate limit ───────────────────────────

_rate = {"remaining": None, "reset": None}


def _note_rate_limit(resp):
    h = resp.headers
    if "X-RateLimit-Remaining" in h:
        _rate["remaining"] = int(h["X-RateLimit-Remaining"])
    if "X-RateLimit-Reset" in h:
        _rate["reset"] = int(h["X-RateLimit-Reset"])


def _wait_until(epoch, why):
    while True:
        left = int(epoch - time.time()) + 2
        if left <= 0:
            break
        print(f"\r  {why}: waiting {left // 60:02d}:{left % 60:02d} ", end="", flush=True)
        time.sleep(min(left, 5))
    print("\r" + " " * 60 + "\r", end="")


def _respect_rate_limit():
    if _rate["remaining"] == 0 and _rate["reset"] and _rate["reset"] > time.time():
        _wait_until(_rate["reset"], "rate limit (5 requests / 15 min)")


def _send(method, path, **kw):
    """One HTTP request, waiting for the rate-limit window if needed."""
    while True:
        _respect_rate_limit()
        r = session.request(method, BASE_URL + path, timeout=TIMEOUT, **kw)
        _note_rate_limit(r)
        if r.status_code == 429:
            reset = _rate["reset"] or time.time() + 15 * 60
            _wait_until(reset, "server said 429 Too Many Requests")
            continue
        return r


# ─────────────────────────── auth ───────────────────────────

CRED_FILE = Path(__file__).resolve().parent / "3a_credentials.json"


def _cred(key, env):
    """Value from the environment, else from 3a_credentials.json (if filled in)."""
    if os.environ.get(env):
        return os.environ[env]
    if CRED_FILE.exists():
        value = json.loads(CRED_FILE.read_text(encoding="utf-8")).get(key, "")
        if value and not value.startswith("PUT_"):
            return value
    return None


def _ask_totp():
    secret = _cred("totp_secret", "THREEA_TOTP_SECRET")
    if secret:
        import pyotp  # pip install pyotp
        return pyotp.TOTP(secret.replace(" ", "")).now()
    while True:
        code = input("Paste a fresh 6-digit TOTP code (it changes every 30 s): ").strip().replace(" ", "")
        if code.isdigit() and len(code) == 6:
            return code
        print("  That doesn't look like a 6-digit code, try again.")


def _extract_token(resp):
    try:
        data = resp.json()
    except ValueError:
        return resp.text.strip().strip('"') or None
    if isinstance(data, str):
        return data
    if isinstance(data, dict):
        for k in ("token", "access_token", "Token"):
            if k in data:
                return data[k]
    return None


def login():
    global _password
    user = _cred("user", "THREEA_USER") or "mohsen"
    if _password is None:
        _password = _cred("password", "THREEA_PASSWORD") or getpass.getpass(
            f"Password for {user} (typing is hidden, press Enter when done): ")
    for attempt in range(3):
        totp = _ask_totp()
        r = _send("POST", "/login", headers={"name": user, "password": _password, "totp": totp})
        token = _extract_token(r) if r.ok else None
        if token:
            session.headers["token"] = token
            TOKEN_FILE.write_text(token)
            print("  ✓ logged in, token cached")
            return
        print(f"  ✗ login failed ({r.status_code}): {r.text[:200]}")
        if r.status_code not in (401, 403):
            break
        print("  The TOTP may have expired or been mistyped.")
    sys.exit("Could not log in.")


def _load_cached_token():
    if TOKEN_FILE.exists():
        session.headers["token"] = TOKEN_FILE.read_text().strip()
        return True
    return False


def call(method, path, body=None):
    """Authenticated API call; logs in (asking for a TOTP) only when needed."""
    if "token" not in session.headers and not _load_cached_token():
        login()
    r = _send(method, path, json=body)
    if r.status_code in (401, 403):
        print(f"  token rejected ({r.status_code}: {r.text[:80]}) -> new login needed")
        TOKEN_FILE.unlink(missing_ok=True)
        session.headers.pop("token", None)
        login()
        r = _send(method, path, json=body)
    if not r.ok:
        sys.exit(f"{method} {path} failed ({r.status_code}): {r.text[:500]}")
    try:
        return r.json()
    except ValueError:
        return r.text


# ─────────────────────────── saving ───────────────────────────

def save(name, data):
    (OUT / f"{name}.json").write_text(json.dumps(data, indent=2, ensure_ascii=False), encoding="utf-8")
    df = to_frame(data)
    if df is not None:
        df.to_csv(OUT / f"{name}.csv", index=False, encoding="utf-8-sig")
    return df


def to_frame(data):
    if isinstance(data, list) and data and all(isinstance(x, dict) for x in data):
        return pd.json_normalize(data)
    if isinstance(data, dict):
        # e.g. {"data": [...]} or {station_id: [...]}
        lists = {k: v for k, v in data.items() if isinstance(v, list)}
        if len(lists) == 1:
            return to_frame(next(iter(lists.values())))
        if lists and all(not isinstance(x, dict) for v in lists.values() for x in v):
            return pd.DataFrame([(k, x) for k, v in lists.items() for x in v], columns=["key", "value"])
    return None


def cached(name, method, path, body=None, refresh=False):
    """Metadata changes rarely and requests are scarce, so reuse what's on disk."""
    f = OUT / f"{name}.json"
    if f.exists() and not refresh:
        return json.loads(f.read_text(encoding="utf-8"))
    print(f"→ {method} {path}")
    data = call(method, path, body)
    save(name, data)
    return data


# ─────────────────────────── commands ───────────────────────────

def get_stations(refresh=False):
    data = cached("stations", "GET", "/getAvailableStations", refresh=refresh)
    df = to_frame(data)
    return data, df


def station_ids(args):
    if args.stations:
        return [int(s) for s in args.stations.split(",")]
    _, df = get_stations(args.refresh)
    col = next((c for c in ("StationId", "id") if df is not None and c in df.columns), None)
    if col is None:
        sys.exit("Could not read station ids from stations.json; pass --stations explicitly.")
    return sorted(int(i) for i in df[col])


def cmd_stations(args):
    data, df = get_stations(args.refresh)
    if df is None:
        print(json.dumps(data, indent=2, ensure_ascii=False)[:3000])
        return
    print(f"\nNumber of stations: {len(df)}")
    with pd.option_context("display.max_rows", 500, "display.width", 200):
        print(df.to_string(index=False))
    print(f"\nSaved → {OUT / 'stations.csv'}")


def cmd_sensors(args):
    ids = station_ids(args)
    for kind in ("Hourly", "Daily"):
        data = cached(f"sensors_{kind.lower()}", "POST", f"/getAvailableStation{kind}Sensors",
                      {"stationsIds": ids}, args.refresh)
        print(f"\n{kind} sensors:")
        print(json.dumps(data, indent=1, ensure_ascii=False)[:4000])
        print(f"  saved → {OUT / f'sensors_{kind.lower()}.json'}")


def cmd_summary(args):
    """Station x sensor tables built from stations.json + sensors_*.json."""
    _, stations = get_stations(args.refresh)
    names = dict(zip(stations["StationId"].astype(int), stations["StationName"]))
    rows = []
    for kind in ("Hourly", "Daily"):
        data = cached(f"sensors_{kind.lower()}", "POST", f"/getAvailableStation{kind}Sensors",
                      {"stationsIds": station_ids(args)}, args.refresh)
        items = data if isinstance(data, list) else [data]
        for item in items:
            for sid, sensors in item.items():
                for s in sensors:
                    rows.append({"StationId": int(sid), "StationName": names.get(int(sid), ""),
                                 "Resolution": kind.lower(), "Sensor": s.get("name"),
                                 "Description": s.get("description"), "Unit": s.get("unit")})
    long = pd.DataFrame(rows).sort_values(["StationId", "Resolution", "Sensor"])
    long.to_csv(OUT / "station_sensors_list.csv", index=False, encoding="utf-8-sig")

    hourly = long[long.Resolution == "hourly"]
    matrix = pd.crosstab([hourly.StationId, hourly.StationName], hourly.Sensor).clip(upper=1)
    matrix = matrix.reindex(pd.MultiIndex.from_arrays(
        [stations["StationId"].astype(int), stations["StationName"]]), fill_value=0)
    matrix.insert(0, "NumSensors", matrix.sum(axis=1))
    matrix.to_csv(OUT / "station_sensors_matrix_hourly.csv", encoding="utf-8-sig")

    with pd.option_context("display.max_rows", 500, "display.width", 250):
        print(f"\n{len(stations)} stations. Hourly sensors per station (1 = present):\n")
        print(matrix.to_string())
        print("\nSensor catalogue (how many stations have each, hourly / daily):\n")
        cat = long.groupby(["Sensor", "Description", "Unit", "Resolution"]).StationId.nunique()
        print(cat.unstack("Resolution").fillna(0).astype(int).to_string())
    print(f"\nSaved → {OUT / 'station_sensors_list.csv'}\n        {OUT / 'station_sensors_matrix_hourly.csv'}")


# ─────────────────────────── catalogue (3a_stations_sensors.json) ───────────────────────────

CATALOGUE_FILE = Path(__file__).resolve().parent / "3a_stations_sensors.json"

# sensor group (name used by the engine's stations.py) -> label, unit, its hourly api_parameters
SENSOR_GROUPS = [
    ("TemperaturaAria", "Air temperature sensor", "°C",
     ["TemperaturaAriaMax", "TemperaturaAriaMedia", "TemperaturaAriaMin"]),
    ("UmiditaRelativa", "Relative humidity sensor", "%",
     ["UmiditaRelativaMax", "UmiditaRelativaMedia", "UmiditaRelativaMin"]),
    ("TBulb", "Wet-bulb temperature (mean, supplied by 3A)", "°C", ["TBulbMed"]),
    ("Pioggia", "Rain gauge (RainRateMax in mm/h)", "mm", ["Pioggia", "RainRateMax"]),
    ("BagnaturaFogliare", "Leaf wetness sensor", "min", ["BagnaturaFogliare"]),
    ("VelocitaVento2m", "Wind speed sensor (anemometer, 2 m)", "m/s",
     ["VelocitaVento2mMax", "VelocitaVento2mMedia"]),
    ("Pressione", "Barometric pressure sensor", "mbar", ["PressioneMax", "PressioneMed", "PressioneMin"]),
    ("RadiazioneSolare", "Solar radiation sensor (pyranometer, irradiance)", "W/m²",
     ["RadiazioneSolareMax", "RadiazioneSolareMin"]),
    ("Suolo1", "Soil moisture sensor 1", "%", ["Suolo1"]),
    ("Suolo2", "Soil moisture sensor 2", "%", ["Suolo2"]),
    ("TemperaturaTerreno", "Soil temperature sensor 1", "°C", ["TemperaturaTerrenoMedia"]),
    ("TemperaturaTerreno2", "Soil temperature sensor 2", "°C", ["TemperaturaTerrenoMedia2"]),
]


def _sensors_by_station(data):
    """{station id: {api_parameter: sensor entry}} from a getAvailableStation*Sensors answer."""
    out = {}
    for item in (data if isinstance(data, list) else [data]):
        for sid, sensors in item.items():
            out.setdefault(int(sid), {}).update({s["name"]: s for s in sensors})
    return out


def build_catalogue(stations, hourly, daily, checked):
    """Catalogue dict from the three metadata answers (no API call).

    Sensors = HOURLY sensors grouped as in SENSOR_GROUPS; a parameter not in
    any group gets its own group (named after it, label/unit from the API).
    DailyOnlyParameters = columns offered only in daily data.
    """
    hourly, daily = _sensors_by_station(hourly), _sensors_by_station(daily)
    known = {p for g in SENSOR_GROUPS for p in g[3]}
    result = []
    for st in sorted(stations, key=lambda s: int(s["StationId"])):
        sid = int(st["StationId"])
        have = hourly.get(sid, {})
        sensors = []
        for key, label, unit, params in SENSOR_GROUPS:
            present = [p for p in params if p in have]
            if present:
                sensors.append({"sensor": key, "name": label, "unit": unit, "api_parameters": present})
        for p in sorted(set(have) - known):
            print(f"  note: station {sid} has new hourly parameter {p!r}; listed as its own sensor")
            sensors.append({"sensor": p, "name": have[p].get("description", p),
                            "unit": have[p].get("unit", ""), "api_parameters": [p]})
        entry = {k: st.get(k) for k in ("StationId", "StationName", "Municipality", "Latitude_WGS84",
                                        "Longitude_WGS84", "Elevation_m", "TimeZone")}
        entry.update(NumSensors=len(sensors), Sensors=sensors,
                     DailyOnlyParameters=sorted(set(daily.get(sid, {})) - set(have)))
        result.append(entry)
    return {
        "source": "https://api.3acloud.it (getAvailableStations, getAvailableStationHourlySensors, "
                  "getAvailableStationDailySensors)",
        "checked": checked,
        "note": "Sensors = HOURLY sensors (what the blocks fetch). DailyOnlyParameters = columns the "
                "station offers only in daily data (/getWeatherDaily). Units from /getAvailableHourlyColumns.",
        "NumStations": len(result),
        "Stations": result,
    }


def cmd_catalogue(args):
    """Ask the server again (3 requests) and rewrite 3a_stations_sensors.json."""
    stations, _ = get_stations(refresh=True)
    ids = sorted(int(s["StationId"]) for s in stations)
    sensors = {kind: cached(f"sensors_{kind.lower()}", "POST", f"/getAvailableStation{kind}Sensors",
                            {"stationsIds": ids}, refresh=True) for kind in ("Hourly", "Daily")}
    old = json.loads(CATALOGUE_FILE.read_text(encoding="utf-8")) if CATALOGUE_FILE.exists() else {}
    new = build_catalogue(stations, sensors["Hourly"], sensors["Daily"], datetime.now().strftime("%Y-%m-%d"))
    params = lambda cat: {s["StationId"]: {p for x in s["Sensors"] for p in x["api_parameters"]}
                          for s in cat.get("Stations", [])}
    before, after = params(old), params(new)
    print("\nChanges in hourly parameters since the last catalogue:")
    for sid in sorted(set(before) | set(after)):
        added, removed = after.get(sid, set()) - before.get(sid, set()), before.get(sid, set()) - after.get(sid, set())
        if added or removed:
            print(f"  {sid}: added {sorted(added)}  removed {sorted(removed)}")
    if before == after:
        print("  none")
    CATALOGUE_FILE.write_text(json.dumps(new, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")
    print(f"\nSaved → {CATALOGUE_FILE}  (update README.md if anything changed)")


def get_columns(kind, refresh=False):
    data = cached(f"columns_{kind.lower()}", "GET", f"/getAvailable{kind}Columns", refresh=refresh)
    if isinstance(data, dict):
        data = next((v for v in data.values() if isinstance(v, list)), data)
    if isinstance(data, list) and data and isinstance(data[0], dict):
        key = next(iter(data[0]))
        data = [d[key] for d in data]
    return data


def cmd_columns(args):
    for kind in ("Hourly", "Daily"):
        cols = get_columns(kind, args.refresh)
        print(f"\n{kind} columns ({len(cols)}):\n  " + ", ".join(map(str, cols)))


def cmd_overview(args):
    cmd_stations(args)
    cmd_sensors(args)
    cmd_columns(args)
    ids = station_ids(args)
    print("\nAvailable periods:")
    for kind in ("Daily", "Hourly"):
        for edge in ("First", "Last"):
            data = cached(f"{edge.lower()}_{kind.lower()}_timestamp", "POST",
                          f"/get{edge}Available{kind}Timestamp", {"stationsIds": ids}, args.refresh)
            print(f"  {edge:5s} {kind:6s}: {json.dumps(data, ensure_ascii=False)[:300]}")


def _parse_date(s):
    for fmt in ("%Y-%m-%d %H:%M", "%Y-%m-%d"):
        try:
            return datetime.strptime(s, fmt)
        except ValueError:
            pass
    sys.exit(f"Bad date {s!r}; use YYYY-MM-DD or 'YYYY-MM-DD HH:MM'")


def _periods(start, end, chunk_days):
    if not chunk_days:
        yield start, end
        return
    cur = start
    while cur <= end:
        nxt = min(cur + timedelta(days=chunk_days) - timedelta(minutes=1), end)
        yield cur, nxt
        cur = nxt + timedelta(minutes=1)


def cmd_weather(args, kind):
    ids = station_ids(args)
    columns = args.columns.split(",") if args.columns else get_columns(kind, args.refresh)
    start, end = _parse_date(args.date_from), _parse_date(args.date_to)
    if len(args.date_to) == 10:
        end = end.replace(hour=23, minute=59)

    frames, raw = [], []
    for a, b in _periods(start, end, args.chunk_days):
        body = {"stationsIds": ids, "dateFrom": a.strftime("%Y-%m-%d %H:%M"),
                "dateTo": b.strftime("%Y-%m-%d %H:%M"), "columns": columns}
        print(f"→ {kind} data {body['dateFrom']} … {body['dateTo']} for {len(ids)} station(s)")
        data = call("POST", f"/getWeather{kind}", body)
        raw.append(data)
        df = to_frame(data)
        if df is not None:
            frames.append(df)
            print(f"   {len(df)} rows")
        else:
            print("   (response not tabular, kept in JSON)", str(data)[:200])

    name = f"weather_{kind.lower()}_{start:%Y%m%d}_{end:%Y%m%d}"
    (OUT / f"{name}.json").write_text(json.dumps(raw if len(raw) > 1 else raw[0], indent=1,
                                                ensure_ascii=False), encoding="utf-8")
    if frames:
        df = pd.concat(frames, ignore_index=True)
        df.to_csv(OUT / f"{name}.csv", index=False, encoding="utf-8-sig")
        print(f"\n{len(df)} rows, {df.shape[1]} columns → {OUT / (name + '.csv')}")
        print(df.head(10).to_string(index=False))


def main():
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("command", choices=["overview", "stations", "sensors", "summary", "catalogue", "columns",
                                        "daily", "hourly", "login"])
    p.add_argument("--from", dest="date_from")
    p.add_argument("--to", dest="date_to")
    p.add_argument("--stations", help="comma-separated station ids (default: all)")
    p.add_argument("--columns", help="comma-separated column names (default: all)")
    p.add_argument("--chunk-days", type=int, help="split the period into N-day requests")
    p.add_argument("--refresh", action="store_true", help="re-download cached metadata")
    args = p.parse_args()

    if args.command in ("daily", "hourly") and not (args.date_from and args.date_to):
        p.error("--from and --to are required for daily/hourly")

    {"overview": cmd_overview, "stations": cmd_stations, "sensors": cmd_sensors, "summary": cmd_summary,
     "catalogue": cmd_catalogue,
     "columns": cmd_columns, "login": lambda a: login(),
     "daily": lambda a: cmd_weather(a, "Daily"), "hourly": lambda a: cmd_weather(a, "Hourly")}[args.command](args)


if __name__ == "__main__":
    main()