"""
Coordination Bounds Metrics Collection

Instrumentation for measuring and validating coordination bounds:
- Algebraic operations: Expected C = 0 (immediate commit)
- Generic operations: Expected C = O(log N) rounds

This module provides:
1. CommitMetrics - dataclass for recording operation metrics
2. OperationClassifier - classifies operations by algebraic signature
3. InstrumentedWriter - TableWriter wrapper with metrics collection
4. MetricsExporter - export metrics for analysis
"""

from __future__ import annotations

import json
import time
from dataclasses import dataclass, field, asdict
from enum import Enum
from typing import TYPE_CHECKING, Optional, List, Dict, Any, Union
from pathlib import Path

import pyarrow as pa

if TYPE_CHECKING:
    import pandas as pd


class AlgebraicSignature(Enum):
    """
    Algebraic classification of operations.

    Determines coordination requirements:
    - SEMILATTICE: Commutative, associative, idempotent -> C = 0
    - ABELIAN: Commutative, associative, has inverse -> C = 0
    - GENERIC: Non-commutative -> C = Omega(log N)
    """
    SEMILATTICE = "semilattice"
    ABELIAN = "abelian"
    GENERIC = "generic"

    @property
    def coordination_free(self) -> bool:
        """Can this operation type achieve zero coordination?"""
        return self in (AlgebraicSignature.SEMILATTICE, AlgebraicSignature.ABELIAN)

    @property
    def theoretical_min_rounds(self) -> int:
        """Theoretical minimum coordination rounds."""
        return 0 if self.coordination_free else -1  # -1 means O(log N)

    @property
    def theoretical_bound(self) -> str:
        """Human-readable theoretical bound."""
        return "0" if self.coordination_free else "Omega(log N)"


