"""
Download hydrometric data from Hub'Eau API v2.
Downloads QmJ (daily discharge), HIJ (daily water level) for all 27 stations.
"""
import requests
import pandas as pd
from pathlib import Path
from tqdm import tqdm
import time

# Configuration
OUTPUT_DIR = Path("datasets/hydrometric")
OUTPUT_DIR.mkdir(exist_ok=True)

BASE_URL = "https://hubeau.eaufrance.fr/api/v2/hydrometrie"

# Load stations
stations_df = pd.read_csv("datasets/station_list.csv")
station_codes = stations_df["station_code"].tolist()

print(f"Downloading hydrometric data for {len(station_codes)} stations...")
print(f"Output: {OUTPUT_DIR}")
print()

def download_observations(station_code, grandeur_hydro="Q", max_results=100000):
    """
    Download observations for a station.

    Args:
        station_code: Station code (e.g., H602021010)
        grandeur_hydro: H (water level) or Q (discharge)
        max_results: Maximum results to fetch

    Returns:
        DataFrame with observations
    """
    url = f"{BASE_URL}/obs_elab"

    all_data = []
    cursor = ""
    page = 1

    while len(all_data) < max_results:
        params = {
            "code_entite": station_code,
            "grandeur_hydro": grandeur_hydro,
            "size": 20000,  # Max per request
            "format": "json",
            "cursor": cursor
        }

        try:
            response = requests.get(url, params=params, timeout=30)

            if response.status_code == 200:
                data = response.json()
                results = data.get("data", [])

                if not results:
                    break

                all_data.extend(results)

                # Check for next page
                next_url = data.get("next")
                if not next_url or len(results) < 20000:
                    break

                # Extract cursor from next URL
                if "cursor=" in next_url:
                    cursor = next_url.split("cursor=")[1].split("&")[0]
                else:
                    break

                page += 1
                time.sleep(0.1)  # Rate limiting

            else:
                print(f"  Error {response.status_code}")
                break

        except Exception as e:
            print(f"  Exception: {e}")
            break

    if all_data:
        return pd.DataFrame(all_data)
    else:
        return pd.DataFrame()


# Download discharge (Q) for all stations
print("="*80)
print("DOWNLOADING DISCHARGE (Q)")
print("="*80)

discharge_data = []

for station_code in tqdm(station_codes, desc="Discharge"):
    df = download_observations(station_code, grandeur_hydro="Q")

    if not df.empty:
        df["station_code"] = station_code
        discharge_data.append(df)
        tqdm.write(f"✓ {station_code}: {len(df)} observations")
    else:
        tqdm.write(f"✗ {station_code}: No data")

    time.sleep(0.2)

if discharge_data:
    discharge_df = pd.concat(discharge_data, ignore_index=True)

    # Save
    discharge_df.to_csv(OUTPUT_DIR / "discharge_observations.csv", index=False)
    discharge_df.to_json(OUTPUT_DIR / "discharge_observations.json", orient="records", indent=2)

    print(f"\n✓ Discharge: {len(discharge_df)} total observations")
    print(f"  Stations: {discharge_df['station_code'].nunique()}")
    print(f"  Date range: {discharge_df['date_obs_elab'].min()} to {discharge_df['date_obs_elab'].max()}")


# Download water level (H) for all stations
print("\n" + "="*80)
print("DOWNLOADING WATER LEVEL (H)")
print("="*80)

waterlevel_data = []

for station_code in tqdm(station_codes, desc="Water Level"):
    df = download_observations(station_code, grandeur_hydro="H")

    if not df.empty:
        df["station_code"] = station_code
        waterlevel_data.append(df)
        tqdm.write(f"✓ {station_code}: {len(df)} observations")
    else:
        tqdm.write(f"✗ {station_code}: No data")

    time.sleep(0.2)

if waterlevel_data:
    waterlevel_df = pd.concat(waterlevel_data, ignore_index=True)

    # Save
    waterlevel_df.to_csv(OUTPUT_DIR / "waterlevel_observations.csv", index=False)
    waterlevel_df.to_json(OUTPUT_DIR / "waterlevel_observations.json", orient="records", indent=2)

    print(f"\n✓ Water level: {len(waterlevel_df)} total observations")
    print(f"  Stations: {waterlevel_df['station_code'].nunique()}")
    print(f"  Date range: {waterlevel_df['date_obs_elab'].min()} to {waterlevel_df['date_obs_elab'].max()}")


# Summary
print("\n" + "="*80)
print("DOWNLOAD COMPLETE")
print("="*80)
print(f"Files saved to: {OUTPUT_DIR}")
print(f"\nFiles created:")
for f in OUTPUT_DIR.glob("*"):
    size_mb = f.stat().st_size / 1024 / 1024
    print(f"  - {f.name} ({size_mb:.1f} MB)")