"""Bootstrap canopy height rasters for an AOI."""

from __future__ import annotations

import argparse
import json
import os
import socket
import time
from collections.abc import Iterable
from dataclasses import dataclass
from pathlib import Path
from typing import Any
from urllib.error import HTTPError, URLError
from urllib.parse import unquote, urlparse
from urllib.request import Request, urlopen

META_CHM_V2_INDEX_BASE_URL = (
    "https://dataforgood-fb-data.s3.amazonaws.com/"
    "forests/v2/global/dinov3_global_chm_v2_ml3"
)
META_CHM_V2_BASE_URL = "https://data.source.coop/tge-labs/meta-chm-v2/chm"
META_CHM_V2_INDEX_URL = f"{META_CHM_V2_INDEX_BASE_URL}/tiles.geojson"
USER_AGENT = "RadioApi canopy bootstrap/0.1"


@dataclass(frozen=True)
class Socks5Proxy:
    scheme: str
    host: str
    port: int
    username: str | None
    password: str | None

    @property
    def endpoint(self) -> str:
        return f"{self.scheme}://{self.host}:{self.port}"


def parse_socks5_proxy(proxy_url: str) -> Socks5Proxy:
    parsed = urlparse(proxy_url)
    if parsed.scheme not in {"socks5", "socks5h"}:
        raise SystemExit(
            f"Unsupported proxy scheme: {parsed.scheme!r}. Use socks5:// or socks5h://"
        )
    if not parsed.hostname:
        raise SystemExit("Proxy URL must include hostname, e.g. socks5://user:pass@host:port")
    return Socks5Proxy(
        scheme=parsed.scheme,
        host=parsed.hostname,
        port=parsed.port or 1080,
        username=unquote(parsed.username) if parsed.username else None,
        password=unquote(parsed.password) if parsed.password else None,
    )


def configure_socks5_proxy(proxy_url: str) -> Socks5Proxy:
    try:
        import socks
    except ImportError as exc:
        raise SystemExit(
            "SOCKS5 proxy requires PySocks. Install with: pip install PySocks"
        ) from exc

    proxy = parse_socks5_proxy(proxy_url)
    # socks5h resolves hostnames on the proxy (like curl --proxy socks5h://).
    rdns = proxy.scheme == "socks5h"
    socks.set_default_proxy(
        socks.SOCKS5,
        proxy.host,
        proxy.port,
        rdns=rdns,
        username=proxy.username,
        password=proxy.password,
    )
    socket.socket = socks.socksocket
    return proxy


def parse_bbox(raw: str) -> tuple[float, float, float, float]:
    parts = [float(part.strip()) for part in raw.split(",")]
    if len(parts) != 4:
        raise argparse.ArgumentTypeError("bbox must be minlon,minlat,maxlon,maxlat")
    minlon, minlat, maxlon, maxlat = parts
    if minlon >= maxlon or minlat >= maxlat:
        raise argparse.ArgumentTypeError("bbox min values must be lower than max values")
    return minlon, minlat, maxlon, maxlat


def intersects(
    a: tuple[float, float, float, float],
    b: tuple[float, float, float, float],
) -> bool:
    return not (a[2] < b[0] or a[0] > b[2] or a[3] < b[1] or a[1] > b[3])


def _walk_coordinates(value: Any) -> Iterable[tuple[float, float]]:
    if (
        isinstance(value, list)
        and len(value) >= 2
        and isinstance(value[0], int | float)
        and isinstance(value[1], int | float)
    ):
        yield float(value[0]), float(value[1])
        return
    if isinstance(value, list):
        for item in value:
            yield from _walk_coordinates(item)


def feature_bbox(feature: dict[str, Any]) -> tuple[float, float, float, float] | None:
    bbox = feature.get("bbox")
    if isinstance(bbox, list) and len(bbox) >= 4:
        return float(bbox[0]), float(bbox[1]), float(bbox[2]), float(bbox[3])

    geometry = feature.get("geometry") or {}
    coordinates = list(_walk_coordinates(geometry.get("coordinates")))
    if not coordinates:
        return None
    lons = [coord[0] for coord in coordinates]
    lats = [coord[1] for coord in coordinates]
    return min(lons), min(lats), max(lons), max(lats)


