"""Open-Meteo daily weather extraction for coffee-growing regions.

Two entry points:

  extract_weather()
      Daily run: fetches the last 10 days for all 12 locations.
      10 days covers the ~5-day ERA5 reanalysis lag plus buffer for missed runs.
      Uses the forecast API (fills the recent window not yet in ERA5 archive).
      12 API calls total. Completes in ~10 seconds.

  extract_weather_backfill()
      Historical fill: fetches 2020-01-01 → yesterday for all 12 locations.
      Uses the archive API. One date-range request per location = 12 total calls.
      Completes in ~30 seconds. No cursor needed.

Landing path: LANDING_DIR/weather/{location_id}/{year}/{date}.json.gz

Each file is a flat JSON object with Open-Meteo variable names:
  {"date": "2020-01-01", "temperature_2m_max": 28.5, "precipitation_sum": 12.5, ...}

No API key required. No rate limiting. Fully idempotent (file existence check).
"""

import gzip
import json
import logging
import os
import sys
import time
from datetime import date, timedelta
from pathlib import Path

import niquests
from extract_core import end_run, landing_path, open_state_db, start_run, write_bytes_atomic

from openmeteo.api import fetch_archive, fetch_recent
from openmeteo.locations import LOCATIONS

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
    handlers=[logging.StreamHandler(sys.stdout)],
)
logger = logging.getLogger("Open-Meteo Extractor")

LANDING_DIR = Path(os.getenv("LANDING_DIR", "data/landing"))
LANDING_SUBDIR = "weather"

EXTRACTOR_DAILY = "openmeteo_daily"
EXTRACTOR_BACKFILL = "openmeteo_backfill"

BACKFILL_START = date(2020, 1, 1)

# Small sleep between location calls — polite usage of the free community API.
SLEEP_BETWEEN_LOCATIONS_SECONDS = 0.5


# ── helpers ───────────────────────────────────────────────────────────────────

def _write_day_file(location_id: str, date_str: str, record: dict) -> int:
    """Write one day's weather record as gzipped JSON. Returns bytes written, or 0 if skipped."""
    assert location_id and date_str and record

    year = date_str[:4]
    dest_dir = landing_path(LANDING_DIR, LANDING_SUBDIR, location_id, year)
    local_file = dest_dir / f"{date_str}.json.gz"

    if local_file.exists():
        return 0

    compressed = gzip.compress(json.dumps(record, separators=(",", ":")).encode("utf-8"))
    bytes_written = write_bytes_atomic(local_file, compressed)
    logger.debug(f"Stored {local_file} ({bytes_written:,} bytes)")
    return bytes_written


def _count_existing_files(location_id: str, start: date, end: date) -> int:
    """Count per-day files already on disk for a location between start and end (inclusive)."""
    count = 0
    d = start
    while d <= end:
        year = d.strftime("%Y")
        dest_dir = landing_path(LANDING_DIR, LANDING_SUBDIR, location_id, year)
        if (dest_dir / f"{d.isoformat()}.json.gz").exists():
            count += 1
        d += timedelta(days=1)
    return count


def _split_and_write(location_id: str, response: dict) -> tuple[int, int, int]:
    """Split an Open-Meteo array response into per-day JSON.gz files.

    Open-Meteo returns parallel arrays per variable under response['daily']['time'].
    We zip these into one flat dict per day and write each as a separate file.

    Returns (files_written, files_skipped, bytes_written).
    """
    daily = response["daily"]
    dates = daily["time"]
    variables = [k for k in daily if k != "time"]

    files_written = 0
    files_skipped = 0
    bytes_written_total = 0

    for i, date_str in enumerate(dates):
        if not date_str:
            continue

        record = {"date": date_str}
        for var in variables:
            values = daily[var]
            record[var] = values[i] if i < len(values) else None

        bw = _write_day_file(location_id, date_str, record)
        if bw > 0:
            files_written += 1
            bytes_written_total += bw
        else:
            files_skipped += 1

    return files_written, files_skipped, bytes_written_total


# ── daily extractor ───────────────────────────────────────────────────────────

def extract_weather() -> None:
    """Fetch the last 10 days of weather for all 12 locations.

    Uses Open-Meteo forecast API (past_days=10). 12 API calls. ~10 seconds.
    """
    conn = open_state_db(LANDING_DIR)
    run_id = start_run(conn, EXTRACTOR_DAILY)
    files_written = 0
    files_skipped = 0
    bytes_written_total = 0

    try:
        with niquests.Session() as session:
            for loc in LOCATIONS:
                logger.info(f"Fetching recent: {loc['id']} ({loc['country']})")
                response = fetch_recent(session, loc["lat"], loc["lon"], past_days=10)
                w, s, bw = _split_and_write(loc["id"], response)
                files_written += w
                files_skipped += s
                bytes_written_total += bw
                time.sleep(SLEEP_BETWEEN_LOCATIONS_SECONDS)

        end_run(
            conn, run_id,
            status="success",
            files_written=files_written,
            files_skipped=files_skipped,
            bytes_written=bytes_written_total,
            cursor_value=date.today().isoformat(),
        )
        logger.info(f"Daily weather complete: {files_written} new, {files_skipped} skipped")
    except Exception as e:
        end_run(conn, run_id, status="failed", error_message=str(e))
        raise
    finally:
        conn.close()


# ── backfill extractor ────────────────────────────────────────────────────────

def extract_weather_backfill() -> None:
    """Fetch full weather history (2020-01-01 → yesterday) for all 12 locations.

    Uses Open-Meteo archive API (ERA5 reanalysis). One date-range request per
    location = 12 calls total. Completes in ~30 seconds.

    Idempotent: per-day files already on disk are skipped when splitting the
    response. Safe to re-run at any time — will skip what already exists.
    """
    yesterday = (date.today() - timedelta(days=1)).isoformat()
    start_date = BACKFILL_START.isoformat()

    conn = open_state_db(LANDING_DIR)
    run_id = start_run(conn, EXTRACTOR_BACKFILL)
    files_written = 0
    files_skipped = 0
    bytes_written_total = 0

    try:
        start = BACKFILL_START
        end = date.fromisoformat(yesterday)
        expected_days = (end - start).days + 1

        with niquests.Session() as session:
            for loc in LOCATIONS:
                logger.info(
                    f"Backfill {loc['id']} ({loc['country']}) "
                    f"{start_date} → {yesterday}"
                )
                existing = _count_existing_files(loc["id"], start, end)
                if existing == expected_days:
                    logger.info(f"  {loc['id']}: 0 new, {existing} already existed (skipped API call)")
                    files_skipped += existing
                    continue

                response = fetch_archive(
                    session, loc["lat"], loc["lon"],
                    start_date=start_date,
                    end_date=yesterday,
                )
                w, s, bw = _split_and_write(loc["id"], response)
                files_written += w
                files_skipped += s
                bytes_written_total += bw
                logger.info(f"  {loc['id']}: {w} new, {s} already existed")
                time.sleep(SLEEP_BETWEEN_LOCATIONS_SECONDS)

        end_run(
            conn, run_id,
            status="success",
            files_written=files_written,
            files_skipped=files_skipped,
            bytes_written=bytes_written_total,
            cursor_value=yesterday,
        )
        logger.info(
            f"Backfill complete: {files_written} new files, "
            f"{files_skipped} already existed"
        )
    except Exception as e:
        end_run(conn, run_id, status="failed", error_message=str(e))
        raise
    finally:
        conn.close()


if __name__ == "__main__":
    extract_weather()
