#!/usr/bin/env python3
"""Download configured geodata tiles into raw/ based on a TOML config."""
from __future__ import annotations

import argparse
import os
import shutil
import sys
import time
import queue
import threading
import zipfile
from concurrent.futures import Future, as_completed
from dataclasses import dataclass
from typing import Dict, Iterable, List, Optional, Tuple
from urllib.parse import urlparse

try:
    import tomllib
except ImportError:  # pragma: no cover - tomllib is required
    raise SystemExit("tomllib is required (Python 3.11+).")

import requests

DEFAULT_CONFIG = "geodata_download.toml"
DEFAULT_OUTPUT_DIR = "raw"

OUTPUT_SUBDIRS = {
    "dgm1": "dgm1",
    "dom1": "dom1",
    "dop20": "dop20",
    "geb3dlo": os.path.join("citygml", "lod2"),
    "citygml": os.path.join("citygml", "lod2"),
    "bdom20rgbi": "bdom20rgbi",
    "lpg": "lpolpg",
    "lpo": "lpolpg",
    "lpolpg": "lpolpg",
}

FILE_TYPE_SUBDIRS = {
    "image": "jp2",
    "worldfile": "j2w",
    "metadata": "meta",
}


@dataclass(frozen=True)
class DownloadTask:
    dataset: str
    url: str
    output_path: str
    download_key: str
    unzip: bool = False
    unzip_dir: Optional[str] = None


class DownloadLogger:
    def __init__(
        self,
        log_file: Optional[str] = None,
        log_format: Optional[str] = None,
        report_progress: bool = True,
    ) -> None:
        self._log_file = log_file
        self._log_format = log_format
        self._report_progress = report_progress

    def log(self, level: str, message: str) -> None:
        line = self._format(level, message)
        if self._report_progress:
            print(line)
        if self._log_file:
            with open(self._log_file, "a", encoding="utf-8") as fh:
                fh.write(line + "\n")

    def _format(self, level: str, message: str) -> str:
        if not self._log_format:
            return f"[{level}] {message}"
        timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
        try:
            return self._log_format.format(timestamp=timestamp, level=level, message=message)
        except Exception:
            return f"[{level}] {message}"


class DownloadProgress:
    def __init__(self, total: int, enabled: bool) -> None:
        self._total = max(total, 1)
        self._enabled = enabled
        self._downloaded = 0
        self._missing = 0
        self._failed = 0
        self._bytes = 0
        self._lock = threading.Lock()
        self._start = time.time()
        self._last_label = ""
        self._last_render = 0.0
        self._min_interval = 0.25

    def add_bytes(self, bytes_delta: int, label: Optional[str] = None) -> None:
        if not self._enabled or bytes_delta <= 0:
            return
        with self._lock:
            self._bytes += bytes_delta
            if label:
                self._last_label = os.path.basename(label)
            self._render_locked(force=False)

    def set_counts(self, downloaded: int, missing: int, failed: int, label: Optional[str] = None) -> None:
        if not self._enabled:
            return
        with self._lock:
            self._downloaded = downloaded
            self._missing = missing
            self._failed = failed
            if label:
                self._last_label = os.path.basename(label)
            self._render_locked(force=True)

    def finish(self) -> None:
        if not self._enabled:
            return
        sys.stderr.write("\n")
        sys.stderr.flush()

    def _render_locked(self, force: bool) -> None:
        now = time.time()
        if not force and (now - self._last_render) < self._min_interval:
            return
        self._last_render = now
        elapsed = max(now - self._start, 0.001)
        done = self._downloaded + self._missing + self._failed
        rate = done / elapsed
        remaining = max(self._total - done, 0)
        eta = int(remaining / rate) if rate > 0 else 0
        bytes_mb = self._bytes / (1024 * 1024)
        bytes_gb = bytes_mb / 1024
        byte_rate = bytes_mb / elapsed
        width = 28
        filled = int(width * done / self._total)
        bar = "#" * filled + "-" * (width - filled)
        line = (
            f"\r[{bar}] {done}/{self._total} "
            f"{rate:.2f}/s eta {eta}s ok={self._downloaded} miss={self._missing} fail={self._failed} "
            f"{bytes_gb:.1f}GB {byte_rate:.1f}MB/s "
            f"{self._last_label}"
        )
        sys.stderr.write(line[:200])
        sys.stderr.flush()