def _asset_url(feature: dict[str, Any], base_url: str = META_CHM_V2_BASE_URL) -> str | None:
    properties = feature.get("properties") or {}
    for key in ("url", "href", "asset_href", "s3_url", "cog_url"):
        value = properties.get(key) or feature.get(key)
        if isinstance(value, str) and value:
            if value.startswith("s3://dataforgood-fb-data/"):
                return value.replace(
                    "s3://dataforgood-fb-data",
                    "https://dataforgood-fb-data.s3.amazonaws.com",
                    1,
                )
            return value

    for key in ("filename", "file", "name", "tile", "quadkey"):
        value = properties.get(key) or feature.get(key)
        if isinstance(value, str) and value:
            filename = value if value.lower().endswith((".tif", ".tiff")) else f"{value}.tif"
            return f"{base_url.rstrip('/')}/{filename}"
    return None


def validate_urls(urls: list[str], sample_size: int = 2, *, timeout_s: int = 60) -> None:
    print(f"validating {min(sample_size, len(urls))} sample URL(s)")
    for url in urls[:sample_size]:
        request = Request(url, headers={"Range": "bytes=0-0", "User-Agent": USER_AGENT})
        try:
            with urlopen(request, timeout=timeout_s) as response:
                if response.status >= 400:
                    raise SystemExit(f"Canopy URL is not reachable: {url} ({response.status})")
        except HTTPError as exc:
            raise SystemExit(
                f"Canopy URL is not reachable: {url} ({exc.code}). "
                "Check --index-url/--urls-file or the dataset URL template."
            ) from exc
        except URLError as exc:
            raise SystemExit(f"Canopy URL validation failed: {url}: {exc}") from exc


def load_index(index_url: str, *, timeout_s: int = 120) -> dict[str, Any]:
    path = Path(index_url)
    if path.exists():
        print(f"loading index from {path}")
        return json.loads(path.read_text(encoding="utf-8"))
    print(f"loading index from {index_url}")
    request = Request(index_url, headers={"User-Agent": USER_AGENT})
    with urlopen(request, timeout=timeout_s) as response:
        payload = response.read()
    print(f"index loaded ({_format_mb(len(payload))})")
    return json.loads(payload.decode("utf-8"))


def select_tile_urls(
    index: dict[str, Any],
    bbox: tuple[float, float, float, float],
    *,
    base_url: str = META_CHM_V2_BASE_URL,
) -> list[str]:
    urls: list[str] = []
    for feature in index.get("features", []):
        tile_bbox = feature_bbox(feature)
        if tile_bbox is None or not intersects(tile_bbox, bbox):
            continue
        url = _asset_url(feature, base_url=base_url)
        if url is not None:
            urls.append(url)
    return sorted(set(urls))


def urls_from_file(path: Path) -> list[str]:
    return [
        line.strip()
        for line in path.read_text(encoding="utf-8").splitlines()
        if line.strip() and not line.lstrip().startswith("#")
    ]


def destination_for(url: str, output_dir: Path) -> Path:
    name = url.rstrip("/").rsplit("/", 1)[-1]
    if not name.lower().endswith((".tif", ".tiff")):
        name = f"{name}.tif"
    return output_dir / name


def _format_mb(value: int) -> str:
    return f"{value / (1024 * 1024):.1f} MB"


def _response_total_size(response: Any, offset: int) -> int | None:
    content_range = response.headers.get("Content-Range")
    if content_range and "/" in content_range:
        total = content_range.rsplit("/", 1)[-1]
        if total.isdigit():
            return int(total)

    total_header = response.headers.get("Content-Length")
    if total_header and total_header.isdigit():
        return offset + int(total_header) if offset else int(total_header)
    return None


