#!/usr/bin/env python3
"""
Binance Vision automatic downloader for USD-M Futures monthly kline archives.

Default job:
- Symbols: ETHUSDC, SOLUSDC
- Interval: 5m
- Market: USD-M Futures
- Downloads every available monthly ZIP
- Downloads and verifies SHA-256 checksum files
- Resumes interrupted downloads
- Extracts CSV files
- Validates the 12-column Binance kline schema
- Merges monthly CSV files into one continuous CSV per symbol/interval

Python 3.10+ is recommended.
No third-party packages are required.
"""

from __future__ import annotations

import argparse
import csv
import hashlib
import os
import shutil
import sys
import time
import urllib.error
import urllib.request
import zipfile
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass
from datetime import date, datetime, timezone
from pathlib import Path
from typing import Iterable, Iterator

BASE_URL = "https://data.binance.vision/data/futures/um/monthly/klines"
STANDARD_COLUMNS = [
    "open_time",
    "open",
    "high",
    "low",
    "close",
    "volume",
    "close_time",
    "quote_volume",
    "count",
    "taker_buy_volume",
    "taker_buy_quote_volume",
    "ignore",
]
USER_AGENT = "Mozilla/5.0 BinanceVisionAutoDownloader/1.0"


@dataclass(frozen=True)
class Job:
    symbol: str
    interval: str
    year_month: str
    url: str
    checksum_url: str
    zip_path: Path
    checksum_path: Path
    extract_dir: Path


@dataclass
class DownloadResult:
    job: Job
    status: str
    message: str
    bytes_downloaded: int = 0
    verified: bool = False
    extracted_csv: str = ""


def parse_year_month(value: str) -> tuple[int, int]:
    try:
        parsed = datetime.strptime(value, "%Y-%m")
    except ValueError as exc:
        raise argparse.ArgumentTypeError(
            f"Invalid year-month '{value}'. Expected YYYY-MM."
        ) from exc
    return parsed.year, parsed.month


def previous_month(today: date | None = None) -> tuple[int, int]:
    current = today or date.today()
    if current.month == 1:
        return current.year - 1, 12
    return current.year, current.month - 1


def month_range(start: str, end: str | None) -> Iterator[str]:
    start_year, start_month = parse_year_month(start)
    if end:
        end_year, end_month = parse_year_month(end)
    else:
        end_year, end_month = previous_month()

    if (start_year, start_month) > (end_year, end_month):
        raise ValueError("Start month must not be after end month.")

    year, month = start_year, start_month
    while (year, month) <= (end_year, end_month):
        yield f"{year:04d}-{month:02d}"
        month += 1
        if month == 13:
            month = 1
            year += 1


def request(
    url: str,
    *,
    method: str = "GET",
    headers: dict[str, str] | None = None,
    timeout: int = 60,
):
    request_headers = {"User-Agent": USER_AGENT}
    if headers:
        request_headers.update(headers)
    req = urllib.request.Request(
        url,
        method=method,
        headers=request_headers,
    )
    return urllib.request.urlopen(req, timeout=timeout)


def remote_file_exists(url: str, timeout: int) -> bool:
    try:
        with request(url, method="HEAD", timeout=timeout) as response:
            return 200 <= response.status < 400
    except urllib.error.HTTPError as exc:
        if exc.code == 404:
            return False
        if exc.code not in {403, 405}:
            return False
    except (urllib.error.URLError, TimeoutError):
        pass

    # Some CDN endpoints do not support HEAD reliably.
    try:
        with request(
            url,
            headers={"Range": "bytes=0-0"},
            timeout=timeout,
        ) as response:
            return response.status in {200, 206}
    except urllib.error.HTTPError as exc:
        return exc.code not in {404}
    except (urllib.error.URLError, TimeoutError):
        return False


def expected_sha256(checksum_path: Path) -> str:
    content = checksum_path.read_text(encoding="utf-8").strip()
    if not content:
        raise ValueError(f"Empty checksum file: {checksum_path}")
    token = content.split()[0].strip().lower()
    if len(token) != 64 or any(ch not in "0123456789abcdef" for ch in token):
        raise ValueError(f"Invalid SHA-256 checksum in {checksum_path}")
    return token


def file_sha256(path: Path, chunk_size: int = 1024 * 1024) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(chunk_size), b""):
            digest.update(chunk)
    return digest.hexdigest()


def download_text(
    url: str,
    destination: Path,
    *,
    retries: int,
    timeout: int,
) -> None:
    destination.parent.mkdir(parents=True, exist_ok=True)
    last_error: Exception | None = None
    for attempt in range(1, retries + 1):
        try:
            with request(url, timeout=timeout) as response:
                payload = response.read()
            destination.write_bytes(payload)
            return
        except Exception as exc:
            last_error = exc
            if attempt < retries:
                time.sleep(min(2**attempt, 15))
    raise RuntimeError(f"Failed to download {url}: {last_error}")


