#!/usr/bin/env python3
"""Utility to download and install the DB-IP datasets used by the IOC checker."""

from __future__ import annotations

import argparse
import datetime as dt
import gzip
import logging
import shutil
import sys
import tempfile
from pathlib import Path
from typing import Iterable

import httpx

from ioc_checker.http_client import build_sync_http_client

LOGGER = logging.getLogger("update-dbip")


def configure_logging(verbose: bool) -> None:
    """Initialise logging for the updater script."""

    level = logging.DEBUG if verbose else logging.INFO
    logging.basicConfig(level=level, format="[%(name)s] %(message)s", force=True)


def default_year_month() -> str:
    """Return the latest available DB-IP dataset identifier (UTC year-month)."""

    today = dt.datetime.now(dt.timezone.utc).date().replace(day=1)
    return today.strftime("%Y-%m")


def parse_args(argv: Iterable[str]) -> argparse.Namespace:
    """Parse command line arguments."""

    parser = argparse.ArgumentParser(
        description="Download the DB-IP city and ASN Lite datasets and install them into the repository.",
    )
    parser.add_argument(
        "year_month",
        nargs="?",
        default=default_year_month(),
        help="Dataset release identifier in YYYY-MM format (defaults to the current month).",
    )
    parser.add_argument(
        "--verbose",
        action="store_true",
        help="Enable debug logging for troubleshooting.",
    )
    return parser.parse_args(list(argv))


def validate_year_month(value: str) -> None:
    """Ensure the provided year-month matches the expected format."""

    try:
        dt.datetime.strptime(value, "%Y-%m")
    except ValueError as exc:  # pragma: no cover - defensive programming
        raise SystemExit(f"Invalid year-month value '{value}'; expected YYYY-MM format.") from exc


def repo_paths() -> tuple[Path, Path]:
    """Return the repository root and target data directory."""

    script_path = Path(__file__).resolve()
    repo_root = script_path.parent.parent
    target_dir = repo_root / "ioc_checker" / "dbipcom"
    return repo_root, target_dir


def download_dataset(client: httpx.Client, url: str, label: str) -> Path:
    """Fetch a gzip-compressed dataset and return the temporary file path."""

    LOGGER.info("Downloading %s dataset from %s", label, url)
    try:
        with client.stream("GET", url) as response:
            response.raise_for_status()
            with tempfile.NamedTemporaryFile(delete=False) as handle:
                for chunk in response.iter_bytes():
                    handle.write(chunk)
                temp_path = Path(handle.name)
    except httpx.HTTPError as exc:
        raise SystemExit(f"Failed to download {label} dataset: {exc}") from exc

    return temp_path


def extract_dataset(source: Path, destination: Path, label: str) -> None:
    """Decompress a gzip archive to the destination path."""

    destination.parent.mkdir(parents=True, exist_ok=True)
    LOGGER.info("Decompressing %s dataset to %s", label, destination)
    try:
        with gzip.open(source, "rb") as compressed, destination.open("wb") as target:
            shutil.copyfileobj(compressed, target)
    except OSError as exc:
        raise SystemExit(f"Failed to decompress {label} dataset: {exc}") from exc


def update_dbip(year_month: str) -> None:
    """Download and install the DB-IP city and ASN datasets for the specified month."""

    validate_year_month(year_month)
    _, target_dir = repo_paths()
    datasets = {
        "DB-IP city": (
            f"https://download.db-ip.com/free/dbip-city-lite-{year_month}.mmdb.gz",
            target_dir / "dbip-city.mmdb",
        ),
        "DB-IP ASN": (
            f"https://download.db-ip.com/free/dbip-asn-lite-{year_month}.mmdb.gz",
            target_dir / "dbip-asn.mmdb",
        ),
    }

    with build_sync_http_client(timeout=60.0) as client:
        for label, (url, destination) in datasets.items():
            temp_file = download_dataset(client, url, label)
            try:
                extract_dataset(temp_file, destination, label)
            finally:
                try:
                    temp_file.unlink()
                except FileNotFoundError:
                    pass

    LOGGER.info("DB-IP datasets updated successfully")


def main(argv: Iterable[str]) -> int:
    """Run the updater script."""

    args = parse_args(argv)
    configure_logging(args.verbose)
    update_dbip(args.year_month)
    return 0


if __name__ == "__main__":  # pragma: no cover - manual execution entry point
    sys.exit(main(sys.argv[1:]))
