#!/usr/bin/env python3

# =================================================================
# CONFIGURATION - modify these values
# =================================================================

# Leave empty and set COINAPI_API_KEY, or paste a key only in a private local copy.
API_KEY = ""

# Public static S3 secret for CoinAPI Flat Files (not a personal credential).
AWS_SECRET_KEY = "coinapi"
AWS_REGION = "us-east-1"
ENDPOINT_URL = "https://s3.flatfiles.coinapi.io"

DEFAULT_BUCKET = "coinapi"  # coinapi or coinapi-daily-tail
DEFAULT_EXCHANGES = ["BINANCE"]  # None or [] = every exchange present in the range
DEFAULT_START_DATE = "2024-01-01"  # YYYY-MM-DD (UTC, inclusive)
DEFAULT_END_DATE = "2024-12-31"    # YYYY-MM-DD (UTC, inclusive)

# Dataset prefixes (without T-). See /products/flat-files/docs/datasets
DEFAULT_DATA_TYPES = ["LIMITBOOK_FULL"]
# TRADES, QUOTES, LIMITBOOK_FULL, LIMITBOOK_SNAPSHOT_50, LIMITBOOK_DEPTH_BANDS,
# OHLCV_ACTIVE_CONSOLIDATED, METRICS-V2, HLORACLEPRICES, HLTWAPSTATUSES,
# HLMISCEVENTS, HLSYSTEMEVENTS

DEFAULT_SYMBOLS = ["BINANCE_SPOT_BTC_USDT"]  # None or [] = all symbols
DEFAULT_ASSETS = None  # e.g. ["BTC", "ETH"]; ignored for OHLCV, Metrics V2, and HL misc/system
DEFAULT_OHLCV_PERIODS = None  # e.g. ["1MIN"]; None = all TP-* files

DEFAULT_OUTPUT_FORMAT = "text"  # text or json
DEFAULT_OUTPUT_FILE = "output.txt"  # None = print to console

# =================================================================
# Implementation
# =================================================================

import io
import json
import os
import re
import subprocess
import sys
from datetime import datetime, timedelta
from typing import Dict, List, Optional

MIB = 1024 ** 2
GIB = 1024 ** 3
HOURLY_START = datetime(2026, 6, 9, 0, 0)  # First hourly prefix: 2026060900
API_KEY_PLACEHOLDERS = {"", "YOUR_COINAPI_KEY", "YOUR_API_KEY", "Flat-Files-API-Key"}
DAILY_ONLY_DATA_TYPES = {"OHLCV_ACTIVE_CONSOLIDATED"}
HOURLY_ALWAYS_DATA_TYPES = {"METRICS-V2"}
EXCHANGE_LEVEL_DATA_TYPES = {
    "OHLCV_ACTIVE_CONSOLIDATED",
    "METRICS-V2",
    "HLMISCEVENTS",
    "HLSYSTEMEVENTS",
}
PERIOD_FILTER_DATA_TYPES = {"OHLCV_ACTIVE_CONSOLIDATED", "METRICS-V2"}


def resolve_api_key() -> str:
    key = os.environ.get("COINAPI_API_KEY", API_KEY).strip()
    if key in API_KEY_PLACEHOLDERS:
        raise SystemExit(
            "Set COINAPI_API_KEY in the environment, or assign API_KEY "
            "in a private local copy. Do not publish real keys."
        )
    return key


def bytes_to_mib(size: int) -> float:
    return size / MIB


def bytes_to_gib(size: int) -> float:
    return size / GIB


def format_size(size: int) -> str:
    return f"{size:,} bytes ({bytes_to_mib(size):.2f} MiB / {bytes_to_gib(size):.2f} GiB)"


def calendar_day(partition: str) -> str:
    return partition[:8]


def iter_partitions(start_date: str, end_date: str, data_type: str, bucket: str) -> List[str]:
    """Daily prefixes before 2026-06-09, hourly prefixes from 2026060900. Metrics V2 is hourly in coinapi."""
    start = datetime.strptime(start_date, "%Y-%m-%d")
    end = datetime.strptime(end_date, "%Y-%m-%d")
    if end < start:
        raise ValueError("DEFAULT_END_DATE must be on or after DEFAULT_START_DATE")

    partitions = []
    current = start
    while current <= end:
        use_hourly = (
            bucket == "coinapi"
            and data_type not in DAILY_ONLY_DATA_TYPES
            and (data_type in HOURLY_ALWAYS_DATA_TYPES or current >= HOURLY_START)
        )
        if use_hourly:
            for hour in range(24):
                partitions.append(current.strftime("%Y%m%d") + f"{hour:02d}")
        else:
            partitions.append(current.strftime("%Y%m%d"))
        current += timedelta(days=1)
    return partitions