class DaemonThreadPool:
    """Minimal daemon-thread pool that supports submit() + shutdown()."""

    def __init__(self, max_workers: int) -> None:
        if max_workers <= 0:
            raise ValueError("max_workers must be greater than 0")
        self._tasks: queue.Queue = queue.Queue()
        self._threads: list[threading.Thread] = []
        self._shutdown = False
        for idx in range(max_workers):
            thread = threading.Thread(
                name=f"download-worker-{idx}",
                target=self._worker,
                daemon=True,
            )
            thread.start()
            self._threads.append(thread)

    def submit(self, fn, *args, **kwargs) -> Future:
        if self._shutdown:
            raise RuntimeError("Thread pool already shutdown")
        future: Future = Future()
        self._tasks.put((future, fn, args, kwargs))
        return future

    def shutdown(self, wait: bool = True, cancel_futures: bool = False) -> None:
        self._shutdown = True
        if cancel_futures:
            while True:
                try:
                    item = self._tasks.get_nowait()
                except queue.Empty:
                    break
                if item is None:
                    continue
                future, _, _, _ = item
                future.cancel()
        for _ in self._threads:
            self._tasks.put(None)
        if wait:
            for thread in self._threads:
                thread.join()

    def _worker(self) -> None:
        while True:
            item = self._tasks.get()
            if item is None:
                return
            future, fn, args, kwargs = item
            if not future.set_running_or_notify_cancel():
                continue
            try:
                result = fn(*args, **kwargs)
            except BaseException as exc:
                future.set_exception(exc)
            else:
                future.set_result(result)


def _load_toml(path: str) -> dict:
    if not os.path.exists(path):
        raise SystemExit(f"Config not found: {path}")
    with open(path, "rb") as fh:
        return tomllib.load(fh)


def _parse_download_targets(cfg: dict) -> Tuple[Dict[str, dict], Optional[str]]:
    download_cfg = cfg.get("download", {})
    if not isinstance(download_cfg, dict):
        raise SystemExit("download must be a table in the TOML config.")
    has_targets = any(isinstance(value, dict) for value in download_cfg.values())
    if has_targets:
        targets = {name: value for name, value in download_cfg.items() if isinstance(value, dict)}
        if not targets:
            raise SystemExit("download targets must be tables, e.g. [download.geobasis].")
        default_key = download_cfg.get("default_target") or download_cfg.get("default")
        return targets, default_key
    return {"default": download_cfg}, "default"


def _parse_tile_ranges(cfg: dict) -> Dict[str, dict]:
    ranges = cfg.get("tile_ranges", {})
    if not isinstance(ranges, dict):
        raise SystemExit("tile_ranges must be a table in the TOML config.")
    return ranges


def _iter_tiles(range_cfg: dict) -> Iterable[Tuple[int, int]]:
    x_start = int(range_cfg["x_start"])
    x_end = int(range_cfg["x_end"])
    x_step = int(range_cfg.get("x_step", 1))
    y_start = int(range_cfg["y_start"])
    y_end = int(range_cfg["y_end"])
    y_step = int(range_cfg.get("y_step", 1))

    for x in range(x_start, x_end + 1, x_step):
        for y in range(y_start, y_end + 1, y_step):
            yield x, y


def _override_range(range_cfg: dict, start: Tuple[int, int], end: Tuple[int, int]) -> dict:
    new_range = dict(range_cfg)
    new_range["x_start"] = start[0]
    new_range["y_start"] = start[1]
    new_range["x_end"] = end[0]
    new_range["y_end"] = end[1]
    return new_range


def _resolve_output_dir(dataset_key: str, dataset_cfg: dict) -> str:
    if dataset_key in OUTPUT_SUBDIRS:
        return OUTPUT_SUBDIRS[dataset_key]
    return dataset_cfg.get("output_subdir", dataset_key)