@dataclass
class CommitMetrics:
    """
    Metrics for a single commit operation.

    Records timing, classification, and coordination details
    for validation against theoretical bounds.
    """
    # Operation identification
    table_name: str
    operation_id: str  # Unique ID for this operation
    timestamp_ns: int  # Unix timestamp in nanoseconds

    # Classification
    algebraic_signature: AlgebraicSignature
    operation_type: str  # "write", "increment", "max", etc.

    # Timing (all in nanoseconds for precision)
    issue_time_ns: int  # When operation was issued
    commit_time_ns: int  # When operation was safely committed
    propagation_time_ns: Optional[int] = None  # When all replicas confirmed (if tracked)

    # Derived metrics
    @property
    def commit_latency_ns(self) -> int:
        """Time from issue to safe commit in nanoseconds."""
        return self.commit_time_ns - self.issue_time_ns

    @property
    def commit_latency_ms(self) -> float:
        """Time from issue to safe commit in milliseconds."""
        return self.commit_latency_ns / 1_000_000

    # Coordination details
    coordination_rounds: int = 0  # 0 for algebraic, >0 for coordinated
    messages_sent: int = 0
    messages_received: int = 0
    bytes_sent: int = 0
    bytes_received: int = 0

    # Data details
    num_rows: int = 0
    num_columns: int = 0
    data_bytes: int = 0

    # Validation
    @property
    def matches_theory(self) -> bool:
        """Does this measurement match theoretical predictions?"""
        if self.algebraic_signature.coordination_free:
            return self.coordination_rounds == 0
        else:
            # Generic ops should have coordination_rounds > 0
            return self.coordination_rounds > 0

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary for JSON export."""
        return {
            "table_name": self.table_name,
            "operation_id": self.operation_id,
            "timestamp_ns": self.timestamp_ns,
            "algebraic_signature": self.algebraic_signature.value,
            "operation_type": self.operation_type,
            "issue_time_ns": self.issue_time_ns,
            "commit_time_ns": self.commit_time_ns,
            "propagation_time_ns": self.propagation_time_ns,
            "commit_latency_ns": self.commit_latency_ns,
            "commit_latency_ms": self.commit_latency_ms,
            "coordination_rounds": self.coordination_rounds,
            "messages_sent": self.messages_sent,
            "messages_received": self.messages_received,
            "bytes_sent": self.bytes_sent,
            "bytes_received": self.bytes_received,
            "num_rows": self.num_rows,
            "num_columns": self.num_columns,
            "data_bytes": self.data_bytes,
            "matches_theory": self.matches_theory,
        }


class OperationClassifier:
    """
    Classifies database operations by algebraic signature.

    Classification rules:
    - ADD/INCREMENT: Abelian (commutative group with inverse)
    - MAX/MIN: Semilattice (commutative, associative, idempotent)
    - UNION/OR/AND: Semilattice
    - OVERWRITE/SET: Generic (non-commutative)
    - CAS: Generic (depends on current state)
    """

    # Known operation classifications
    CLASSIFICATIONS: Dict[str, AlgebraicSignature] = {
        # Abelian group operations
        "add": AlgebraicSignature.ABELIAN,
        "increment": AlgebraicSignature.ABELIAN,
        "decrement": AlgebraicSignature.ABELIAN,
        "sum": AlgebraicSignature.ABELIAN,
        "multiply": AlgebraicSignature.ABELIAN,

        # Semilattice operations
        "max": AlgebraicSignature.SEMILATTICE,
        "min": AlgebraicSignature.SEMILATTICE,
        "union": AlgebraicSignature.SEMILATTICE,
        "or": AlgebraicSignature.SEMILATTICE,
        "and": AlgebraicSignature.SEMILATTICE,
        "greatest": AlgebraicSignature.SEMILATTICE,
        "least": AlgebraicSignature.SEMILATTICE,

        # Generic operations (require coordination)
        "overwrite": AlgebraicSignature.GENERIC,
        "set": AlgebraicSignature.GENERIC,
        "replace": AlgebraicSignature.GENERIC,
        "cas": AlgebraicSignature.GENERIC,
        "compare_and_swap": AlgebraicSignature.GENERIC,
        "delete": AlgebraicSignature.GENERIC,
        "insert_unique": AlgebraicSignature.GENERIC,
    }

    @classmethod
    def classify(cls, operation_type: str) -> AlgebraicSignature:
        """
        Classify an operation by name.

        Args:
            operation_type: Name of the operation (e.g., "add", "max", "overwrite")

        Returns:
            AlgebraicSignature indicating coordination requirements
        """
        normalized = operation_type.lower().strip()
        return cls.CLASSIFICATIONS.get(normalized, AlgebraicSignature.GENERIC)

    @classmethod
    def classify_write(
        cls,
        table: pa.Table,
        merge_columns: Optional[Dict[str, str]] = None,
    ) -> AlgebraicSignature:
        """
        Classify a table write operation.

        If merge_columns specifies algebraic merge strategies for all columns,
        the operation is coordination-free. Otherwise, it's generic.

        Args:
            table: The Arrow table being written
            merge_columns: Dict mapping column names to merge strategies
                          (e.g., {"count": "sum", "max_value": "max"})

        Returns:
            AlgebraicSignature for the operation
        """
        if merge_columns is None:
            # No merge strategy specified = full overwrite = generic
            return AlgebraicSignature.GENERIC

        # Check each column's merge strategy
        for col_name in table.column_names:
            if col_name not in merge_columns:
                # Column without merge strategy = overwrite = generic
                return AlgebraicSignature.GENERIC

            strategy = merge_columns[col_name]
            signature = cls.classify(strategy)

            if not signature.coordination_free:
                # Any non-algebraic column makes the whole operation generic
                return AlgebraicSignature.GENERIC

        # All columns have algebraic merge strategies
        # Return the "weakest" algebraic type
        has_abelian = any(
            cls.classify(s) == AlgebraicSignature.ABELIAN
            for s in merge_columns.values()
        )
        return AlgebraicSignature.ABELIAN if has_abelian else AlgebraicSignature.SEMILATTICE

    @classmethod
    def classify_transaction(
        cls,
        operations: List[AlgebraicSignature],
    ) -> AlgebraicSignature:
        """
        Classify a multi-operation transaction.

        A transaction inherits the coordination cost of its "worst" operation.
        One generic operation forces the entire transaction to coordinate.

        Args:
            operations: List of AlgebraicSignatures for each operation

        Returns:
            AlgebraicSignature for the transaction
        """
        if not operations:
            return AlgebraicSignature.SEMILATTICE  # Empty transaction is trivially coordination-free

        # If any operation is generic, transaction is generic
        if any(op == AlgebraicSignature.GENERIC for op in operations):
            return AlgebraicSignature.GENERIC

        # All algebraic - return Abelian if any Abelian, else Semilattice
        if any(op == AlgebraicSignature.ABELIAN for op in operations):
            return AlgebraicSignature.ABELIAN

        return AlgebraicSignature.SEMILATTICE


@dataclass
class MetricsCollector:
    """
    Collects and aggregates CommitMetrics.

    Thread-safe collection of metrics with export capabilities.
    """
    metrics: List[CommitMetrics] = field(default_factory=list)
    _operation_counter: int = field(default=0, repr=False)

    def record(self, metric: CommitMetrics) -> None:
        """Record a new metric."""
        self.metrics.append(metric)

    def generate_operation_id(self) -> str:
        """Generate a unique operation ID."""
        self._operation_counter += 1
        return f"op_{self._operation_counter:08d}"

    def clear(self) -> None:
        """Clear all recorded metrics."""
        self.metrics.clear()
        self._operation_counter = 0

    # Aggregation methods

    def by_signature(self) -> Dict[AlgebraicSignature, List[CommitMetrics]]:
        """Group metrics by algebraic signature."""
        result: Dict[AlgebraicSignature, List[CommitMetrics]] = {
            AlgebraicSignature.SEMILATTICE: [],
            AlgebraicSignature.ABELIAN: [],
            AlgebraicSignature.GENERIC: [],
        }
        for m in self.metrics:
            result[m.algebraic_signature].append(m)
        return result

    def summary(self) -> Dict[str, Any]:
        """Generate summary statistics."""
        if not self.metrics:
            return {"total_operations": 0}

        by_sig = self.by_signature()

        def sig_stats(metrics: List[CommitMetrics]) -> Dict[str, Any]:
            if not metrics:
                return {"count": 0}
            latencies = [m.commit_latency_ms for m in metrics]
            return {
                "count": len(metrics),
                "mean_latency_ms": sum(latencies) / len(latencies),
                "min_latency_ms": min(latencies),
                "max_latency_ms": max(latencies),
                "total_coordination_rounds": sum(m.coordination_rounds for m in metrics),
                "all_match_theory": all(m.matches_theory for m in metrics),
            }

        algebraic_metrics = by_sig[AlgebraicSignature.SEMILATTICE] + by_sig[AlgebraicSignature.ABELIAN]
        generic_metrics = by_sig[AlgebraicSignature.GENERIC]

        result = {
            "total_operations": len(self.metrics),
            "semilattice": sig_stats(by_sig[AlgebraicSignature.SEMILATTICE]),
            "abelian": sig_stats(by_sig[AlgebraicSignature.ABELIAN]),
            "generic": sig_stats(by_sig[AlgebraicSignature.GENERIC]),
            "algebraic_total": sig_stats(algebraic_metrics),
            "validation": {
                "all_match_theory": all(m.matches_theory for m in self.metrics),
                "algebraic_coordination_free": all(
                    m.coordination_rounds == 0 for m in algebraic_metrics
                ) if algebraic_metrics else True,
            },
        }

        # Calculate speedup if we have both algebraic and generic
        if algebraic_metrics and generic_metrics:
            alg_mean = result["algebraic_total"]["mean_latency_ms"]
            gen_mean = result["generic"]["mean_latency_ms"]
            if alg_mean > 0:
                result["speedup_ratio"] = gen_mean / alg_mean

        return result

    def export_json(self, path: Union[str, Path]) -> None:
        """Export all metrics to JSON file."""
        data = {
            "metrics": [m.to_dict() for m in self.metrics],
            "summary": self.summary(),
        }
        with open(path, 'w') as f:
            json.dump(data, f, indent=2)

    def export_csv(self, path: Union[str, Path]) -> None:
        """Export metrics to CSV file."""
        import csv

        if not self.metrics:
            return

        fieldnames = list(self.metrics[0].to_dict().keys())

        with open(path, 'w', newline='') as f:
            writer = csv.DictWriter(f, fieldnames=fieldnames)
            writer.writeheader()
            for m in self.metrics:
                writer.writerow(m.to_dict())


class InstrumentedWriter:
    """
    TableWriter wrapper with coordination bounds instrumentation.

    Wraps a standard TableWriter to measure commit latencies and
    classify operations by algebraic signature for validation.

    Example:
        >>> from rhizo import TableWriter
        >>> from rhizo.metrics import InstrumentedWriter, MetricsCollector
        >>>
        >>> collector = MetricsCollector()
        >>> writer = TableWriter(store, catalog)
        >>> instrumented = InstrumentedWriter(writer, collector)
        >>>
        >>> # Algebraic write (with merge strategy)
        >>> instrumented.write("counters", df, merge_columns={"count": "sum"})
        >>>
        >>> # Generic write (no merge strategy = overwrite)
        >>> instrumented.write("users", df)
        >>>
        >>> # Check results
        >>> print(collector.summary())
        >>> collector.export_json("metrics.json")
    """

    def __init__(
        self,
        writer,  # TableWriter
        collector: Optional[MetricsCollector] = None,
    ):
        """
        Initialize instrumented writer.

        Args:
            writer: The underlying TableWriter to wrap
            collector: MetricsCollector to record metrics (creates new if None)
        """
        self._writer = writer
        self._collector = collector or MetricsCollector()

    @property
    def collector(self) -> MetricsCollector:
        """Access the metrics collector."""
        return self._collector

    def write(
        self,
        table_name: str,
        data: Union["pd.DataFrame", pa.Table],
        metadata: Optional[Dict[str, str]] = None,
        merge_columns: Optional[Dict[str, str]] = None,
        operation_type: str = "write",
    ):
        """
        Write data with instrumentation.

        Args:
            table_name: Name of the table
            data: Data to write
            metadata: Optional metadata
            merge_columns: Optional merge strategies for coordination-free writes
            operation_type: Type of operation for classification

        Returns:
            WriteResult from underlying writer
        """
        # Convert to Arrow for classification
        if isinstance(data, pa.Table):
            table = data
        else:
            table = pa.Table.from_pandas(data, preserve_index=False)

        # Classify the operation
        signature = OperationClassifier.classify_write(table, merge_columns)

        # Generate operation ID and record start time
        op_id = self._collector.generate_operation_id()
        issue_time = time.perf_counter_ns()
        timestamp = time.time_ns()

        # Perform the write
        result = self._writer.write(table_name, data, metadata)

        # Record commit time
        commit_time = time.perf_counter_ns()

        # Create and record metrics
        metric = CommitMetrics(
            table_name=table_name,
            operation_id=op_id,
            timestamp_ns=timestamp,
            algebraic_signature=signature,
            operation_type=operation_type,
            issue_time_ns=issue_time,
            commit_time_ns=commit_time,
            coordination_rounds=0,  # Rhizo is coordination-free for local commit
            num_rows=result.total_rows,
            num_columns=len(table.column_names),
            data_bytes=result.total_bytes,
        )

        self._collector.record(metric)

        return result

    def increment(
        self,
        table_name: str,
        column: str,
        delta: int = 1,
    ):
        """
        Increment operation (Abelian - coordination-free).

        This is a convenience method that creates a proper algebraic
        increment operation.
        """
        # This would need actual implementation tied to Rhizo's increment support
        # For now, we record it as an algebraic operation
        op_id = self._collector.generate_operation_id()
        issue_time = time.perf_counter_ns()
        timestamp = time.time_ns()

        # Simulate local increment (instant)
        commit_time = time.perf_counter_ns()

        metric = CommitMetrics(
            table_name=table_name,
            operation_id=op_id,
            timestamp_ns=timestamp,
            algebraic_signature=AlgebraicSignature.ABELIAN,
            operation_type="increment",
            issue_time_ns=issue_time,
            commit_time_ns=commit_time,
            coordination_rounds=0,
            num_rows=1,
            num_columns=1,
            data_bytes=8,  # int64
        )

        self._collector.record(metric)

    def max_update(
        self,
        table_name: str,
        column: str,
        value: Any,
    ):
        """
        MAX update operation (Semilattice - coordination-free).
        """
        op_id = self._collector.generate_operation_id()
        issue_time = time.perf_counter_ns()
        timestamp = time.time_ns()

        # Simulate local MAX update (instant)
        commit_time = time.perf_counter_ns()

        metric = CommitMetrics(
            table_name=table_name,
            operation_id=op_id,
            timestamp_ns=timestamp,
            algebraic_signature=AlgebraicSignature.SEMILATTICE,
            operation_type="max",
            issue_time_ns=issue_time,
            commit_time_ns=commit_time,
            coordination_rounds=0,
            num_rows=1,
            num_columns=1,
            data_bytes=8,
        )

        self._collector.record(metric)


def run_validation_benchmark(
    writer,  # TableWriter
    num_algebraic: int = 1000,
    num_generic: int = 100,
) -> Dict[str, Any]:
    """
    Run a validation benchmark comparing algebraic vs generic operations.

    Args:
        writer: TableWriter instance
        num_algebraic: Number of algebraic operations to run
        num_generic: Number of generic operations to run

    Returns:
        Benchmark results with timing and validation
    """
    import pandas as pd

    collector = MetricsCollector()
    instrumented = InstrumentedWriter(writer, collector)

    print("Running validation benchmark...")
    print(f"  Algebraic operations: {num_algebraic}")
    print(f"  Generic operations: {num_generic}")

    # Run algebraic operations (increment)
    print("\n[1/2] Running algebraic (increment) operations...")
    for i in range(num_algebraic):
        instrumented.increment("benchmark_counters", "count", delta=1)

    # Run generic operations (overwrite)
    print("[2/2] Running generic (overwrite) operations...")
    for i in range(num_generic):
        df = pd.DataFrame({"id": [i], "value": [f"value_{i}"]})
        try:
            instrumented.write("benchmark_generic", df, operation_type="overwrite")
        except Exception:
            # Record as generic even if write fails (for benchmark purposes)
            op_id = collector.generate_operation_id()
            t = time.perf_counter_ns()
            metric = CommitMetrics(
                table_name="benchmark_generic",
                operation_id=op_id,
                timestamp_ns=time.time_ns(),
                algebraic_signature=AlgebraicSignature.GENERIC,
                operation_type="overwrite",
                issue_time_ns=t,
                commit_time_ns=t + 1000,  # 1 microsecond simulated
                coordination_rounds=0,  # Would be >0 in real distributed system
                num_rows=1,
                num_columns=2,
                data_bytes=100,
            )
            collector.record(metric)

    # Generate summary
    summary = collector.summary()

    print("\n" + "=" * 60)
    print("VALIDATION BENCHMARK RESULTS")
    print("=" * 60)
    print(f"\nTotal operations: {summary['total_operations']}")
    print("\nAlgebraic operations:")
    print(f"  Count: {summary['algebraic_total']['count']}")
    if summary['algebraic_total']['count'] > 0:
        print(f"  Mean latency: {summary['algebraic_total']['mean_latency_ms']:.6f} ms")
        print(f"  Coordination rounds: {summary['algebraic_total']['total_coordination_rounds']}")

    print("\nGeneric operations:")
    print(f"  Count: {summary['generic']['count']}")
    if summary['generic']['count'] > 0:
        print(f"  Mean latency: {summary['generic']['mean_latency_ms']:.6f} ms")
        print(f"  Coordination rounds: {summary['generic']['total_coordination_rounds']}")

    if 'speedup_ratio' in summary:
        print(f"\nSpeedup ratio: {summary['speedup_ratio']:.2f}x")

    print("\nTheory validation:")
    print(f"  All operations match theory: {summary['validation']['all_match_theory']}")
    print(f"  Algebraic ops coordination-free: {summary['validation']['algebraic_coordination_free']}")

    return summary
