"""Parquet adapter with Rust-accelerated I/O.

Uses Rust for fast parquet reading and data processing.
Falls back to pandas for operations not yet ported to Rust.
"""

import json
from pathlib import Path
from typing import Dict, List, Any, Optional
from datetime import datetime
from decimal import Decimal
from collections import defaultdict

import pandas as pd
from loguru import logger

from chainswarm_core.constants import PatternTypes

# Import Rust I/O functions
try:
    from chainswarm_analyzers_baseline._rust import (
        load_transfers_and_build_timestamps,
        load_money_flows,
        load_transfers,
    )
    RUST_IO_AVAILABLE = True
except ImportError:
    RUST_IO_AVAILABLE = False
    logger.warning("Rust I/O not available, falling back to pandas")


class ParquetAdapter:
    """Parquet adapter with Rust-accelerated I/O."""

    PATTERN_TYPE_FILES = {
        PatternTypes.CYCLE: 'patterns_cycle.parquet',
        PatternTypes.LAYERING_PATH: 'patterns_layering.parquet',
        PatternTypes.SMURFING_NETWORK: 'patterns_network.parquet',
        PatternTypes.PROXIMITY_RISK: 'patterns_proximity.parquet',
        PatternTypes.MOTIF_FANIN: 'patterns_motif.parquet',
        PatternTypes.MOTIF_FANOUT: 'patterns_motif.parquet',
        PatternTypes.TEMPORAL_BURST: 'patterns_burst.parquet',
        PatternTypes.THRESHOLD_EVASION: 'patterns_threshold.parquet',
    }

    def __init__(self, input_path: Path, output_path: Path, use_rust: bool = True):
        self.input_path = Path(input_path)
        self.output_path = Path(output_path)
        self.use_rust = use_rust and RUST_IO_AVAILABLE

        # Create output directory if it doesn't exist
        self.output_path.mkdir(parents=True, exist_ok=True)

        # Cache for loaded dataframes (used when Rust not available)
        self._transfers_df: Optional[pd.DataFrame] = None
        self._labels_df: Optional[pd.DataFrame] = None
        self._assets_df: Optional[pd.DataFrame] = None
        self._prices_df: Optional[pd.DataFrame] = None
        self._money_flows_df: Optional[pd.DataFrame] = None

        if self.use_rust:
            logger.info("Using Rust-accelerated I/O")
        else:
            logger.info("Using pandas I/O (Rust not available)")

    def read_transfers(
        self,
        start_timestamp_ms: int,
        end_timestamp_ms: int
    ) -> List[Dict[str, Any]]:
        """Read transfers from parquet."""
        transfers_file = self.input_path / "transfers.parquet"

        if not transfers_file.exists():
            raise ValueError(f"Transfers file not found: {transfers_file}")

        if self.use_rust:
            logger.info(f"Loading transfers from {transfers_file} (Rust)")
            raw_transfers = load_transfers(
                str(transfers_file),
                start_timestamp_ms,
                end_timestamp_ms
            )
            # Convert tuples to dicts
            transfers = [
                {
                    'from_address': t[0],
                    'to_address': t[1],
                    'block_timestamp': t[2],
                    'amount_usd': t[3],
                }
                for t in raw_transfers
            ]
            logger.info(f"Loaded {len(transfers)} transfers (Rust)")
            return transfers

        # Fallback to pandas
        if self._transfers_df is None:
            logger.info(f"Loading transfers from {transfers_file}")
            self._transfers_df = pd.read_parquet(transfers_file, engine='pyarrow')
            logger.info(f"Loaded {len(self._transfers_df)} transfers")

        df = self._transfers_df

        mask = (df['block_timestamp'] >= start_timestamp_ms) & \
               (df['block_timestamp'] < end_timestamp_ms)
        filtered_df = df[mask]

        logger.info(f"Filtered to {len(filtered_df)} transfers in time range")

        return filtered_df.to_dict(orient='records')

    def read_money_flows(
        self,
        start_timestamp_ms: int,
        end_timestamp_ms: int
    ) -> List[Dict[str, Any]]:
        """Read money flows from parquet."""
        money_flows_file = self.input_path / "money_flows.parquet"

        if not money_flows_file.exists():
            raise ValueError(f"Money flows file not found: {money_flows_file}")

        if self.use_rust:
            logger.info(f"Loading money flows from {money_flows_file} (Rust)")
            sources, targets, weights = load_money_flows(
                str(money_flows_file),
                start_timestamp_ms,
                end_timestamp_ms
            )
            # Convert to list of dicts
            flows = [
                {
                    'from_address': src,
                    'to_address': tgt,
                    'amount_usd_sum': weight,
                }
                for src, tgt, weight in zip(sources, targets, weights)
            ]
            logger.info(f"Loaded {len(flows)} money flows (Rust)")
            if not flows:
                raise ValueError(f"No money flows in time range {start_timestamp_ms} to {end_timestamp_ms}")
            return flows

        # Fallback to pandas
        if self._money_flows_df is None:
            logger.info(f"Loading money flows from {money_flows_file}")
            self._money_flows_df = pd.read_parquet(money_flows_file, engine='pyarrow')
            logger.info(f"Loaded {len(self._money_flows_df)} money flows")

        df = self._money_flows_df

        # Filter by last_seen_timestamp to get flows active in the window
        mask = (df['last_seen_timestamp'] >= start_timestamp_ms) & \
               (df['last_seen_timestamp'] < end_timestamp_ms)
        filtered_df = df[mask]

        if filtered_df.empty:
            raise ValueError(f"No money flows in time range {start_timestamp_ms} to {end_timestamp_ms}")

        logger.info(f"Filtered to {len(filtered_df)} money flows in time range")

        return filtered_df.to_dict(orient='records')

    def read_money_flows_raw(
        self,
        start_timestamp_ms: int,
        end_timestamp_ms: int
    ) -> tuple:
        """Read money flows and return raw arrays for direct graph building.

        Returns:
            (sources, targets, weights) tuple for RustGraph.add_edges()
        """
        money_flows_file = self.input_path / "money_flows.parquet"

        if not money_flows_file.exists():
            raise ValueError(f"Money flows file not found: {money_flows_file}")

        if self.use_rust:
            logger.info(f"Loading money flows from {money_flows_file} (Rust raw)")
            return load_money_flows(
                str(money_flows_file),
                start_timestamp_ms,
                end_timestamp_ms
            )

        # Fallback: load via pandas and convert
        flows = self.read_money_flows(start_timestamp_ms, end_timestamp_ms)
        sources = [f['from_address'] for f in flows]
        targets = [f['to_address'] for f in flows]
        weights = [f['amount_usd_sum'] for f in flows]
        return sources, targets, weights

    def read_assets(self, network: Optional[str] = None) -> List[Dict[str, Any]]:
        """Read assets from parquet."""
        assets_file = self.input_path / "assets.parquet"

        if not assets_file.exists():
            raise ValueError(f"Assets file not found: {assets_file}")

        if self._assets_df is None:
            logger.info(f"Loading assets from {assets_file}")
            self._assets_df = pd.read_parquet(assets_file, engine='pyarrow')
            logger.info(f"Loaded {len(self._assets_df)} assets")

        df = self._assets_df

        if network:
            df = df[df['network'] == network]

        logger.info(f"Returning {len(df)} assets")

        return df.to_dict(orient='records')

    def read_asset_prices(
        self,
        start_timestamp_ms: int,
        end_timestamp_ms: int
    ) -> List[Dict[str, Any]]:
        """Read asset prices from parquet."""
        prices_file = self.input_path / "asset_prices.parquet"

        if not prices_file.exists():
            raise ValueError(f"Asset prices file not found: {prices_file}")

        if self._prices_df is None:
            logger.info(f"Loading asset prices from {prices_file}")
            self._prices_df = pd.read_parquet(prices_file, engine='pyarrow')
            logger.info(f"Loaded {len(self._prices_df)} asset price records")

        df = self._prices_df

        start_date = pd.Timestamp(start_timestamp_ms, unit='ms')
        end_date = pd.Timestamp(end_timestamp_ms, unit='ms')

        if 'price_date' in df.columns:
            price_dates = pd.to_datetime(df['price_date'])
            mask = (price_dates >= start_date) & (price_dates < end_date)
            filtered_df = df[mask]
        else:
            filtered_df = df

        logger.info(f"Returning {len(filtered_df)} asset prices in time range")

        return filtered_df.to_dict(orient='records')

    def read_address_labels(
        self,
        addresses: List[str]
    ) -> Dict[str, Dict[str, Any]]:
        """Read address labels from parquet."""
        labels_file = self.input_path / "address_labels.parquet"

        if not labels_file.exists():
            logger.info(f"Address labels file not found: {labels_file}, returning empty labels")
            return {}

        if self._labels_df is None:
            logger.info(f"Loading address labels from {labels_file}")
            self._labels_df = pd.read_parquet(labels_file, engine='pyarrow')
            logger.info(f"Loaded {len(self._labels_df)} address labels")

        df = self._labels_df

        # Use set for O(1) lookup
        addresses_set = set(addresses)
        mask = df['address'].isin(addresses_set)
        filtered_df = df[mask]

        # Convert to dict efficiently using vectorized operations
        result = {}
        if not filtered_df.empty:
            for col in ['label', 'address_type', 'trust_level', 'source']:
                if col not in filtered_df.columns:
                    filtered_df[col] = ''

            for row in filtered_df.itertuples(index=False):
                result[row.address] = {
                    'label': getattr(row, 'label', ''),
                    'address_type': getattr(row, 'address_type', ''),
                    'trust_level': getattr(row, 'trust_level', ''),
                    'source': getattr(row, 'source', ''),
                }

        logger.info(f"Found labels for {len(result)} of {len(addresses)} addresses")

        return result

    def read_transfer_timestamps(
        self,
        start_timestamp_ms: int,
        end_timestamp_ms: int,
        addresses: List[str]
    ) -> Dict[str, List[Dict[str, Any]]]:
        """Extract timestamp data for burst detection from transfers.

        Uses Rust for fast parallel processing when available.

        Args:
            start_timestamp_ms: Start of time window in milliseconds
            end_timestamp_ms: End of time window in milliseconds
            addresses: List of addresses to extract timestamp data for

        Returns:
            Dict mapping address -> list of timestamp records.
            Each record: {'timestamp': int, 'volume': float, 'counterparty': str}
        """
        if not addresses:
            return {}

        transfers_file = self.input_path / "transfers.parquet"

        if not transfers_file.exists():
            raise ValueError(f"Transfers file not found: {transfers_file}")

        if self.use_rust:
            logger.info(f"Building timestamp data from {transfers_file} (Rust)")
            raw_data = load_transfers_and_build_timestamps(
                str(transfers_file),
                start_timestamp_ms,
                end_timestamp_ms,
                addresses
            )
            # Convert tuples to dicts
            result = {}
            for addr, records in raw_data.items():
                result[addr] = [
                    {'timestamp': ts, 'volume': vol, 'counterparty': cp}
                    for ts, vol, cp in records
                ]
            logger.info(f"Built timestamps for {len(result)} addresses (Rust)")
            return result

        # Fallback to pandas
        transfers = self.read_transfers(start_timestamp_ms, end_timestamp_ms)

        address_timestamps = defaultdict(list)
        addresses_set = set(addresses)

        for transfer in transfers:
            from_addr = transfer['from_address']
            to_addr = transfer['to_address']
            timestamp = transfer['block_timestamp']
            volume = transfer.get('amount_usd', 0) or 0

            if from_addr in addresses_set:
                address_timestamps[from_addr].append({
                    'timestamp': timestamp,
                    'volume': float(volume),
                    'counterparty': to_addr
                })

            if to_addr in addresses_set:
                address_timestamps[to_addr].append({
                    'timestamp': timestamp,
                    'volume': float(volume),
                    'counterparty': from_addr
                })

        logger.info(f"Read timestamps for {len(address_timestamps)} addresses")

        return dict(address_timestamps)

    LIST_COLUMNS = ['hourly_activity', 'daily_activity']

    def _convert_value(self, value: Any, as_json: bool = False) -> Any:
        if isinstance(value, Decimal):
            return float(value)
        if isinstance(value, list):
            if as_json:
                return json.dumps(value)
            return [self._convert_value(v) for v in value]
        if isinstance(value, dict):
            return {k: self._convert_value(v) for k, v in value.items()}
        return value

    def _normalize_features(self, features: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        normalized = []
        for feature in features:
            normalized_feature = {}
            for key, value in feature.items():
                is_list_col = key in self.LIST_COLUMNS
                normalized_feature[key] = self._convert_value(value, as_json=is_list_col)
            normalized.append(normalized_feature)
        return normalized

    def write_features(
        self,
        features: List[Dict[str, Any]],
        window_days: int,
        processing_date: str
    ) -> None:
        """Write features to parquet."""
        if not features:
            raise ValueError("No features to write")

        date_obj = datetime.strptime(processing_date, '%Y-%m-%d').date()

        for feature in features:
            feature['window_days'] = window_days
            feature['processing_date'] = date_obj

        features = self._normalize_features(features)

        output_file = self.output_path / "features.parquet"

        df = pd.DataFrame(features)
        df.to_parquet(output_file, engine='pyarrow', index=False)

        logger.info(f"Wrote {len(features)} features to {output_file}")

    def write_patterns(
        self,
        patterns: List[Dict[str, Any]],
        window_days: int,
        processing_date: str
    ) -> None:
        """Write patterns to parquet files by type."""
        if not patterns:
            logger.info("No patterns to write")
            return

        date_obj = datetime.strptime(processing_date, '%Y-%m-%d').date()

        for pattern in patterns:
            pattern['window_days'] = window_days
            pattern['processing_date'] = date_obj

        patterns_by_type: Dict[str, List[Dict]] = defaultdict(list)
        for pattern in patterns:
            pattern_type = pattern.get('pattern_type', 'UNKNOWN')
            patterns_by_type[pattern_type].append(pattern)

        for pattern_type, type_patterns in patterns_by_type.items():
            if pattern_type not in self.PATTERN_TYPE_FILES:
                logger.warning(f"Unknown pattern type: {pattern_type}, skipping {len(type_patterns)} patterns")
                continue

            output_filename = self.PATTERN_TYPE_FILES[pattern_type]
            output_file = self.output_path / output_filename

            df = pd.DataFrame(type_patterns)

            if output_file.exists() and pattern_type in [PatternTypes.MOTIF_FANIN, PatternTypes.MOTIF_FANOUT]:
                existing_df = pd.read_parquet(output_file, engine='pyarrow')
                df = pd.concat([existing_df, df], ignore_index=True)

            df.to_parquet(output_file, engine='pyarrow', index=False)

            logger.info(f"Wrote {len(type_patterns)} {pattern_type} patterns to {output_file}")

        total_patterns = sum(len(p) for p in patterns_by_type.values())
        logger.info(f"Total patterns written: {total_patterns} across {len(patterns_by_type)} types")

    def clear_cache(self) -> None:
        """Clear cached dataframes."""
        self._transfers_df = None
        self._labels_df = None
        self._assets_df = None
        self._prices_df = None
        self._money_flows_df = None
        logger.debug("Cleared Parquet adapter cache")