def _format_url(
    template: str,
    base_url: str,
    x: Optional[int],
    y: Optional[int],
    extra: dict,
) -> str:
    format_vars = {"base_url": base_url}
    if x is not None:
        format_vars["x"] = x
    if y is not None:
        format_vars["y"] = y
    for key, value in extra.items():
        if isinstance(value, (str, int, float)):
            format_vars[key] = value
    return template.format(**format_vars)


def _select_download_target(
    download_targets: Dict[str, dict],
    default_key: Optional[str],
    dataset_key: str,
    dataset_cfg: dict,
) -> Tuple[str, dict]:
    requested = dataset_cfg.get("download")
    if requested:
        if requested not in download_targets:
            raise SystemExit(f"{dataset_key}: download target not found: {requested}")
        return requested, download_targets[requested]
    if default_key:
        if default_key not in download_targets:
            raise SystemExit(f"download.default_target not found: {default_key}")
        return default_key, download_targets[default_key]
    if len(download_targets) == 1:
        name = next(iter(download_targets))
        return name, download_targets[name]
    raise SystemExit(
        f"{dataset_key}: multiple download targets defined; set datasets.{dataset_key}.download "
        "or download.default_target.",
    )


def _resolve_url(source_cfg: dict, base_url: str, x: Optional[int], y: Optional[int]) -> str:
    if "url" in source_cfg:
        return str(source_cfg["url"])
    template = source_cfg.get("url_template")
    if not template:
        raise SystemExit("url_template or url must be defined for the dataset.")
    return _format_url(str(template), base_url, x, y, source_cfg)


def _build_tasks(
    cfg: dict,
    datasets: Dict[str, dict],
    tile_ranges: Dict[str, dict],
    download_targets: Dict[str, dict],
    default_target: Optional[str],
    start_override: Optional[Tuple[int, int]],
    end_override: Optional[Tuple[int, int]],
) -> List[DownloadTask]:
    tasks: List[DownloadTask] = []

    for dataset_key, dataset_cfg in datasets.items():
        download_key, download_cfg = _select_download_target(
            download_targets,
            default_target,
            dataset_key,
            dataset_cfg,
        )
        base_url = download_cfg.get("base_url", "").rstrip("/")
        base_output_dir = download_cfg.get("output_directory", DEFAULT_OUTPUT_DIR)
        base_subdir = _resolve_output_dir(dataset_key, dataset_cfg)
        dataset_out_dir = os.path.join(base_output_dir, base_subdir)
        unzip = bool(dataset_cfg.get("unzip", False))
        unzip_dir = dataset_out_dir

        tile_range_key = dataset_cfg.get("tile_range")
        if tile_range_key:
            if tile_range_key not in tile_ranges:
                raise SystemExit(f"{dataset_key}: tile_range not found: {tile_range_key}")
            range_cfg = tile_ranges[tile_range_key]
            if start_override and end_override:
                range_cfg = _override_range(range_cfg, start_override, end_override)

            for x, y in _iter_tiles(range_cfg):
                if "files" in dataset_cfg:
                    for file_cfg in dataset_cfg["files"]:
                        file_type = file_cfg.get("type", "file")
                        file_subdir = FILE_TYPE_SUBDIRS.get(file_type, file_type)
                        out_dir = os.path.join(dataset_out_dir, file_subdir)
                        url = _resolve_url(file_cfg, base_url, x, y)
                        filename = os.path.basename(urlparse(url).path)
                        tasks.append(
                            DownloadTask(
                                dataset_key,
                                url,
                                os.path.join(out_dir, filename),
                                download_key,
                                unzip,
                                unzip_dir,
                            )
                        )
                else:
                    url = _resolve_url(dataset_cfg, base_url, x, y)
                    filename = os.path.basename(urlparse(url).path)
                    tasks.append(
                        DownloadTask(
                            dataset_key,
                            url,
                            os.path.join(dataset_out_dir, filename),
                            download_key,
                            unzip,
                            unzip_dir,
                        )
                    )
        else:
            if "files" in dataset_cfg:
                for file_cfg in dataset_cfg["files"]:
                    file_type = file_cfg.get("type", "file")
                    file_subdir = FILE_TYPE_SUBDIRS.get(file_type, file_type)
                    out_dir = os.path.join(dataset_out_dir, file_subdir)
                    url = _resolve_url(file_cfg, base_url, None, None)
                    filename = os.path.basename(urlparse(url).path)
                    tasks.append(
                        DownloadTask(
                            dataset_key,
                            url,
                            os.path.join(out_dir, filename),
                            download_key,
                            unzip,
                            unzip_dir,
                        )
                    )
            else:
                url = _resolve_url(dataset_cfg, base_url, None, None)
                filename = os.path.basename(urlparse(url).path)
                tasks.append(
                    DownloadTask(
                        dataset_key,
                        url,
                        os.path.join(dataset_out_dir, filename),
                        download_key,
                        unzip,
                        unzip_dir,
                    )
                )

    return tasks