def download_with_resume(
    url: str,
    destination: Path,
    *,
    retries: int,
    timeout: int,
) -> int:
    destination.parent.mkdir(parents=True, exist_ok=True)
    partial_path = destination.with_suffix(destination.suffix + ".part")
    last_error: Exception | None = None

    for attempt in range(1, retries + 1):
        existing_size = partial_path.stat().st_size if partial_path.exists() else 0
        headers: dict[str, str] = {}
        mode = "wb"

        if existing_size > 0:
            headers["Range"] = f"bytes={existing_size}-"
            mode = "ab"

        try:
            with request(url, headers=headers, timeout=timeout) as response:
                if existing_size > 0 and response.status != 206:
                    existing_size = 0
                    mode = "wb"

                with partial_path.open(mode) as output:
                    while True:
                        chunk = response.read(1024 * 1024)
                        if not chunk:
                            break
                        output.write(chunk)

            os.replace(partial_path, destination)
            return destination.stat().st_size
        except Exception as exc:
            last_error = exc
            if attempt < retries:
                time.sleep(min(2**attempt, 15))

    raise RuntimeError(f"Failed to download {url}: {last_error}")


def extract_zip(zip_path: Path, extract_dir: Path) -> list[Path]:
    extract_dir.mkdir(parents=True, exist_ok=True)
    extracted: list[Path] = []

    with zipfile.ZipFile(zip_path, "r") as archive:
        bad_member = archive.testzip()
        if bad_member is not None:
            raise ValueError(f"Corrupted ZIP member: {bad_member}")
        archive.extractall(extract_dir)
        for name in archive.namelist():
            path = extract_dir / name
            if path.is_file():
                extracted.append(path)

    return extracted


def looks_like_header(row: list[str]) -> bool:
    if not row:
        return False
    first = row[0].strip().lower()
    return first in {"open_time", "open time"} or not first.replace(".", "").isdigit()


def validate_csv(csv_path: Path, expected_interval_minutes: int) -> dict[str, object]:
    rows = 0
    duplicates = 0
    gaps = 0
    malformed = 0
    first_open_time: int | None = None
    last_open_time: int | None = None
    previous_open_time: int | None = None
    expected_delta = expected_interval_minutes * 60 * 1000

    with csv_path.open("r", encoding="utf-8-sig", newline="") as handle:
        reader = csv.reader(handle)
        for index, row in enumerate(reader):
            if index == 0 and looks_like_header(row):
                continue
            if len(row) != 12:
                malformed += 1
                continue

            try:
                open_time = int(float(row[0]))
            except ValueError:
                malformed += 1
                continue

            if first_open_time is None:
                first_open_time = open_time

            if previous_open_time is not None:
                delta = open_time - previous_open_time
                if delta == 0:
                    duplicates += 1
                elif delta != expected_delta:
                    gaps += 1

            previous_open_time = open_time
            last_open_time = open_time
            rows += 1

    return {
        "file": csv_path.name,
        "rows": rows,
        "first_open_time": first_open_time or "",
        "last_open_time": last_open_time or "",
        "duplicates": duplicates,
        "nonstandard_intervals": gaps,
        "malformed_rows": malformed,
        "valid": malformed == 0 and duplicates == 0,
    }


def interval_to_minutes(interval: str) -> int:
    if interval.endswith("m"):
        return int(interval[:-1])
    if interval.endswith("h"):
        return int(interval[:-1]) * 60
    if interval.endswith("d"):
        return int(interval[:-1]) * 1440
    raise ValueError(
        f"Validation currently supports minute/hour/day intervals, not '{interval}'."
    )


def build_jobs(
    symbols: list[str],
    intervals: list[str],
    start: str,
    end: str | None,
    output_dir: Path,
) -> list[Job]:
    jobs: list[Job] = []
    for symbol in symbols:
        symbol = symbol.upper().strip()
        for interval in intervals:
            interval = interval.strip()
            for year_month in month_range(start, end):
                filename = f"{symbol}-{interval}-{year_month}.zip"
                url = f"{BASE_URL}/{symbol}/{interval}/{filename}"
                checksum_url = f"{url}.CHECKSUM"
                raw_dir = output_dir / symbol / interval / "raw"
                extract_dir = output_dir / symbol / interval / "csv"
                jobs.append(
                    Job(
                        symbol=symbol,
                        interval=interval,
                        year_month=year_month,
                        url=url,
                        checksum_url=checksum_url,
                        zip_path=raw_dir / filename,
                        checksum_path=raw_dir / f"{filename}.CHECKSUM",
                        extract_dir=extract_dir,
                    )
                )
    return jobs


