#!/usr/bin/env python3
"""
Download completed NSRDB county polygon archive ZIP files.

This shared engine reads JSON acknowledgements produced by the polygon request
workflow and downloads each outputs.downloadUrl. Metric-specific wrappers
supply the response and archive directories.
"""

from __future__ import annotations

import argparse
import json
import time
import urllib.error
import urllib.parse
import urllib.request
from dataclasses import dataclass, field
from enum import Enum
from pathlib import Path
from typing import Any


DEFAULT_TIMEOUT = 300
DEFAULT_CHUNK_SIZE = 1024 * 1024


class ArchiveDownloadState(Enum):
    """Distinct lifecycle phases for one NSRDB archive download."""

    QUEUED = "queued"
    CHECKING = "checking"
    DOWNLOADING = "downloading"
    SKIPPED = "skipped"
    PENDING = "pending"
    DOWNLOADED = "downloaded"
    FAILED = "failed"


class ArchiveDownloadEvent(Enum):
    """Events that may move an archive download to another lifecycle phase."""

    START = "start"
    EXISTING_FILE_FOUND = "existing_file_found"
    DOWNLOAD_STARTED = "download_started"
    ARCHIVE_NOT_READY = "archive_not_ready"
    DOWNLOAD_SUCCEEDED = "download_succeeded"
    DOWNLOAD_FAILED = "download_failed"


ARCHIVE_DOWNLOAD_TRANSITIONS = {
    (ArchiveDownloadState.QUEUED, ArchiveDownloadEvent.START): ArchiveDownloadState.CHECKING,
    (
        ArchiveDownloadState.CHECKING,
        ArchiveDownloadEvent.EXISTING_FILE_FOUND,
    ): ArchiveDownloadState.SKIPPED,
    (
        ArchiveDownloadState.CHECKING,
        ArchiveDownloadEvent.DOWNLOAD_STARTED,
    ): ArchiveDownloadState.DOWNLOADING,
    (
        ArchiveDownloadState.DOWNLOADING,
        ArchiveDownloadEvent.ARCHIVE_NOT_READY,
    ): ArchiveDownloadState.PENDING,
    (
        ArchiveDownloadState.DOWNLOADING,
        ArchiveDownloadEvent.DOWNLOAD_SUCCEEDED,
    ): ArchiveDownloadState.DOWNLOADED,
    (
        ArchiveDownloadState.QUEUED,
        ArchiveDownloadEvent.DOWNLOAD_FAILED,
    ): ArchiveDownloadState.FAILED,
    (
        ArchiveDownloadState.CHECKING,
        ArchiveDownloadEvent.DOWNLOAD_FAILED,
    ): ArchiveDownloadState.FAILED,
    (
        ArchiveDownloadState.DOWNLOADING,
        ArchiveDownloadEvent.DOWNLOAD_FAILED,
    ): ArchiveDownloadState.FAILED,
}


class InvalidArchiveDownloadTransition(RuntimeError):
    """Reject an event that is not valid for the current download state."""


@dataclass
class ArchiveDownloadStateMachine:
    """Track and validate the lifecycle of one archive download."""

    label: str
    state: ArchiveDownloadState = ArchiveDownloadState.QUEUED
    history: list[ArchiveDownloadState] = field(
        default_factory=lambda: [ArchiveDownloadState.QUEUED]
    )

    def transition(self, event: ArchiveDownloadEvent) -> ArchiveDownloadState:
        """Apply one event and return the resulting state."""
        next_state = ARCHIVE_DOWNLOAD_TRANSITIONS.get((self.state, event))
        if next_state is None:
            raise InvalidArchiveDownloadTransition(
                f"Invalid archive download transition for {self.label}: "
                f"{event.value} while {self.state.value}."
            )
        self.state = next_state
        self.history.append(next_state)
        return next_state


class PolygonArchiveDownloadError(RuntimeError):
    """Store a failed archive download with useful context."""

    def __init__(self, message: str, status_code: int | None = None) -> None:
        super().__init__(message)
        self.status_code = status_code


def is_pending_s3_archive_response(url: str, status_code: int) -> bool:
    """Return whether an S3 HTTP response likely means the archive is pending."""
    if status_code != 403:
        return False

    host = urllib.parse.urlsplit(url).netloc.lower()
    return host == "s3.amazonaws.com" or host.endswith(".amazonaws.com")


def read_response(response_path: Path) -> dict[str, Any]:
    """Read a saved NSRDB archive response JSON file."""
    with response_path.open(encoding="utf-8") as handle:
        response = json.load(handle)

    if not isinstance(response, dict):
        raise ValueError("Response JSON must contain an object.")

    return response


def download_url_from_response(response: dict[str, Any]) -> str:
    """Return the download URL from a saved NSRDB archive response."""
    outputs = response.get("outputs")
    if not isinstance(outputs, dict):
        raise ValueError("Response JSON is missing an outputs object.")

    download_url = outputs.get("downloadUrl")
    if not isinstance(download_url, str) or not download_url.strip():
        raise ValueError("Response JSON is missing outputs.downloadUrl.")

    return download_url.strip()


def output_name_for_response(response_path: Path) -> str:
    """Convert a response JSON filename into the corresponding ZIP filename."""
    stem = response_path.stem
    if stem.endswith("_response"):
        stem = stem[: -len("_response")]
    return f"{stem}.zip"


