"""Download GDELT Web NGrams 3.0 JSON.GZ files for a time range.

Example:
    from gdeltnews_rs import download

    download(
        "2025-03-15T00:00:00",
        "2025-03-15T00:10:00",
        outdir="gdeltdata",
    )
"""

from __future__ import annotations

import datetime as dt
import os
from dataclasses import dataclass
from typing import Iterable, Optional, Union

import requests
from tqdm import tqdm

GDELT_BASE_URL = "http://data.gdeltproject.org/gdeltv3/webngrams"

TimestampLike = Union[str, dt.datetime]


@dataclass(frozen=True)
class DownloadStats:
    """Download summary."""
    requested_minutes: int
    downloaded_gz: int


def parse_timestamp(ts: str) -> dt.datetime:
    """Parse a timestamp string into a naive UTC datetime."""
    ts = ts.strip()
    if len(ts) == 14 and ts.isdigit():
        return dt.datetime.strptime(ts, "%Y%m%d%H%M%S")
    if ts.endswith("Z"):
        ts = ts[:-1]
    for fmt in ("%Y-%m-%dT%H:%M:%S", "%Y-%m-%d %H:%M:%S"):
        try:
            return dt.datetime.strptime(ts, fmt)
        except ValueError:
            continue
    raise ValueError(f"Unrecognized timestamp format: {ts}")


def _coerce_timestamp(value: TimestampLike) -> dt.datetime:
    if isinstance(value, dt.datetime):
        return value
    return parse_timestamp(str(value))


def iter_minutes(start: dt.datetime, end: dt.datetime) -> Iterable[dt.datetime]:
    """Yield every minute from start to end inclusive."""
    if end < start:
        raise ValueError("End time must be >= start time")
    current = start
    step = dt.timedelta(minutes=1)
    while current <= end:
        yield current
        current += step


def _download_one(
    ts: dt.datetime,
    dest_dir: str,
    overwrite: bool = False,
    timeout: int = 30,
) -> Optional[str]:
    """Download a single GDELT Web NGrams file. Returns path or None."""
    fname = ts.strftime("%Y%m%d%H%M%S") + ".webngrams.json.gz"
    url = f"{GDELT_BASE_URL}/{fname}"
    gz_path = os.path.join(dest_dir, fname)

    if not overwrite and os.path.exists(gz_path):
        return gz_path

    try:
        resp = requests.get(url, stream=True, timeout=timeout)
    except requests.RequestException:
        return None

    if resp.status_code != 200:
        return None

    with open(gz_path, "wb") as f:
        for chunk in resp.iter_content(chunk_size=1 << 20):
            if chunk:
                f.write(chunk)
    return gz_path


def download(
    start: TimestampLike,
    end: TimestampLike,
    *,
    outdir: str = "gdeltdata",
    overwrite: bool = False,
    timeout: int = 30,
    show_progress: bool = True,
) -> DownloadStats:
    """Download GDELT Web NGrams .gz files for the given time range.

    Note: Unlike the original gdeltnews package, this does NOT decompress
    files. The Rust reconstruction engine reads .gz files directly.

    Args:
        start: Start timestamp (datetime or string).
        end: End timestamp (datetime or string).
        outdir: Destination directory.
        overwrite: Re-download even if file exists.
        timeout: HTTP request timeout in seconds.
        show_progress: Show tqdm progress bar.

    Returns:
        DownloadStats with counts.
    """
    start_dt = _coerce_timestamp(start)
    end_dt = _coerce_timestamp(end)
    minutes = list(iter_minutes(start_dt, end_dt))
    total = len(minutes)

    print(f"Downloading {total} files from {start_dt} to {end_dt}")
    os.makedirs(outdir, exist_ok=True)

    downloaded = 0
    iterator = tqdm(minutes, desc="Downloading", unit="file") if show_progress else minutes

    for ts in iterator:
        path = _download_one(ts, outdir, overwrite=overwrite, timeout=timeout)
        if path is not None:
            downloaded += 1

    print(f"Downloaded {downloaded}/{total} files to {outdir}")
    return DownloadStats(requested_minutes=total, downloaded_gz=downloaded)


__all__ = ["DownloadStats", "download", "parse_timestamp"]