def process_job(
    job: Job,
    *,
    retries: int,
    timeout: int,
    extract: bool,
    verify: bool,
    force: bool,
) -> DownloadResult:
    try:
        if not force and job.zip_path.exists() and job.checksum_path.exists():
            if verify:
                expected = expected_sha256(job.checksum_path)
                actual = file_sha256(job.zip_path)
                if actual == expected:
                    status = "cached"
                else:
                    job.zip_path.unlink(missing_ok=True)
                    status = "redownload"
            else:
                status = "cached"
        else:
            status = "download"

        if status in {"download", "redownload"}:
            if not remote_file_exists(job.url, timeout):
                return DownloadResult(
                    job=job,
                    status="not_available",
                    message="Monthly archive is not available.",
                )

            download_text(
                job.checksum_url,
                job.checksum_path,
                retries=retries,
                timeout=timeout,
            )
            bytes_downloaded = download_with_resume(
                job.url,
                job.zip_path,
                retries=retries,
                timeout=timeout,
            )
        else:
            bytes_downloaded = 0

        verified = False
        if verify:
            expected = expected_sha256(job.checksum_path)
            actual = file_sha256(job.zip_path)
            if actual != expected:
                raise ValueError(
                    f"Checksum mismatch for {job.zip_path.name}: "
                    f"expected {expected}, got {actual}"
                )
            verified = True

        extracted_csv = ""
        if extract:
            extracted_files = extract_zip(job.zip_path, job.extract_dir)
            csv_files = [path for path in extracted_files if path.suffix.lower() == ".csv"]
            if csv_files:
                extracted_csv = csv_files[0].name

        return DownloadResult(
            job=job,
            status="ok" if status != "cached" else "cached",
            message="Completed.",
            bytes_downloaded=bytes_downloaded,
            verified=verified,
            extracted_csv=extracted_csv,
        )
    except Exception as exc:
        return DownloadResult(
            job=job,
            status="error",
            message=str(exc),
        )


def merge_monthly_csvs(
    symbol: str,
    interval: str,
    output_dir: Path,
) -> Path | None:
    csv_dir = output_dir / symbol / interval / "csv"
    monthly_files = sorted(csv_dir.glob(f"{symbol}-{interval}-*.csv"))
    if not monthly_files:
        return None

    merged_path = output_dir / symbol / interval / f"{symbol}-{interval}-FULL.csv"
    last_open_time: int | None = None

    with merged_path.open("w", encoding="utf-8", newline="") as output:
        writer = csv.writer(output)
        writer.writerow(STANDARD_COLUMNS)

        for source_path in monthly_files:
            with source_path.open("r", encoding="utf-8-sig", newline="") as source:
                reader = csv.reader(source)
                for index, row in enumerate(reader):
                    if index == 0 and looks_like_header(row):
                        continue
                    if len(row) != 12:
                        continue
                    try:
                        open_time = int(float(row[0]))
                    except ValueError:
                        continue

                    if last_open_time is not None and open_time <= last_open_time:
                        continue

                    writer.writerow(row)
                    last_open_time = open_time

    return merged_path


def write_reports(
    results: list[DownloadResult],
    validations: list[dict[str, object]],
    output_dir: Path,
) -> None:
    report_path = output_dir / "download_report.csv"
    with report_path.open("w", encoding="utf-8", newline="") as handle:
        fieldnames = [
            "symbol",
            "interval",
            "year_month",
            "status",
            "verified",
            "bytes_downloaded",
            "zip_path",
            "extracted_csv",
            "message",
        ]
        writer = csv.DictWriter(handle, fieldnames=fieldnames)
        writer.writeheader()
        for result in results:
            writer.writerow(
                {
                    "symbol": result.job.symbol,
                    "interval": result.job.interval,
                    "year_month": result.job.year_month,
                    "status": result.status,
                    "verified": result.verified,
                    "bytes_downloaded": result.bytes_downloaded,
                    "zip_path": str(result.job.zip_path),
                    "extracted_csv": result.extracted_csv,
                    "message": result.message,
                }
            )

    validation_path = output_dir / "validation_report.csv"
    if validations:
        with validation_path.open("w", encoding="utf-8", newline="") as handle:
            writer = csv.DictWriter(handle, fieldnames=list(validations[0].keys()))
            writer.writeheader()
            writer.writerows(validations)


