from __future__ import annotations

from datetime import date, timedelta
import json
from urllib.parse import urlencode, urlparse
from urllib.request import Request, urlopen

import pendulum
import psycopg2
from airflow.sdk import Connection, Param, dag, get_current_context, task


BANGKOK_LAT = 13.7563
BANGKOK_LON = 100.5018


def postgres_connection_kwargs() -> dict[str, object]:
    """Read a named Airflow Connection without embedding credentials in DAG code."""
    connection = Connection.get("demo_postgres")
    parsed = urlparse(connection.get_uri())
    return {
        "host": parsed.hostname,
        "port": parsed.port or 5432,
        "dbname": parsed.path.lstrip("/"),
        "user": parsed.username,
        "password": parsed.password,
    }


@dag(
    dag_id="daily_weather_to_postgres",
    schedule=None,  # Manual trigger only: there is no cron schedule.
    start_date=pendulum.datetime(2026, 1, 1, tz="Asia/Bangkok"),
    catchup=False,
    tags=["warin.me", "teaching", "manual", "database"],
    params={
        "data_date": Param(
            (date.today() - timedelta(days=7)).isoformat(),
            type="string",
            format="date",
            title="Data date",
            description="One historical day to download (YYYY-MM-DD)",
        ),
        "latitude": Param(BANGKOK_LAT, type="number", minimum=-90, maximum=90),
        "longitude": Param(BANGKOK_LON, type="number", minimum=-180, maximum=180),
    },
)
def daily_weather_to_postgres():
    """Download one requested day and idempotently load it into PostgreSQL."""

    @task
    def extract() -> dict:
        params = get_current_context()["params"]
        selected_date = date.fromisoformat(str(params["data_date"]))
        if selected_date >= date.today():
            raise ValueError("Choose a completed historical day before today")

        query = urlencode(
            {
                "latitude": float(params["latitude"]),
                "longitude": float(params["longitude"]),
                "start_date": selected_date.isoformat(),
                "end_date": selected_date.isoformat(),
                "daily": ",".join(
                    [
                        "temperature_2m_max",
                        "temperature_2m_min",
                        "precipitation_sum",
                        "wind_speed_10m_max",
                    ]
                ),
                "timezone": "Asia/Bangkok",
            }
        )
        url = f"https://archive-api.open-meteo.com/v1/archive?{query}"
        request = Request(url, headers={"User-Agent": "warin-airflow-teaching-lab/1.0"})
        with urlopen(request, timeout=30) as response:
            payload = json.load(response)

        daily = payload.get("daily", {})
        if not daily.get("time"):
            raise ValueError(f"The source returned no daily record for {selected_date}")
        return {
            "observed_date": daily["time"][0],
            "latitude": round(float(payload["latitude"]), 4),
            "longitude": round(float(payload["longitude"]), 4),
            "temperature_max_c": daily["temperature_2m_max"][0],
            "temperature_min_c": daily["temperature_2m_min"][0],
            "precipitation_mm": daily["precipitation_sum"][0],
            "wind_speed_max_kmh": daily["wind_speed_10m_max"][0],
            "source_url": url,
        }

    @task
    def validate(row: dict) -> dict:
        if row["temperature_min_c"] > row["temperature_max_c"]:
            raise ValueError("Minimum temperature exceeds maximum temperature")
        if row["precipitation_mm"] < 0 or row["wind_speed_max_kmh"] < 0:
            raise ValueError("Precipitation and wind speed cannot be negative")
        return row

    @task
    def load(row: dict) -> dict:
        create_sql = """
            CREATE TABLE IF NOT EXISTS teaching_daily_weather (
                observed_date date NOT NULL,
                latitude numeric(8,4) NOT NULL,
                longitude numeric(8,4) NOT NULL,
                temperature_max_c numeric(6,2),
                temperature_min_c numeric(6,2),
                precipitation_mm numeric(8,2),
                wind_speed_max_kmh numeric(8,2),
                source_url text NOT NULL,
                loaded_at timestamptz NOT NULL DEFAULT now(),
                PRIMARY KEY (observed_date, latitude, longitude)
            )
        """
        upsert_sql = """
            INSERT INTO teaching_daily_weather (
                observed_date, latitude, longitude, temperature_max_c,
                temperature_min_c, precipitation_mm, wind_speed_max_kmh, source_url
            ) VALUES (
                %(observed_date)s, %(latitude)s, %(longitude)s, %(temperature_max_c)s,
                %(temperature_min_c)s, %(precipitation_mm)s, %(wind_speed_max_kmh)s,
                %(source_url)s
            )
            ON CONFLICT (observed_date, latitude, longitude) DO UPDATE SET
                temperature_max_c = EXCLUDED.temperature_max_c,
                temperature_min_c = EXCLUDED.temperature_min_c,
                precipitation_mm = EXCLUDED.precipitation_mm,
                wind_speed_max_kmh = EXCLUDED.wind_speed_max_kmh,
                source_url = EXCLUDED.source_url,
                loaded_at = now()
        """
        with psycopg2.connect(**postgres_connection_kwargs()) as connection:
            with connection.cursor() as cursor:
                cursor.execute(create_sql)
                cursor.execute(upsert_sql, row)
                cursor.execute(
                    """SELECT count(*) FROM teaching_daily_weather
                       WHERE observed_date = %s AND latitude = %s AND longitude = %s""",
                    (row["observed_date"], row["latitude"], row["longitude"]),
                )
                stored_rows = cursor.fetchone()[0]
        return {"table": "teaching_daily_weather", "stored_rows_for_key": stored_rows, **row}

    load(validate(extract()))


daily_weather_to_postgres()