"""Performance monitoring module.

This module provides decorators and utilities for tracking performance metrics
across cleave operations. Metrics are stored in ~/.cleave/performance.yaml.

Key features:
- @timed decorator for execution time tracking
- @memory_profiled decorator for memory usage tracking
- Manual metric recording via record_metric()
- Persistence to YAML for analysis
"""

from __future__ import annotations

import logging
import time
import tracemalloc
from functools import wraps
from pathlib import Path
from typing import Any, Callable, TypeVar

from cleave.core.file_utils import atomic_write_text
from cleave.core.yaml_utils import parse_yaml_simple, to_yaml

logger = logging.getLogger(__name__)

# Type variable for generic decorator
F = TypeVar("F", bound=Callable[..., Any])

# Default storage path
DEFAULT_PERFORMANCE_FILE = Path.home() / ".cleave" / "performance.yaml"

# In-memory cache for performance metrics
_metrics_cache: dict[str, Any] = {}
_cache_loaded = False


# =============================================================================
# METRICS STORAGE
# =============================================================================


def _get_storage_path(storage_path: Path | None = None) -> Path:
    """Get the storage path for performance metrics.

    Args:
        storage_path: Optional override path

    Returns:
        Path to performance metrics file
    """
    path = storage_path or DEFAULT_PERFORMANCE_FILE
    path.parent.mkdir(parents=True, exist_ok=True)
    return path


def _load_metrics(storage_path: Path | None = None) -> dict:
    """Load metrics from storage.

    Args:
        storage_path: Optional override path

    Returns:
        Dict of metrics data
    """
    global _metrics_cache, _cache_loaded

    if _cache_loaded and storage_path is None:
        return _metrics_cache

    path = _get_storage_path(storage_path)

    if path.exists():
        try:
            metrics = parse_yaml_simple(path.read_text())
            if storage_path is None:
                _metrics_cache = metrics
                _cache_loaded = True
            return metrics
        except Exception as e:
            logger.warning(f"Failed to load metrics from {path}: {e}")
            return {}
    else:
        return {}


def _save_metrics(metrics: dict, storage_path: Path | None = None) -> None:
    """Save metrics to storage.

    Args:
        metrics: Dict of metrics data
        storage_path: Optional override path
    """
    global _metrics_cache, _cache_loaded

    path = _get_storage_path(storage_path)

    try:
        atomic_write_text(path, to_yaml(metrics))
        if storage_path is None:
            _metrics_cache = metrics
            _cache_loaded = True
    except Exception as e:
        logger.error(f"Failed to save metrics to {path}: {e}")


def get_metrics(storage_path: Path | None = None) -> dict:
    """Get all performance metrics.

    Args:
        storage_path: Optional override path

    Returns:
        Dict of all metrics
    """
    return _load_metrics(storage_path)


def clear_metrics(storage_path: Path | None = None) -> None:
    """Clear all performance metrics.

    Args:
        storage_path: Optional override path
    """
    global _metrics_cache, _cache_loaded

    _save_metrics({}, storage_path)
    _metrics_cache = {}
    _cache_loaded = False


def record_metric(
    name: str,
    data: dict[str, Any],
    storage_path: Path | None = None,
) -> None:
    """Record a performance metric.

    Args:
        name: Metric name
        data: Metric data (dict of key-value pairs)
        storage_path: Optional override path
    """
    metrics = _load_metrics(storage_path)

    if name not in metrics:
        metrics[name] = {}

    # Merge data into existing metric
    metrics[name].update(data)

    _save_metrics(metrics, storage_path)


# =============================================================================
# DECORATORS
# =============================================================================


def timed(operation_name: str, storage_path: Path | None = None) -> Callable[[F], F]:
    """Decorator to track function execution time.

    Usage:
        @timed("my_operation")
        def slow_function():
            time.sleep(1)

    Args:
        operation_name: Name of the operation being timed
        storage_path: Optional override path for metrics storage

    Returns:
        Decorated function
    """

    def decorator(func: F) -> F:
        @wraps(func)
        def wrapper(*args, **kwargs):
            start_time = time.perf_counter()
            try:
                result = func(*args, **kwargs)
                return result
            finally:
                elapsed = time.perf_counter() - start_time

                # Update metrics
                metrics = _load_metrics(storage_path)

                if operation_name not in metrics:
                    metrics[operation_name] = {
                        "count": 0,
                        "total_time": 0.0,
                        "avg_time": 0.0,
                        "min_time": float("inf"),
                        "max_time": 0.0,
                    }

                op_metrics = metrics[operation_name]
                op_metrics["count"] += 1
                op_metrics["total_time"] += elapsed
                op_metrics["avg_time"] = op_metrics["total_time"] / op_metrics["count"]
                op_metrics["min_time"] = min(op_metrics["min_time"], elapsed)
                op_metrics["max_time"] = max(op_metrics["max_time"], elapsed)
                op_metrics["last_time"] = elapsed

                _save_metrics(metrics, storage_path)

        return wrapper  # type: ignore

    return decorator