def _download_one(
    url: str,
    destination: Path,
    timeout_s: int,
    chunk_size: int,
    retries: int,
) -> None:
    partial = destination.with_suffix(f"{destination.suffix}.part")
    attempt = 0

    while True:
        attempt += 1
        offset = partial.stat().st_size if partial.exists() else 0
        headers = {"User-Agent": USER_AGENT}
        if offset:
            headers["Range"] = f"bytes={offset}-"
            print(f"  resume from {_format_mb(offset)}")

        request = Request(url, headers=headers)
        try:
            with urlopen(request, timeout=timeout_s) as response:
                append = offset > 0 and response.status == 206
                if offset > 0 and not append:
                    print("  server ignored Range header, restarting this tile")
                    offset = 0

                total = _response_total_size(response, offset)
                read = offset
                last_report = time.monotonic()
                mode = "ab" if append else "wb"

                with partial.open(mode) as output:
                    while True:
                        chunk = response.read(chunk_size)
                        if not chunk:
                            break
                        output.write(chunk)
                        read += len(chunk)
                        now = time.monotonic()
                        if now - last_report >= 5:
                            if total:
                                pct = (read / total) * 100
                                print(
                                    f"  {_format_mb(read)} / {_format_mb(total)} "
                                    f"({pct:.1f}%)"
                                )
                            else:
                                print(f"  {_format_mb(read)}")
                            last_report = now
        except (TimeoutError, URLError, OSError) as exc:
            if attempt > retries:
                raise
            print(f"  retry {attempt}/{retries} after {type(exc).__name__}: {exc}")
            time.sleep(min(30, attempt * 5))
            continue

        partial.replace(destination)
        return


def download_urls(
    urls: list[str],
    output_dir: Path,
    overwrite: bool,
    *,
    timeout_s: int = 60,
    chunk_size: int = 1024 * 1024,
    retries: int = 5,
) -> tuple[int, int, list[str]]:
    downloaded = 0
    skipped = 0
    failed: list[str] = []
    output_dir.mkdir(parents=True, exist_ok=True)

    for url in urls:
        destination = destination_for(url, output_dir)
        if destination.exists() and not overwrite:
            skipped += 1
            print(f"skip existing {destination}")
            continue

        partial = destination.with_suffix(f"{destination.suffix}.part")
        if partial.exists() and overwrite:
            partial.unlink()

        print(f"download {url}")
        try:
            _download_one(url, destination, timeout_s, chunk_size, retries)
        except (HTTPError, URLError, TimeoutError, OSError) as exc:
            failed.append(f"{url}: {exc}")
            continue
        downloaded += 1

    return downloaded, skipped, failed


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Bootstrap Meta/WRI canopy height COGs")
    parser.add_argument("--bbox", type=parse_bbox, help="AOI bbox: minlon,minlat,maxlon,maxlat")
    parser.add_argument("--output-dir", default=Path("data/canopy"), type=Path)
    parser.add_argument("--source", default="meta-chm-v2", choices=["meta-chm-v2"])
    parser.add_argument("--index-url", default=META_CHM_V2_INDEX_URL)
    parser.add_argument("--urls-file", type=Path)
    parser.add_argument("--skip-url-check", action="store_true")
    parser.add_argument("--timeout-s", default=60, type=int)
    parser.add_argument(
        "--retries",
        default=5,
        type=int,
        help="Retry count per tile after network timeouts; partial files are resumed",
    )
    parser.add_argument(
        "--proxy",
        default=os.environ.get("CANOPY_PROXY"),
        help="SOCKS5 proxy URL, e.g. socks5://user:pass@host:port. "
        "Default: CANOPY_PROXY env var",
    )
    parser.add_argument("--overwrite", action="store_true")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    if args.urls_file is None and args.bbox is None:
        raise SystemExit("--bbox is required unless --urls-file is provided")

    if args.proxy:
        proxy = configure_socks5_proxy(args.proxy)
        print(f"using proxy {proxy.endpoint}")

    if args.urls_file is not None:
        urls = urls_from_file(args.urls_file)
    else:
        index = load_index(args.index_url, timeout_s=max(args.timeout_s, 120))
        urls = select_tile_urls(index, args.bbox)

    if not urls:
        raise SystemExit("No canopy tiles intersect requested AOI")

    if not args.skip_url_check:
        validate_urls(urls, timeout_s=args.timeout_s)

    print(f"selected_tiles={len(urls)}")
    downloaded, skipped, failed = download_urls(
        urls,
        args.output_dir,
        args.overwrite,
        timeout_s=args.timeout_s,
        retries=args.retries,
    )
    print(f"downloaded={downloaded} skipped={skipped} failed={len(failed)}")
    if failed:
        for item in failed:
            print(f"failed {item}")
        raise SystemExit(1)


if __name__ == "__main__":
    main()