def _select_datasets(cfg: dict, requested: Optional[List[str]]) -> Dict[str, dict]:
    datasets = cfg.get("datasets", {})
    if not isinstance(datasets, dict) or not datasets:
        raise SystemExit("datasets must be defined in the TOML config.")

    if requested:
        missing = [name for name in requested if name not in datasets]
        if missing:
            raise SystemExit(f"Unknown dataset(s): {', '.join(missing)}")
        selected = {name: datasets[name] for name in requested}
    else:
        selected = {name: ds for name, ds in datasets.items() if ds.get("enabled", True)}

    if not selected:
        raise SystemExit("No datasets selected for download.")
    return selected


def _safe_remove_dir(base_dir: str, rel_dir: str) -> None:
    target = os.path.abspath(os.path.join(base_dir, rel_dir))
    base = os.path.abspath(base_dir)
    if os.path.commonpath([target, base]) != base:
        raise SystemExit(f"Refusing to delete outside base dir: {target}")
    if target == base:
        raise SystemExit(f"Refusing to delete base dir: {target}")
    if os.path.exists(target):
        shutil.rmtree(target)


def _unzipped_marker_path(task: DownloadTask) -> str:
    unzip_dir = task.unzip_dir or os.path.dirname(task.output_path)
    marker_name = f".{os.path.basename(task.output_path)}.unzipped"
    return os.path.join(unzip_dir, marker_name)


def _needs_unzip(task: DownloadTask) -> bool:
    if not task.unzip:
        return False
    if not os.path.exists(task.output_path):
        return False
    return not os.path.exists(_unzipped_marker_path(task))


def _maybe_unzip(task: DownloadTask, logger: DownloadLogger) -> bool:
    if not task.unzip:
        return True
    if not os.path.exists(task.output_path):
        logger.log("ERROR", f"Unzip failed; missing file: {task.output_path}")
        return False
    unzip_dir = task.unzip_dir or os.path.dirname(task.output_path)
    os.makedirs(unzip_dir, exist_ok=True)
    try:
        with zipfile.ZipFile(task.output_path, "r") as zf:
            zf.extractall(unzip_dir)
    except (zipfile.BadZipFile, OSError) as exc:
        logger.log("ERROR", f"Unzip failed: {task.output_path} ({exc})")
        return False
    marker = _unzipped_marker_path(task)
    try:
        with open(marker, "w", encoding="utf-8") as fh:
            fh.write(time.strftime("%Y-%m-%d %H:%M:%S"))
    except OSError as exc:
        logger.log("WARN", f"Could not write unzip marker: {marker} ({exc})")
    logger.log("INFO", f"Unzipped: {task.output_path} -> {unzip_dir}")
    return True