def memory_profiled(
    operation_name: str,
    storage_path: Path | None = None,
) -> Callable[[F], F]:
    """Decorator to track memory usage.

    Usage:
        @memory_profiled("my_operation")
        def allocate_memory():
            data = [0] * 1000000

    Args:
        operation_name: Name of the operation being profiled
        storage_path: Optional override path for metrics storage

    Returns:
        Decorated function
    """

    def decorator(func: F) -> F:
        @wraps(func)
        def wrapper(*args, **kwargs):
            tracemalloc.start()
            try:
                result = func(*args, **kwargs)
                return result
            finally:
                current, peak = tracemalloc.get_traced_memory()
                tracemalloc.stop()

                # Update metrics
                metrics = _load_metrics(storage_path)

                if operation_name not in metrics:
                    metrics[operation_name] = {
                        "count": 0,
                        "peak_memory_mb": 0.0,
                        "avg_memory_mb": 0.0,
                        "total_memory_mb": 0.0,
                    }

                op_metrics = metrics[operation_name]
                peak_mb = peak / (1024 * 1024)

                op_metrics["count"] += 1
                op_metrics["total_memory_mb"] += peak_mb
                op_metrics["avg_memory_mb"] = (
                    op_metrics["total_memory_mb"] / op_metrics["count"]
                )
                op_metrics["peak_memory_mb"] = max(
                    op_metrics.get("peak_memory_mb", 0.0), peak_mb
                )
                op_metrics["last_memory_mb"] = peak_mb

                _save_metrics(metrics, storage_path)

        return wrapper  # type: ignore

    return decorator


# =============================================================================
# INSTRUMENTATION HELPERS
# =============================================================================


def record_workspace_metrics(workspace_path: Path) -> None:
    """Record metrics about a workspace.

    Args:
        workspace_path: Path to workspace directory
    """
    if not workspace_path.exists():
        logger.warning(f"Workspace not found: {workspace_path}")
        return

    total_bytes = 0
    file_count = 0

    for file_path in workspace_path.rglob("*"):
        if file_path.is_file():
            file_count += 1
            total_bytes += file_path.stat().st_size

    record_metric("workspace_size", {
        "total_bytes": total_bytes,
        "file_count": file_count,
        "total_mb": total_bytes / (1024 * 1024),
    })


def record_llm_metrics(
    first_token_latency: float,
    total_time: float,
    tokens: int,
) -> None:
    """Record LLM response metrics.

    Args:
        first_token_latency: Time to first token (seconds)
        total_time: Total response time (seconds)
        tokens: Number of tokens generated
    """
    tokens_per_sec = tokens / total_time if total_time > 0 else 0

    record_metric("llm_response", {
        "first_token_latency": first_token_latency,
        "total_time": total_time,
        "tokens": tokens,
        "tokens_per_sec": tokens_per_sec,
    })


# =============================================================================
# CONTEXT MANAGERS
# =============================================================================


class TimeContext:
    """Context manager for timing a block of code.

    Usage:
        with TimeContext("operation_name") as timer:
            # do work
            pass
        print(f"Elapsed: {timer.elapsed}")
    """

    def __init__(self, operation_name: str, storage_path: Path | None = None):
        self.operation_name = operation_name
        self.storage_path = storage_path
        self.start_time = 0.0
        self.elapsed = 0.0

    def __enter__(self) -> TimeContext:
        self.start_time = time.perf_counter()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.elapsed = time.perf_counter() - self.start_time

        # Record metric
        metrics = _load_metrics(self.storage_path)

        if self.operation_name not in metrics:
            metrics[self.operation_name] = {
                "count": 0,
                "total_time": 0.0,
                "avg_time": 0.0,
            }

        op_metrics = metrics[self.operation_name]
        op_metrics["count"] += 1
        op_metrics["total_time"] += self.elapsed
        op_metrics["avg_time"] = op_metrics["total_time"] / op_metrics["count"]

        _save_metrics(metrics, self.storage_path)

        return False  # Don't suppress exceptions
