"""gharchive_hourly — Module 2 Lesson 06 sample DAG (L12: OpenLineage + freshness).

Idempotent: extract writes to a staging path then atomically renames to the
final path. Replaying any interval is safe — the final file is never
half-written. Retry budgets are configured per task character
(external API vs pure compute). XCom passes the extract output path to
transform (metadata only, not data).

L12 additions:
  - Every PythonOperator task emits OpenLineage START / COMPLETE events via
    pre_execute / post_execute hooks. The backend is configured via the
    OPENLINEAGE_URL env var; defaults to file:// fallback when Marquez is
    not running.
  - A final `freshness_check` PythonOperator runs `observability.freshness`
    and raises RuntimeError on breach, failing the DAG run visibly.
"""

from __future__ import annotations

import os
import sys
import pathlib
from datetime import datetime

from airflow import DAG
from airflow.operators.python import PythonOperator

_SAMPLE_PIPELINE_ROOT = pathlib.Path("/opt/airflow/sample-pipeline")
if _SAMPLE_PIPELINE_ROOT.exists() and str(_SAMPLE_PIPELINE_ROOT) not in sys.path:
    sys.path.insert(0, str(_SAMPLE_PIPELINE_ROOT))

# _helpers lives next to this DAG file in /opt/airflow/dags
sys.path.insert(0, str(pathlib.Path(__file__).parent))
from _helpers import (  # noqa: E402
    atomic_replace,
    transactional_write_path,
    EXTERNAL_API_RETRIES,
    PURE_COMPUTE_RETRIES,
)

# OpenLineage helpers — imported lazily so the DAG loads even when the
# observability package is not on sys.path (e.g. unit-test environments).
def _ol_emit_start(context: dict) -> None:  # pre_execute hook
    try:
        from observability.openlineage_emit import emit_start, new_run_id
        task_id = context["task_instance"].task_id
        run_id = new_run_id()
        context["task_instance"].xcom_push(key=f"_ol_run_id_{task_id}", value=run_id)
        emit_start(run_id, f"gharchive_hourly.{task_id}")
    except Exception:  # noqa: BLE001 — never fail the pipeline on OL errors
        pass


def _ol_emit_complete(context: dict) -> None:  # post_execute hook
    try:
        from observability.openlineage_emit import emit_complete
        task_id = context["task_instance"].task_id
        run_id = context["task_instance"].xcom_pull(key=f"_ol_run_id_{task_id}")
        if run_id:
            emit_complete(run_id, f"gharchive_hourly.{task_id}")
    except Exception:  # noqa: BLE001
        pass


def _hour_key(context) -> str:
    """Format the data interval start as YYYY-MM-DD-H (gharchive filename shape)."""
    return context["data_interval_start"].strftime("%Y-%m-%d-%-H")


def _extract(**context) -> str:
    """Download one hour from gharchive, write via staging-then-rename."""
    from extract.gharchive import fetch_hour

    hour = _hour_key(context)
    data_dir = pathlib.Path("/opt/airflow/sample-pipeline/data/raw")
    data_dir.mkdir(parents=True, exist_ok=True)

    # fetch_hour writes <hour>.json.gz into out_dir and returns the path
    final = data_dir / f"{hour}.json.gz"
    staging_dir = data_dir / f".staging"
    staging_dir.mkdir(parents=True, exist_ok=True)
    raw_staging = fetch_hour(hour=hour, out_dir=str(staging_dir))
    staging = pathlib.Path(raw_staging)
    atomic_replace(staging, final)

    # Push the final path to XCom so the transform task knows where to read.
    return str(final)


def _transform(**context) -> None:
    """Count events by type; write output via staging-then-rename."""
    from transform.events import count_by_type

    hour = _hour_key(context)
    input_path = context["task_instance"].xcom_pull(task_ids="extract")

    out_dir = pathlib.Path("/opt/airflow/sample-pipeline/data/out")
    out_dir.mkdir(parents=True, exist_ok=True)
    final = out_dir / f"{hour}.csv"
    staging = transactional_write_path(final, context["run_id"])

    count_by_type(input_path=input_path, output_path=str(staging))
    atomic_replace(staging, final)


def _freshness_check(**_context) -> None:
    """Run the observability freshness SLA check; raise on breach.

    L12: this task is the final step in the DAG. A BREACH causes the task
    — and therefore the DAG run — to fail visibly, ensuring on-call is paged
    rather than a silent stale-data window opening.
    """
    try:
        from observability.freshness import check
        from observability.alerts import alert
    except ImportError as exc:
        # observability package not on path in this environment — skip silently.
        print(f"freshness_check: observability package unavailable ({exc}), skipping.")
        return

    rc = check(sla_minutes=90)
    if rc == 1:
        alert("ERROR", "freshness", "events table SLA breach", dag="gharchive_hourly")
        raise RuntimeError("Freshness SLA breach: events table age exceeds 90 minutes.")
    if rc == 2:
        alert("WARN", "freshness", "events table UNKNOWN — may be empty",
              dag="gharchive_hourly")


with DAG(
    dag_id="gharchive_hourly",
    start_date=datetime(2026, 5, 1),
    schedule="@hourly",
    catchup=True,
    max_active_runs=2,
    tags=["data-engineering", "module-2", "module-5"],
    default_args={"owner": "data-eng"},
) as dag:
    extract = PythonOperator(
        task_id="extract",
        python_callable=_extract,
        pre_execute=_ol_emit_start,
        post_execute=_ol_emit_complete,
        **EXTERNAL_API_RETRIES,
    )
    transform = PythonOperator(
        task_id="transform",
        python_callable=_transform,
        pre_execute=_ol_emit_start,
        post_execute=_ol_emit_complete,
        **PURE_COMPUTE_RETRIES,
    )
    freshness_check = PythonOperator(
        task_id="freshness_check",
        python_callable=_freshness_check,
        pre_execute=_ol_emit_start,
        post_execute=_ol_emit_complete,
        **PURE_COMPUTE_RETRIES,
    )
    extract >> transform >> freshness_check