def _download_task(
    session: requests.Session,
    task: DownloadTask,
    timeout: int,
    verify: bool | str,
    retries: int,
    stop_event: threading.Event,
    progress: DownloadProgress,
) -> Tuple[str, DownloadTask, Optional[str]]:
    os.makedirs(os.path.dirname(task.output_path), exist_ok=True)
    tmp_path = f"{task.output_path}.part"

    if stop_event.is_set():
        return "aborted", task, "Interrupted"

    for attempt in range(retries + 1):
        if stop_event.is_set():
            return "aborted", task, "Interrupted"
        try:
            with session.get(task.url, stream=True, timeout=timeout, verify=verify) as resp:
                if resp.status_code in (404, 410):
                    return "missing", task, f"HTTP {resp.status_code}"
                if resp.status_code >= 400:
                    return "failed", task, f"HTTP {resp.status_code}"
                resp.raise_for_status()
                with open(tmp_path, "wb") as fh:
                    for chunk in resp.iter_content(chunk_size=1024 * 1024):
                        if stop_event.is_set():
                            return "aborted", task, "Interrupted"
                        if chunk:
                            fh.write(chunk)
                            progress.add_bytes(len(chunk), task.output_path)
            os.replace(tmp_path, task.output_path)
            return "downloaded", task, None
        except requests.RequestException as exc:
            if attempt >= retries:
                return "failed", task, str(exc)
            time.sleep(1.0 + attempt * 0.5)
        except OSError as exc:
            return "failed", task, str(exc)

    if os.path.exists(tmp_path):
        try:
            os.remove(tmp_path)
        except OSError:
            pass
    return "failed", task, "Unknown error"


def run_download(
    config_path: str,
    requested_datasets: Optional[List[str]] = None,
    start_override: Optional[Tuple[int, int]] = None,
    end_override: Optional[Tuple[int, int]] = None,
    clean_downloads: bool = False,
    ca_bundle_override: Optional[str] = None,
) -> int:
    cfg = _load_toml(config_path)
    download_targets, default_target = _parse_download_targets(cfg)
    tile_ranges = _parse_tile_ranges(cfg)
    datasets = _select_datasets(cfg, requested_datasets)

    logging_cfg = cfg.get("logging", {})
    progress_enabled = bool(logging_cfg.get("report_progress", True))
    logger = DownloadLogger(
        logging_cfg.get("log_file"),
        logging_cfg.get("log_format"),
        progress_enabled,
    )

    if len(download_targets) > 1 and not default_target:
        default_target = next(iter(download_targets))
        logger.log(
            "WARN",
            "No download.default_target set; using "
            f"{default_target} as default for datasets without a download key.",
        )

    if clean_downloads:
        for dataset_key, dataset_cfg in datasets.items():
            download_key, download_cfg = _select_download_target(
                download_targets,
                default_target,
                dataset_key,
                dataset_cfg,
            )
            base_output_dir = download_cfg.get("output_directory", DEFAULT_OUTPUT_DIR)
            _safe_remove_dir(base_output_dir, _resolve_output_dir(dataset_key, dataset_cfg))

    tasks = _build_tasks(
        cfg,
        datasets,
        tile_ranges,
        download_targets,
        default_target,
        start_override,
        end_override,
    )

    if not tasks:
        logger.log("INFO", "No download tasks generated.")
        return 0

    tasks_by_target: Dict[str, List[DownloadTask]] = {}
    for task in tasks:
        tasks_by_target.setdefault(task.download_key, []).append(task)

    total_downloaded = 0
    total_missing = 0
    total_failed = 0
    total_skipped = 0
    total_unzipped_existing = 0

    for target_key, target_cfg in download_targets.items():
        target_tasks = tasks_by_target.get(target_key)
        if not target_tasks:
            continue

        target_label = f"[{target_key}] "
        skip_existing = not clean_downloads
        pending: List[DownloadTask] = []
        unzip_only: List[DownloadTask] = []
        skipped = 0

        for task in target_tasks:
            if skip_existing and os.path.exists(task.output_path):
                if _needs_unzip(task):
                    unzip_only.append(task)
                else:
                    skipped += 1
                continue
            pending.append(task)

        if skipped:
            logger.log("INFO", f"{target_label}Skipped {skipped} existing file(s).")
        if unzip_only:
            logger.log("INFO", f"{target_label}Queued {len(unzip_only)} unzip(s) for existing files.")

        verify_ssl = target_cfg.get("verify_ssl", True)
        ca_bundle = ca_bundle_override or target_cfg.get("ca_bundle")
        if verify_ssl and ca_bundle:
            if os.path.exists(ca_bundle):
                verify = ca_bundle
                source = "CLI" if ca_bundle_override else "config"
                logger.log("INFO", f"{target_label}Using CA bundle ({source}): {ca_bundle}")
            else:
                verify = True
                logger.log("WARN", f"{target_label}CA bundle not found, using system trust: {ca_bundle}")
        else:
            verify = bool(verify_ssl)
            if not verify_ssl:
                logger.log("WARN", f"{target_label}TLS verification disabled by config.")

        timeout = int(target_cfg.get("timeout_seconds", 300))
        retries = int(target_cfg.get("retry_attempts", 3))
        parallel = int(target_cfg.get("parallel_downloads", 4))
        user_agent = target_cfg.get("user_agent", "geodata-download/1.0")

        downloaded = 0
        missing = 0
        failed = 0

        if pending:
            progress = DownloadProgress(len(pending), progress_enabled)
            stop_event = threading.Event()
            interrupted = False

            with requests.Session() as session:
                session.headers.update({"User-Agent": user_agent})
                executor = DaemonThreadPool(max_workers=parallel)
                futures = [
                    executor.submit(_download_task, session, task, timeout, verify, retries, stop_event, progress)
                    for task in pending
                ]
                try:
                    for future in as_completed(futures):
                        status, task, detail = future.result()
                        if status == "downloaded":
                            downloaded += 1
                            if task.unzip and not _maybe_unzip(task, logger):
                                failed += 1
                        elif status == "missing":
                            missing += 1
                            logger.log("WARN", f"{target_label}Missing tile: {task.url} ({detail})")
                        elif status == "aborted":
                            failed += 1
                        else:
                            failed += 1
                            extra = f" ({detail})" if detail else ""
                            logger.log("ERROR", f"{target_label}Failed: {task.url}{extra}")
                        progress.set_counts(downloaded, missing, failed, task.output_path)
                except KeyboardInterrupt:
                    interrupted = True
                    stop_event.set()
                    logger.log("WARN", f"{target_label}Interrupted; stopping downloads.")
                finally:
                    if interrupted:
                        executor.shutdown(wait=False, cancel_futures=True)
                    else:
                        executor.shutdown(wait=True)

            progress.finish()
            if interrupted:
                return 130
        else:
            logger.log("INFO", f"{target_label}Nothing to download.")

        if unzip_only:
            for task in unzip_only:
                if _maybe_unzip(task, logger):
                    total_unzipped_existing += 1
                else:
                    failed += 1

        total_downloaded += downloaded
        total_missing += missing
        total_failed += failed
        total_skipped += skipped

    logger.log(
        "INFO",
        "Done. "
        f"Downloaded={total_downloaded}, Missing={total_missing}, Failed={total_failed}, "
        f"Skipped={total_skipped}, UnzippedExisting={total_unzipped_existing}.",
    )
    return 1 if total_failed else 0