def main() -> int:
    parser = argparse.ArgumentParser(
        description=(
            "Download all available Binance Vision USD-M Futures monthly kline archives."
        )
    )
    parser.add_argument(
        "--symbols",
        nargs="+",
        default=["ETHUSDC", "SOLUSDC"],
        help="Symbols to download. Default: ETHUSDC SOLUSDC",
    )
    parser.add_argument(
        "--intervals",
        nargs="+",
        default=["5m"],
        help="Kline intervals. Default: 5m",
    )
    parser.add_argument(
        "--start",
        default="2020-01",
        help="First month to probe, formatted YYYY-MM. Default: 2020-01",
    )
    parser.add_argument(
        "--end",
        default=None,
        help="Last month to probe, formatted YYYY-MM. Default: previous month",
    )
    parser.add_argument(
        "--output",
        default="binance_data",
        help="Output directory. Default: binance_data",
    )
    parser.add_argument(
        "--workers",
        type=int,
        default=4,
        help="Parallel download workers. Default: 4",
    )
    parser.add_argument(
        "--retries",
        type=int,
        default=5,
        help="Retries for failed downloads. Default: 5",
    )
    parser.add_argument(
        "--timeout",
        type=int,
        default=60,
        help="Network timeout in seconds. Default: 60",
    )
    parser.add_argument(
        "--no-extract",
        action="store_true",
        help="Keep ZIP files without extracting CSV files.",
    )
    parser.add_argument(
        "--no-verify",
        action="store_true",
        help="Skip SHA-256 checksum verification.",
    )
    parser.add_argument(
        "--no-validate",
        action="store_true",
        help="Skip CSV schema and continuity validation.",
    )
    parser.add_argument(
        "--no-merge",
        action="store_true",
        help="Skip creation of one merged CSV per symbol/interval.",
    )
    parser.add_argument(
        "--force",
        action="store_true",
        help="Download archives again even when verified files already exist.",
    )
    args = parser.parse_args()

    output_dir = Path(args.output).resolve()
    output_dir.mkdir(parents=True, exist_ok=True)

    print("Binance Vision automatic downloader")
    print(f"Output directory: {output_dir}")
    print(f"Symbols: {', '.join(args.symbols)}")
    print(f"Intervals: {', '.join(args.intervals)}")
    print(f"Month range: {args.start} to {args.end or 'previous month'}")
    print()

    jobs = build_jobs(
        symbols=args.symbols,
        intervals=args.intervals,
        start=args.start,
        end=args.end,
        output_dir=output_dir,
    )

    results: list[DownloadResult] = []
    with ThreadPoolExecutor(max_workers=max(1, args.workers)) as executor:
        futures = {
            executor.submit(
                process_job,
                job,
                retries=max(1, args.retries),
                timeout=max(5, args.timeout),
                extract=not args.no_extract,
                verify=not args.no_verify,
                force=args.force,
            ): job
            for job in jobs
        }

        completed = 0
        for future in as_completed(futures):
            result = future.result()
            results.append(result)
            completed += 1
            print(
                f"[{completed}/{len(jobs)}] "
                f"{result.job.symbol} {result.job.interval} "
                f"{result.job.year_month}: {result.status}"
            )
            if result.status == "error":
                print(f"    {result.message}")

    validations: list[dict[str, object]] = []
    if not args.no_extract and not args.no_validate:
        print()
        print("Validating extracted CSV files...")
        for symbol in args.symbols:
            for interval in args.intervals:
                minutes = interval_to_minutes(interval)
                csv_dir = output_dir / symbol.upper() / interval / "csv"
                for csv_path in sorted(csv_dir.glob("*.csv")):
                    validations.append(validate_csv(csv_path, minutes))

    if not args.no_extract and not args.no_merge:
        print()
        print("Merging monthly CSV files...")
        for symbol in args.symbols:
            for interval in args.intervals:
                merged = merge_monthly_csvs(
                    symbol.upper(),
                    interval,
                    output_dir,
                )
                if merged:
                    print(f"Created: {merged}")

    results.sort(key=lambda item: (
        item.job.symbol,
        item.job.interval,
        item.job.year_month,
    ))
    write_reports(results, validations, output_dir)

    ok_count = sum(result.status in {"ok", "cached"} for result in results)
    missing_count = sum(result.status == "not_available" for result in results)
    error_count = sum(result.status == "error" for result in results)

    print()
    print("Finished.")
    print(f"Available archives completed: {ok_count}")
    print(f"Months not available: {missing_count}")
    print(f"Errors: {error_count}")
    print(f"Download report: {output_dir / 'download_report.csv'}")
    if validations:
        print(f"Validation report: {output_dir / 'validation_report.csv'}")

    return 1 if error_count else 0


if __name__ == "__main__":
    raise SystemExit(main())