class CoinAPIS3Calculator:
    def __init__(self, api_key: str, endpoint_url: str):
        self.endpoint_url = endpoint_url
        self.api_key = api_key

    def aws_ls(self, path: str) -> str:
        print(f"Checking {path}", file=sys.stderr)
        cmd = [
            "aws",
            "--endpoint-url", self.endpoint_url,
            "--region", AWS_REGION,
            "s3", "ls", path,
        ]
        env = os.environ.copy()
        env["AWS_ACCESS_KEY_ID"] = self.api_key
        env["AWS_SECRET_ACCESS_KEY"] = AWS_SECRET_KEY
        try:
            return subprocess.run(
                cmd, capture_output=True, text=True, check=True, env=env
            ).stdout
        except subprocess.CalledProcessError as error:
            print(f"Error getting S3 listing: {error}", file=sys.stderr)
            if error.stderr:
                print(error.stderr, file=sys.stderr)
            return ""

    def discover_exchanges(self, partition: str, data_type: str, bucket: str) -> List[str]:
        listing = self.aws_ls(f"s3://{bucket}/T-{data_type}/D-{partition}/")
        found = []
        for line in listing.split("\n"):
            parts = line.split()
            if len(parts) >= 2 and parts[0] == "PRE" and parts[1].startswith("E-"):
                found.append(parts[1].strip("/").replace("E-", "", 1))
        return found

    def parse_file_row(
        self, line: str, exchange: str, partition: str
    ) -> Optional[Dict]:
        parts = line.strip().split()
        if len(parts) < 4:
            return None
        filename = parts[3]
        size = int(parts[2])
        symbol_match = re.search(r"SC-([^+]+)", filename)
        period_match = re.search(r"TP-([A-Z0-9]+)", filename)
        if symbol_match:
            symbol = symbol_match.group(1)
        elif period_match:
            symbol = f"TP-{period_match.group(1)}"
        else:
            symbol = "EXCHANGE_FILE"
        return {
            "exchange": exchange,
            "partition": partition,
            "size": size,
            "filename": filename,
            "symbol": symbol,
        }

    def filter_lines(
        self,
        listing: str,
        data_type: str,
        symbols: Optional[List[str]] = None,
        assets: Optional[List[str]] = None,
        ohlcv_periods: Optional[List[str]] = None,
    ) -> List[str]:
        skip_symbol_filters = data_type in EXCHANGE_LEVEL_DATA_TYPES
        filtered = []
        for line in listing.split("\n"):
            if not line.strip():
                continue
            if data_type in PERIOD_FILTER_DATA_TYPES and ohlcv_periods:
                if not any(f"TP-{period}" in line for period in ohlcv_periods):
                    continue
            if not skip_symbol_filters and symbols:
                if not any(f"SC-{symbol}+" in line for symbol in symbols):
                    continue
            if not skip_symbol_filters and assets:
                upper_line = line.upper()
                if not any(
                    f"_{asset.upper()}_" in upper_line or f"_{asset.upper()}+" in upper_line
                    for asset in assets
                ):
                    continue
            filtered.append(line)
        return filtered

    def calculate_partition_size(
        self,
        exchange: str,
        partition: str,
        data_type: str,
        symbols: Optional[List[str]],
        assets: Optional[List[str]],
        ohlcv_periods: Optional[List[str]],
        bucket: str,
    ) -> Dict:
        listing = self.aws_ls(
            f"s3://{bucket}/T-{data_type}/D-{partition}/E-{exchange}/"
        )
        total_size = 0
        symbol_sizes = {}
        files_list = []

        for line in self.filter_lines(listing, data_type, symbols, assets, ohlcv_periods):
            row = self.parse_file_row(line, exchange, partition)
            if not row:
                continue
            total_size += row["size"]
            symbol_sizes[row["symbol"]] = symbol_sizes.get(row["symbol"], 0) + row["size"]
            files_list.append(row)

        return {
            "total_size": total_size,
            "symbols": symbol_sizes,
            "file_count": len(files_list),
            "files": files_list,
        }

    def calculate_date_range(
        self,
        exchanges: Optional[List[str]],
        start_date: str,
        end_date: str,
        data_types: Optional[List[str]] = None,
        symbols: Optional[List[str]] = None,
        assets: Optional[List[str]] = None,
        ohlcv_periods: Optional[List[str]] = None,
        bucket: str = DEFAULT_BUCKET,
    ) -> Dict:
        data_types = data_types or ["LIMITBOOK_FULL"]
        configured_exchanges = [item for item in (exchanges or []) if item]
        results = {
            "bucket": bucket,
            "start_date": start_date,
            "end_date": end_date,
            "total_size": 0,
            "by_type": {
                dt: {"size": 0, "files": 0, "files_list": []} for dt in data_types
            },
            "by_exchange": {},
            "by_symbol": {},
        }

        for data_type in data_types:
            last_day = None
            for partition in iter_partitions(start_date, end_date, data_type, bucket):
                day = calendar_day(partition)
                if day != last_day:
                    kind = "hourly" if len(partition) == 10 else "daily"
                    print(f"Processed {day} ({data_type}, {kind})", file=sys.stderr)
                    last_day = day

                if configured_exchanges:
                    partition_exchanges = configured_exchanges
                else:
                    partition_exchanges = self.discover_exchanges(
                        partition, data_type, bucket
                    )

                for exchange in partition_exchanges:
                    part = self.calculate_partition_size(
                        exchange,
                        partition,
                        data_type,
                        symbols,
                        assets,
                        ohlcv_periods,
                        bucket,
                    )
                    if part["file_count"] == 0:
                        continue

                    results["total_size"] += part["total_size"]
                    results["by_type"][data_type]["size"] += part["total_size"]
                    results["by_type"][data_type]["files"] += part["file_count"]
                    results["by_type"][data_type]["files_list"].extend(part["files"])

                    exchange_info = results["by_exchange"].setdefault(
                        exchange, {"size": 0, "files": 0}
                    )
                    exchange_info["size"] += part["total_size"]
                    exchange_info["files"] += part["file_count"]

                    for symbol, size in part["symbols"].items():
                        results["by_symbol"][symbol] = (
                            results["by_symbol"].get(symbol, 0) + size
                        )

        results["total_mib"] = bytes_to_mib(results["total_size"])
        results["total_gib"] = bytes_to_gib(results["total_size"])
        for info in results["by_type"].values():
            info["size_mib"] = bytes_to_mib(info["size"])
            info["size_gib"] = bytes_to_gib(info["size"])
            info["files_list"].sort(
                key=lambda row: (row["partition"], row["exchange"], row["filename"])
            )
        for info in results["by_exchange"].values():
            info["size_mib"] = bytes_to_mib(info["size"])
            info["size_gib"] = bytes_to_gib(info["size"])
        return results


