#!/usr/bin/env python3
"""Download L2 raw packages referenced by scene-grouped event JSON records."""

from __future__ import annotations

import argparse
import json
import os
import re
import shutil
import subprocess
import sys
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from typing import Any

try:
    from dotenv import load_dotenv
except ImportError:
    def load_dotenv(*args: Any, **kwargs: Any) -> bool:
        return False


FILE = Path(__file__).resolve()
ROOT = FILE.parents[3]
if str(ROOT) not in sys.path:
    sys.path.append(str(ROOT))

from tools.model_inference.adapters.eventid_clip_resolver import (  # noqa: E402
    DEFAULT_EVENT_CACHE_FILE,
    ResolvedEventRecord,
    _extract_condition_values,
    _sanitize_identifier_for_path,
    _select_event_records_by_condition,
    resolve_event_records,
)


TIMESTAMP_RE = re.compile(r"^\d{14}$")


@dataclass
class DownloadResult:
    scene: str
    record_index: int
    rawid: str
    rawid_field_used: str
    condition_values: dict[str, str]
    l2_timestamp: str | None
    mdi_key: str | None
    source_data_path: str | None
    output_dir: str
    expected_download_dir: str | None
    status: str
    detail: str
    command: list[str] | None = None

    def to_dict(self) -> dict[str, Any]:
        return {
            "scene": self.scene,
            "record_index": self.record_index,
            "rawid": self.rawid,
            "rawid_field_used": self.rawid_field_used,
            "condition_values": self.condition_values,
            "l2_timestamp": self.l2_timestamp,
            "mdi_key": self.mdi_key,
            "source_data_path": self.source_data_path,
            "output_dir": self.output_dir,
            "expected_download_dir": self.expected_download_dir,
            "status": self.status,
            "detail": self.detail,
            "command": self.command,
        }


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Download L2 data packages by resolving rawid metadata from an event JSON file."
    )
    parser.add_argument("--event-json-file", required=True)
    parser.add_argument("--scene", default="")
    parser.add_argument("--event-id-field", default="rawid")
    parser.add_argument("--event-clip-ids-field", default="clips")
    parser.add_argument("--condition-fields", nargs="*", default=[])
    parser.add_argument("--max-records-per-condition", type=int, default=0)
    parser.add_argument("--condition-select-strategy", default="first")
    parser.add_argument("--max-events", type=int, default=0)
    parser.add_argument("--event-cache-file", default=str(DEFAULT_EVENT_CACHE_FILE))
    parser.add_argument("--event-resolve-workers", type=int, default=4)
    parser.add_argument("--event-request-timeout", type=float, default=60.0)
    parser.add_argument("--event-request-retries", type=int, default=3)
    parser.add_argument("--event-request-retry-backoff-sec", type=float, default=2.0)
    parser.add_argument("--output-root", required=True)
    parser.add_argument("--manifest-path", default="")
    parser.add_argument("--dry-run", action="store_true")
    parser.add_argument("--skip-done", action="store_true")
    parser.add_argument("--strict", action="store_true")
    return parser.parse_args()


def ensure_pdcl_auth_defaults() -> None:
    load_dotenv()
    os.environ.setdefault("STS_UID", "dis-uploader")
    os.environ.setdefault("STS_SECRET_KEY", "277310cc09724d315514a79701fecb0f")


def log_progress(message: str) -> None:
    timestamp = datetime.now().astimezone().strftime("%Y-%m-%d %H:%M:%S")
    print(f"[download_rawid_l2 {timestamp}] {message}", flush=True)


def extract_l2_timestamp_from_meta(meta: dict[str, Any]) -> tuple[str, str]:
    data_info = meta.get("data_info")
    if isinstance(data_info, dict):
        items = data_info.get("items")
        if isinstance(items, list):
            for item in items:
                if not isinstance(item, dict):
                    continue
                if str(item.get("data_type", "")).strip() != "onboard":
                    continue
                if item.get("available") is False:
                    continue
                data_path = str(item.get("data_path", "")).strip()
                timestamp = Path(data_path).stem
                if TIMESTAMP_RE.fullmatch(timestamp):
                    return timestamp, data_path

    store_path = str(meta.get("store_path", "")).strip()
    timestamp = Path(store_path).stem
    if TIMESTAMP_RE.fullmatch(timestamp):
        return timestamp, store_path

    return "", ""


def load_raw_meta(rawid: str) -> dict[str, Any]:
    ensure_pdcl_auth_defaults()
    from pdcl_dss import Raw

    with Raw(rawid) as raw:
        return dict(raw.meta)


def build_output_dir(output_root: Path, record: ResolvedEventRecord) -> Path:
    rawid_dir = _sanitize_identifier_for_path(record.event_id, prefix=record.event_id_field_used or "rawid")
    return output_root / record.scene / rawid_dir


def write_manifest(
    manifest_path: Path,
    *,
    args: argparse.Namespace,
    selected_records: list[ResolvedEventRecord],
    selection_summary: dict[str, Any],
    results: list[DownloadResult],
) -> None:
    manifest_path.parent.mkdir(parents=True, exist_ok=True)
    payload = {
        "event_json_file": str(Path(args.event_json_file).resolve()),
        "scene": args.scene,
        "event_id_field": args.event_id_field,
        "event_clip_ids_field": args.event_clip_ids_field,
        "output_root": str(Path(args.output_root).resolve()),
        "dry_run": bool(args.dry_run),
        "skip_done": bool(args.skip_done),
        "strict": bool(args.strict),
        "selected_record_count": len(selected_records),
        "selection": selection_summary,
        "summary": summarize_results(results),
        "results": [result.to_dict() for result in results],
    }
    manifest_path.write_text(json.dumps(payload, indent=2, ensure_ascii=False), encoding="utf-8")