def parse_args(argv: Optional[Iterable[str]] = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Download geodata tiles from TOML config.")
    parser.add_argument(
        "--config",
        default=DEFAULT_CONFIG,
        help="Path to geodata_download.toml.",
    )
    parser.add_argument(
        "--datasets",
        help="Comma-separated dataset keys to download (default: enabled datasets).",
    )
    parser.add_argument(
        "--start",
        nargs=2,
        type=int,
        metavar=("X", "Y"),
        help="Override tile range start (x y).",
    )
    parser.add_argument(
        "--end",
        nargs=2,
        type=int,
        metavar=("X", "Y"),
        help="Override tile range end (x y).",
    )
    parser.add_argument(
        "--clean-downloads",
        action="store_true",
        help="Delete selected dataset folders before downloading.",
    )
    parser.add_argument(
        "--ca-bundle",
        help="Path to a CA bundle file to override the config.",
    )
    return parser.parse_args(argv)


def main(argv: Optional[Iterable[str]] = None) -> int:
    args = parse_args(argv)
    datasets = [name.strip() for name in args.datasets.split(",")] if args.datasets else None
    start = tuple(args.start) if args.start else None
    end = tuple(args.end) if args.end else None
    if (start is None) != (end is None):
        raise SystemExit("--start and --end must be provided together.")

    return run_download(
        config_path=args.config,
        requested_datasets=datasets,
        start_override=start,
        end_override=end,
        clean_downloads=args.clean_downloads,
        ca_bundle_override=args.ca_bundle,
    )


if __name__ == "__main__":
    sys.exit(main())