def print_size_report(results: Dict) -> None:
    print("\nCoinAPI Flat Files Size Report")
    print("=" * 80)
    print()
    print(f"Total Size: {format_size(results['total_size'])}")
    print()
    print("By Exchange:")
    print("-" * 40)
    if not results["by_exchange"]:
        print("No files found.")
    for exchange, info in sorted(results["by_exchange"].items()):
        print(f"{exchange}:")
        print(f"  Size: {format_size(info['size'])}")
        print(f"  Files: {info['files']}")
        print()

    print("By Data Type:")
    print("-" * 40)
    for data_type, info in results["by_type"].items():
        print()
        print(f"{data_type}:")
        print(f"  Size: {format_size(info['size'])}")
        print(f"  Files: {info['files']}")
        if info["files_list"]:
            print()
            print("  Files:")
            for row in info["files_list"]:
                print(
                    f"    {row['exchange']} {row['partition']} "
                    f"{row['size']:,} {row['filename']}"
                )

    print()
    print(
        "Only files that were found are listed. Missing days or hours are omitted. "
        "Use these sizes with the Flat Files pricing page. "
        "Transfer tiers reset at the end of each UTC day, so cost is not included here."
    )


def main():
    results = CoinAPIS3Calculator(resolve_api_key(), ENDPOINT_URL).calculate_date_range(
        exchanges=DEFAULT_EXCHANGES,
        start_date=DEFAULT_START_DATE,
        end_date=DEFAULT_END_DATE,
        data_types=DEFAULT_DATA_TYPES,
        symbols=DEFAULT_SYMBOLS,
        assets=DEFAULT_ASSETS,
        ohlcv_periods=DEFAULT_OHLCV_PERIODS,
        bucket=DEFAULT_BUCKET,
    )

    if DEFAULT_OUTPUT_FORMAT == "json":
        output_content = json.dumps(results, indent=2)
    else:
        buffer = io.StringIO()
        original_stdout = sys.stdout
        sys.stdout = buffer
        try:
            print_size_report(results)
        finally:
            sys.stdout = original_stdout
        output_content = buffer.getvalue()

    if DEFAULT_OUTPUT_FILE:
        with open(DEFAULT_OUTPUT_FILE, "w", encoding="utf-8") as handle:
            handle.write(output_content)
        print(f"Results written to {DEFAULT_OUTPUT_FILE}", file=sys.stderr)
    else:
        print(output_content)


if __name__ == "__main__":
    main()

CoinAPI Flat Files Size Report
================================================================================

Total Size: 1,046,195,958,815 bytes (997730.22 MiB / 974.35 GiB)

By Exchange:
----------------------------------------
BINANCEFTS:
  Size: 1,046,195,958,815 bytes (997730.22 MiB / 974.35 GiB)
  Files: 4236

By Data Type:
----------------------------------------

LIMITBOOK_FULL:
  Size: 1,046,195,958,815 bytes (997730.22 MiB / 974.35 GiB)
  Files: 4236

  Files:
    BINANCEFTS 20191007 19,973,468 IDDI-5677816+SC-BINANCEFTS_PERP_BTC_USDT+S-BTCUSDT.csv.gz
    BINANCEFTS 20191008 12,508,716 IDDI-5677816+SC-BINANCEFTS_PERP_BTC_USDT+S-BTCUSDT.csv.gz
    BINANCEFTS 20191009 13,506,394 IDDI-5677816+SC-BINANCEFTS_PERP_BTC_USDT+S-BTCUSDT.csv.gz
    BINANCEFTS 20191010 13,050,459 IDDI-5677816+SC-BINANCEFTS_PERP_BTC_USDT+S-BTCUSDT.csv.gz
    BINANCEFTS 20191011 14,142,177 IDDI-5677816+SC-BINANCEFTS_PERP_BTC_USDT+S-BTCUSDT.csv.gz