"""
IcebergClient - Query OHLC data from Iceberg warehouse (Iceberg 1.10.1)

Tickers use Nautilus format: "BTC/USDT.BINANCE"
All timestamps are nanoseconds since epoch.
"""

import tracemalloc
from pathlib import Path
from typing import Optional, List, Tuple
import pandas as pd
import logging
from pyiceberg.catalog import load_catalog
from pyiceberg.expressions import (
    And,
    EqualTo,
    GreaterThanOrEqual,
    LessThan,
)

log = logging.getLogger(__name__)


def _rss_mb() -> str:
    """Return current VmRSS and VmPeak from /proc/self/status as a short string."""
    try:
        info = {}
        for line in Path("/proc/self/status").read_text().splitlines():
            for key in ("VmRSS", "VmPeak", "VmSize"):
                if line.startswith(f"{key}:"):
                    info[key] = int(line.split()[1]) // 1024  # kB → MB
        return f"RSS={info.get('VmRSS','?')}MB peak={info.get('VmPeak','?')}MB virt={info.get('VmSize','?')}MB"
    except Exception:
        return "?"


class IcebergClient:
    """
    Client for querying OHLC data from Iceberg warehouse (Iceberg 1.10.1).

    Note: Iceberg 1.x does not enforce primary keys at the table level.
    Deduplication is handled by:
    - Flink upsert mode with equality delete files
    - PyIceberg automatically filters deleted rows during queries
    - Last-write-wins semantics for duplicates

    Provides:
    - Query OHLC data by ticker, period, and time range
    - Identify missing data gaps
    - Efficient partition pruning for large datasets
    """

    def __init__(
        self,
        catalog_uri: str,
        namespace: str = "trading",
        s3_endpoint: Optional[str] = None,
        s3_access_key: Optional[str] = None,
        s3_secret_key: Optional[str] = None,
        s3_region: Optional[str] = None,
    ):
        """
        Initialize Iceberg client.

        Args:
            catalog_uri: URI of the Iceberg catalog (e.g., "http://iceberg-catalog:8181")
            namespace: Iceberg namespace (default: "trading")
            s3_endpoint: S3/MinIO endpoint URL (e.g., "http://localhost:9000")
            s3_access_key: S3/MinIO access key
            s3_secret_key: S3/MinIO secret key
            s3_region: S3/MinIO region (e.g., "us-east-1")
        """
        self.catalog_uri = catalog_uri
        self.namespace = namespace

        catalog_props = {"uri": catalog_uri}
        if s3_endpoint:
            catalog_props["s3.endpoint"] = s3_endpoint
            catalog_props["s3.path-style-access"] = "true"
        if s3_access_key:
            catalog_props["s3.access-key-id"] = s3_access_key
        if s3_secret_key:
            catalog_props["s3.secret-access-key"] = s3_secret_key
        if s3_region:
            catalog_props["s3.region"] = s3_region

        self.catalog = load_catalog("trading", **catalog_props)
        self.table = self.catalog.load_table(f"{namespace}.ohlc")

    def query_ohlc(
        self,
        ticker: str,
        period_seconds: int,
        start_time: int,
        end_time: int,
        columns: Optional[List[str]] = None
    ) -> pd.DataFrame:
        """
        Query OHLC data for a specific ticker, period, and time range.

        Args:
            ticker: Market identifier in Nautilus format (e.g., "BTC/USDT.BINANCE")
            period_seconds: OHLC period in seconds (60, 300, 3600, etc.)
            start_time: Start timestamp in nanoseconds
            end_time: End timestamp in nanoseconds (exclusive)
            columns: Optional list of columns to select.

        Returns:
            DataFrame with OHLC data sorted by timestamp, with a DatetimeIndex (UTC).
        """
        # Reload table metadata to pick up snapshots committed after this client was initialized
        self.table = self.catalog.load_table(f"{self.namespace}.ohlc")

        # Always fetch ingested_at alongside requested columns so we can deduplicate
        # duplicate parquet files left by repeated Flink job runs on the same key range.
        fetch_columns = None
        if columns is not None:
            fetch_columns = list(columns) if "ingested_at" in columns else list(columns) + ["ingested_at"]

        scan = self.table.scan(
            row_filter=And(
                EqualTo("ticker", ticker),
                EqualTo("period_seconds", period_seconds),
                GreaterThanOrEqual("timestamp", start_time),
                LessThan("timestamp", end_time)  # end_time is exclusive
            )
        )

        if fetch_columns is not None:
            scan = scan.select(*fetch_columns)

        if not tracemalloc.is_tracing():
            tracemalloc.start()
        tm_before = tracemalloc.take_snapshot()
        log.info("MEM before scan.to_pandas(): %s", _rss_mb())

        df = scan.to_pandas()

        log.info("MEM after scan.to_pandas(): %s | rows=%d cols=%s mem=%dMB",
                 _rss_mb(), len(df), list(df.columns),
                 df.memory_usage(deep=True).sum() // (1024 * 1024))
        tm_after = tracemalloc.take_snapshot()
        top = tm_after.compare_to(tm_before, "lineno")
        for stat in top[:5]:
            log.info("TRACEMALLOC: %s", stat)

        if not df.empty:
            # Deduplicate: keep the most-recently-ingested row per timestamp.
            if "ingested_at" in df.columns:
                df = (
                    df.sort_values("ingested_at", ascending=False)
                      .drop_duplicates(subset=["timestamp"])
                )
            log.info("MEM after dedup: %s | rows=%d", _rss_mb(), len(df))
            # Drop ingested_at if the caller did not ask for it
            if columns is not None and "ingested_at" not in columns and "ingested_at" in df.columns:
                df = df.drop(columns=["ingested_at"])

            df = df.sort_values("timestamp")
            # Convert integer nanosecond timestamps to DatetimeIndex
            df.index = pd.to_datetime(df["timestamp"], unit="ns", utc=True)
            df.index.name = "datetime"
            df = df.drop(columns=["timestamp"])

        return df

    def find_missing_ranges(
        self,
        ticker: str,
        period_seconds: int,
        start_time: int,
        end_time: int
    ) -> List[Tuple[int, int]]:
        """
        Identify missing data ranges in the requested time period.

        Args:
            ticker: Market identifier in Nautilus format
            period_seconds: OHLC period in seconds
            start_time: Start timestamp in nanoseconds
            end_time: End timestamp in nanoseconds

        Returns:
            List of (start_time, end_time) tuples for missing ranges (nanoseconds)
        """
        df = self.query_ohlc(ticker, period_seconds, start_time, end_time)

        if df.empty:
            return [(start_time, end_time)]

        # Convert period to nanoseconds
        period_nanos = period_seconds * 1_000_000_000

        # Generate expected timestamps — end_time is exclusive
        expected_timestamps = list(range(start_time, end_time, period_nanos))
        # DatetimeIndex backed by nanoseconds — view as int64 directly
        actual_timestamps = set(df.index.view('int64'))

        missing = sorted(set(expected_timestamps) - actual_timestamps)

        if not missing:
            return []

        # Consolidate consecutive missing timestamps into ranges
        ranges = []
        range_start = missing[0]
        prev_ts = missing[0]

        for ts in missing[1:]:
            if ts > prev_ts + period_nanos:
                ranges.append((range_start, prev_ts))
                range_start = ts
            prev_ts = ts

        ranges.append((range_start, prev_ts))
        return ranges

    def has_data(
        self,
        ticker: str,
        period_seconds: int,
        start_time: int,
        end_time: int
    ) -> bool:
        """
        Check if any data exists for the given parameters.

        Args:
            ticker: Market identifier in Nautilus format
            period_seconds: OHLC period in seconds
            start_time: Start timestamp in nanoseconds
            end_time: End timestamp in nanoseconds

        Returns:
            True if at least one candle exists, False otherwise
        """
        df = self.query_ohlc(ticker, period_seconds, start_time, end_time)
        return not df.empty