def download_file(
    url: str,
    output_path: Path,
    timeout: int,
    overwrite: bool,
    machine: ArchiveDownloadStateMachine,
) -> ArchiveDownloadState:
    """Download a URL to disk and return the terminal download state."""
    output_path.parent.mkdir(parents=True, exist_ok=True)

    if output_path.exists() and not overwrite:
        return machine.transition(ArchiveDownloadEvent.EXISTING_FILE_FOUND)

    machine.transition(ArchiveDownloadEvent.DOWNLOAD_STARTED)
    temp_path = output_path.with_suffix(f"{output_path.suffix}.part")
    if temp_path.exists():
        temp_path.unlink()

    request = urllib.request.Request(url, headers={"User-Agent": "county-climate-explorer/0.1"})
    try:
        with urllib.request.urlopen(request, timeout=timeout) as response:
            with temp_path.open("wb") as handle:
                while True:
                    chunk = response.read(DEFAULT_CHUNK_SIZE)
                    if not chunk:
                        break
                    handle.write(chunk)
    except urllib.error.HTTPError as error:
        if temp_path.exists():
            temp_path.unlink()
        if is_pending_s3_archive_response(url, error.code):
            return machine.transition(ArchiveDownloadEvent.ARCHIVE_NOT_READY)
        machine.transition(ArchiveDownloadEvent.DOWNLOAD_FAILED)
        raise PolygonArchiveDownloadError(f"HTTP {error.code} {error.reason}", error.code) from error
    except urllib.error.URLError as error:
        if temp_path.exists():
            temp_path.unlink()
        machine.transition(ArchiveDownloadEvent.DOWNLOAD_FAILED)
        raise PolygonArchiveDownloadError(str(error.reason)) from error

    temp_path.replace(output_path)
    return machine.transition(ArchiveDownloadEvent.DOWNLOAD_SUCCEEDED)


def response_paths(response_dir: Path, start: int, limit: int | None) -> list[Path]:
    """Return saved response JSON files in deterministic order."""
    paths = sorted(response_dir.glob("*.json"))

    if start < 0:
        raise ValueError("--start must be 0 or greater.")
    if limit is not None and limit < 1:
        raise ValueError("--limit must be 1 or greater.")

    if limit is None:
        return paths[start:]
    return paths[start : start + limit]


def run(args: argparse.Namespace) -> None:
    """Download all requested NSRDB polygon archives."""
    paths = response_paths(args.response_dir, args.start, args.limit)
    if not paths:
        print(f"No response JSON files found in {args.response_dir}")
        return

    counts = {
        ArchiveDownloadState.DOWNLOADED: 0,
        ArchiveDownloadState.SKIPPED: 0,
        ArchiveDownloadState.PENDING: 0,
        ArchiveDownloadState.FAILED: 0,
    }
    total = len(paths)

    for index, response_path in enumerate(paths, start=1):
        output_path = args.output_dir / output_name_for_response(response_path)
        label = response_path.name
        machine = ArchiveDownloadStateMachine(label)
        machine.transition(ArchiveDownloadEvent.START)

        try:
            response = read_response(response_path)
            download_url = download_url_from_response(response)
            state = download_file(
                download_url,
                output_path,
                args.timeout,
                args.overwrite,
                machine,
            )
        except (OSError, ValueError, PolygonArchiveDownloadError) as error:
            if machine.state not in {
                ArchiveDownloadState.FAILED,
                ArchiveDownloadState.SKIPPED,
                ArchiveDownloadState.PENDING,
                ArchiveDownloadState.DOWNLOADED,
            }:
                machine.transition(ArchiveDownloadEvent.DOWNLOAD_FAILED)
            counts[machine.state] += 1
            print(f"[{index}/{total}] Failed {label}: {error}")
            continue

        counts[state] += 1
        if state is ArchiveDownloadState.SKIPPED:
            print(f"[{index}/{total}] Skipped existing {output_path}")
        elif state is ArchiveDownloadState.PENDING:
            print(
                f"[{index}/{total}] Pending {label}: archive URL returned "
                "S3 HTTP 403; retry after NSRDB finishes generating it."
            )
        else:
            print(f"[{index}/{total}] Downloaded {output_path}")

        if args.delay > 0 and index < total:
            time.sleep(args.delay)

    print(
        f"Finished: downloaded={counts[ArchiveDownloadState.DOWNLOADED]}, "
        f"skipped={counts[ArchiveDownloadState.SKIPPED]}, "
        f"pending={counts[ArchiveDownloadState.PENDING]}, "
        f"failed={counts[ArchiveDownloadState.FAILED]}, output_dir={args.output_dir}"
    )


def parse_args(
    argv: list[str] | None = None,
    description: str | None = None,
) -> argparse.Namespace:
    """Parse command-line arguments."""
    parser = argparse.ArgumentParser(description=description or __doc__)
    parser.add_argument("--response-dir", type=Path)
    parser.add_argument("--output-dir", type=Path)
    parser.add_argument("--start", type=int, default=0, help="Zero-based response file offset.")
    parser.add_argument("--limit", type=int, help="Number of response files to download.")
    parser.add_argument("--timeout", type=int, default=DEFAULT_TIMEOUT, help="Download timeout in seconds.")
    parser.add_argument("--delay", type=float, default=0.0, help="Seconds to wait between downloads.")
    parser.add_argument("--overwrite", action="store_true", help="Download even when the ZIP already exists.")
    args = parser.parse_args(argv)

    missing = [
        option
        for option, value in {
            "--response-dir": args.response_dir,
            "--output-dir": args.output_dir,
        }.items()
        if not value
    ]
    if missing:
        parser.error(f"the following arguments are required: {', '.join(missing)}")
    if args.timeout <= 0:
        parser.error("--timeout must be greater than 0.")
    if args.delay < 0:
        parser.error("--delay must be 0 or greater.")

    return args


def main(
    argv: list[str] | None = None,
    description: str | None = None,
) -> None:
    """Run the archive downloader."""
    run(parse_args(argv, description))


if __name__ == "__main__":
    main()