def summarize_results(results: list[DownloadResult]) -> dict[str, int]:
    summary: dict[str, int] = {}
    for result in results:
        summary[result.status] = summary.get(result.status, 0) + 1
    return dict(sorted(summary.items()))


def run_mdi_download(mdi_key: str, output_dir: Path, expected_download_dir: Path, *, dry_run: bool, skip_done: bool) -> tuple[str, str, list[str]]:
    command = ["mdi", "raw", "-r", mdi_key, "-s", str(output_dir)]

    if skip_done and expected_download_dir.exists():
        return "exists", f"target already exists: {expected_download_dir}", command

    if dry_run:
        return "planned", f"would run {' '.join(command)}", command

    output_dir.mkdir(parents=True, exist_ok=True)
    if shutil.which("mdi") is None:
        return "failed", "mdi command not found in PATH", command

    completed = subprocess.run(
        command,
        check=False,
        capture_output=True,
        text=True,
        encoding="utf-8",
    )
    if completed.returncode == 0:
        return "downloaded", completed.stdout.strip() or "mdi raw completed", command

    detail = completed.stderr.strip() or completed.stdout.strip() or "mdi raw failed"
    return "failed", detail, command


def download_one_record(
    record: ResolvedEventRecord,
    *,
    output_root: Path,
    condition_fields: list[str],
    dry_run: bool,
    skip_done: bool,
) -> DownloadResult:
    rawid = record.event_id
    output_dir = build_output_dir(output_root, record)
    condition_values = _extract_condition_values(record.source_record, condition_fields)

    try:
        meta = load_raw_meta(rawid)
        l2_timestamp, source_data_path = extract_l2_timestamp_from_meta(meta)
    except Exception as exc:
        return DownloadResult(
            scene=record.scene,
            record_index=record.record_index,
            rawid=rawid,
            rawid_field_used=record.event_id_field_used,
            condition_values=condition_values,
            l2_timestamp=None,
            mdi_key=None,
            source_data_path=None,
            output_dir=str(output_dir),
            expected_download_dir=None,
            status="failed_meta",
            detail=f"{type(exc).__name__}: {exc}",
        )

    if not l2_timestamp:
        return DownloadResult(
            scene=record.scene,
            record_index=record.record_index,
            rawid=rawid,
            rawid_field_used=record.event_id_field_used,
            condition_values=condition_values,
            l2_timestamp=None,
            mdi_key=None,
            source_data_path=source_data_path or None,
            output_dir=str(output_dir),
            expected_download_dir=None,
            status="failed_no_l2_timestamp",
            detail="no onboard L2 timestamp was found in Raw.meta",
        )

    mdi_key = f"{rawid}::{l2_timestamp}"
    expected_download_dir = output_dir / l2_timestamp
    status, detail, command = run_mdi_download(
        mdi_key,
        output_dir,
        expected_download_dir,
        dry_run=dry_run,
        skip_done=skip_done,
    )
    return DownloadResult(
        scene=record.scene,
        record_index=record.record_index,
        rawid=rawid,
        rawid_field_used=record.event_id_field_used,
        condition_values=condition_values,
        l2_timestamp=l2_timestamp,
        mdi_key=mdi_key,
        source_data_path=source_data_path,
        output_dir=str(output_dir),
        expected_download_dir=str(expected_download_dir),
        status=status,
        detail=detail,
        command=command,
    )


def main() -> None:
    args = parse_args()
    condition_fields = [str(field).strip() for field in args.condition_fields if str(field).strip()]
    output_root = Path(args.output_root).resolve()
    manifest_path = (
        Path(args.manifest_path).resolve()
        if args.manifest_path
        else output_root / "download_manifest.json"
    )

    resolved_records, _ = resolve_event_records(
        json_file=args.event_json_file,
        scene_filter=args.scene or None,
        event_id_field=args.event_id_field,
        clip_ids_field=args.event_clip_ids_field,
        max_events=args.max_events,
        timeout=args.event_request_timeout,
        cache_file=args.event_cache_file,
        workers=args.event_resolve_workers,
        max_retries=args.event_request_retries,
        retry_backoff_sec=args.event_request_retry_backoff_sec,
    )
    selected_records, selection_summary = _select_event_records_by_condition(
        resolved_records,
        condition_fields=condition_fields,
        max_records_per_condition=args.max_records_per_condition,
        selection_strategy=args.condition_select_strategy,
    )

    log_progress(
        f"selected {len(selected_records)} rawid record(s) from {len(resolved_records)} resolved record(s)"
    )

    results: list[DownloadResult] = []
    for index, record in enumerate(selected_records, start=1):
        log_progress(f"[{index}/{len(selected_records)}] {record.scene} {record.event_id}")
        result = download_one_record(
            record,
            output_root=output_root,
            condition_fields=condition_fields,
            dry_run=args.dry_run,
            skip_done=args.skip_done,
        )
        results.append(result)
        log_progress(f"  -> {result.status}: {result.mdi_key or result.detail}")
        write_manifest(
            manifest_path,
            args=args,
            selected_records=selected_records,
            selection_summary=selection_summary,
            results=results,
        )

    summary = summarize_results(results)
    log_progress(f"manifest: {manifest_path}")
    log_progress("summary: " + ", ".join(f"{key}={value}" for key, value in summary.items()))

    if args.strict and any(result.status.startswith("failed") for result in results):
        raise SystemExit(1)


if __name__ == "__main__":
    main()
