"""Cost Tracking for Empathy Framework

Tracks API costs across model tiers and calculates savings from
smart model routing (Haiku/Sonnet/Opus selection).

Features:
- Log each API request with model, tokens, and task type
- Calculate actual cost vs baseline (if all requests used premium model)
- Generate weekly/monthly reports
- Integrate with `empathy costs` and `empathy morning` commands
- **Performance optimized**: Batch writes (50 requests) + JSONL format

Model pricing is sourced from attune.models.MODEL_REGISTRY.

Copyright 2025 Smart-AI-Memory
Licensed under Fair Source License 0.9
"""

import atexit
import heapq
import json
from datetime import datetime, timedelta
from pathlib import Path
from typing import Any

# Import pricing from unified registry
from attune.config import _validate_file_path
from attune.models import MODEL_REGISTRY
from attune.models.registry import TIER_PRICING


def _build_model_pricing() -> dict[str, dict[str, float]]:
    """Build MODEL_PRICING from unified registry."""
    pricing: dict[str, dict[str, float]] = {}

    # Add all models from registry
    for provider_models in MODEL_REGISTRY.values():
        for model_info in provider_models.values():
            pricing[model_info.id] = {
                "input": model_info.input_cost_per_million,
                "output": model_info.output_cost_per_million,
            }

    # Add tier aliases from registry
    pricing.update(TIER_PRICING)

    # Add legacy model names for backward compatibility
    legacy_models = {
        "claude-3-haiku-20240307": {"input": 0.25, "output": 1.25},
        "claude-3-5-sonnet-20241022": {"input": 3.00, "output": 15.00},
        "claude-opus-4-20250514": {"input": 15.00, "output": 75.00},
        "gpt-4-turbo": {"input": 10.00, "output": 30.00},
    }
    pricing.update(legacy_models)

    return pricing


# Pricing per million tokens - sourced from unified registry
MODEL_PRICING = _build_model_pricing()

# Default premium model for baseline comparison
BASELINE_MODEL = "claude-opus-4-6"


class CostTracker:
    """Tracks API costs and calculates savings from model routing.

    **Performance Optimized:**
    - Batch writes (flush every 50 requests)
    - JSONL append-only format for new data
    - Backward compatible with JSON format
    - Zero data loss (atexit handler)
    - Lazy loading: Full request history only loaded when accessed
    - Separate summary file: Fast init (80-90% faster for large histories)

    Usage:
        tracker = CostTracker()
        tracker.log_request("claude-3-haiku-20240307", 1000, 500, "summarize")
        report = tracker.get_report()
    """

    @property
    def requests(self) -> list[dict]:
        """Access request history (lazy-loaded for performance).

        Returns:
            List of request records. Triggers lazy loading on first access.
        """
        self._load_requests()
        return self.data.get("requests", [])

    def __init__(self, storage_dir: str = ".empathy", batch_size: int = 50):
        """Initialize cost tracker.

        Args:
            storage_dir: Directory for cost data storage
            batch_size: Number of requests to buffer before flushing (default: 50)

        Performance optimizations:
            - Lazy loading: Only load summary data on init, defer full request history
            - Separate summary file: Fast access to daily_totals without parsing JSONL
            - Init time reduced by 80-90% for large history files
        """
        self.storage_dir = Path(storage_dir)
        self.storage_dir.mkdir(parents=True, exist_ok=True)
        self.costs_file = self.storage_dir / "costs.json"
        self.costs_jsonl = self.storage_dir / "costs.jsonl"
        self.costs_summary = self.storage_dir / "costs_summary.json"
        self.batch_size = batch_size
        self._buffer: list[dict] = []  # Buffered requests not yet flushed
        self._requests_loaded = False  # Track if full history is loaded
        self._load_summary()  # Only load summary on init (fast)

        # Register cleanup handler to flush on exit
        atexit.register(self._cleanup)

    def _cleanup(self) -> None:
        """Cleanup handler - flush buffer on exit."""
        try:
            if self._buffer:
                self.flush()
        except Exception:  # noqa: BLE001
            # INTENTIONAL: Best-effort flush, don't break shutdown
            pass

    def _load_summary(self) -> None:
        """Load only summary data on init (fast path).

        This loads daily_totals from the summary file without parsing
        the full request history. Full history is lazy-loaded only when needed.

        Performance: 80-90% faster init for large history files.
        """
        # Initialize with default structure (no requests loaded yet)
        self.data = self._default_data()
        self.data["requests"] = []  # Start empty, lazy-load later

        # Try loading pre-computed summary first (fastest)
        if self.costs_summary.exists():
            try:
                with open(self.costs_summary) as f:
                    summary_data = json.load(f)
                    self.data["daily_totals"] = summary_data.get("daily_totals", {})
                    self.data["created_at"] = summary_data.get(
                        "created_at", self.data["created_at"]
                    )
                    self.data["last_updated"] = summary_data.get(
                        "last_updated", self.data["last_updated"]
                    )
                    return  # Summary loaded, done
            except (OSError, json.JSONDecodeError):
                pass  # Fall through to JSON fallback

        # Fallback: Load daily_totals from costs.json (backward compatibility)
        if self.costs_file.exists():
            try:
                with open(self.costs_file) as f:
                    json_data = json.load(f)
                    self.data["daily_totals"] = json_data.get("daily_totals", {})
                    self.data["created_at"] = json_data.get("created_at", self.data["created_at"])
                    self.data["last_updated"] = json_data.get(
                        "last_updated", self.data["last_updated"]
                    )
                    # Don't load requests here - they'll be lazy-loaded
            except (OSError, json.JSONDecodeError):
                pass  # Use defaults

    def _load_requests(self) -> None:
        """Lazy-load full request history (only when needed).

        Called automatically when request history is accessed.
        Most operations use daily_totals and don't need this.
        """
        if self._requests_loaded:
            return  # Already loaded

        # Load from JSON first
        if self.costs_file.exists():
            try:
                with open(self.costs_file) as f:
                    json_data = json.load(f)
                    self.data["requests"] = json_data.get("requests", [])
            except (OSError, json.JSONDecodeError):
                self.data["requests"] = []

        # Append from JSONL if it exists
        if self.costs_jsonl.exists():
            try:
                with open(self.costs_jsonl) as f:
                    for line in f:
                        if line.strip():
                            request = json.loads(line)
                            self.data["requests"].append(request)
            except (OSError, json.JSONDecodeError):
                pass  # Ignore errors, use what we have

        self._requests_loaded = True

    def _load(self) -> None:
        """Load cost data from storage (supports both JSON and JSONL).

        Deprecated: Use _load_summary() for fast init, _load_requests() for full history.
        Kept for backward compatibility.
        """
        self._load_summary()
        self._load_requests()

    def _default_data(self) -> dict:
        """Return default data structure."""
        return {
            "requests": [],
            "daily_totals": {},
            "created_at": datetime.now().isoformat(),
            "last_updated": datetime.now().isoformat(),
        }

    def _update_daily_totals(self, request: dict) -> None:
        """Update daily totals from a request.

        Args:
            request: Request record with cost information

        """
        timestamp = request.get("timestamp", datetime.now().isoformat())
        date = timestamp[:10]  # Extract YYYY-MM-DD

        if date not in self.data["daily_totals"]:
            self.data["daily_totals"][date] = {
                "requests": 0,
                "input_tokens": 0,
                "output_tokens": 0,
                "actual_cost": 0,
                "baseline_cost": 0,
                "savings": 0,
            }

        daily = self.data["daily_totals"][date]
        daily["requests"] += 1
        daily["input_tokens"] += request["input_tokens"]
        daily["output_tokens"] += request["output_tokens"]
        daily["actual_cost"] = round(daily["actual_cost"] + request["actual_cost"], 6)
        daily["baseline_cost"] = round(daily["baseline_cost"] + request["baseline_cost"], 6)
        daily["savings"] = round(daily["savings"] + request["savings"], 6)

    def _save(self) -> None:
        """Save cost data to storage (legacy JSON format).

        **Note:** This is only used for backward compatibility.
        New data is written to JSONL format via flush().
        """
        try:
            self.data["last_updated"] = datetime.now().isoformat()
            validated_path = _validate_file_path(str(self.costs_file))
            with open(validated_path, "w") as f:
                json.dump(self.data, f, indent=2)
        except (OSError, ValueError):
            pass  # Best effort - cost tracking shouldn't crash the app

    def _save_summary(self) -> None:
        """Save summary data (daily_totals) to separate file for fast loading.

        This enables 80-90% faster init by avoiding full JSONL parsing on startup.
        """
        summary_data = {
            "daily_totals": self.data.get("daily_totals", {}),
            "created_at": self.data.get("created_at", datetime.now().isoformat()),
            "last_updated": datetime.now().isoformat(),
        }
        try:
            validated_path = _validate_file_path(str(self.costs_summary))
            with open(validated_path, "w") as f:
                json.dump(summary_data, f, indent=2)
        except (OSError, ValueError):
            pass  # Best effort - summary is an optimization, not critical

    def flush(self) -> None:
        """Flush buffered requests to disk (JSONL format).

        This is called automatically:
        - Every `batch_size` requests
        - On process exit (atexit handler)
        - Manually by calling tracker.flush()
        """
        if not self._buffer:
            return

        # Append buffered requests to JSONL file
        try:
            validated_path = _validate_file_path(str(self.costs_jsonl))
            with open(validated_path, "a") as f:
                for request in self._buffer:
                    f.write(json.dumps(request) + "\n")

            # Update daily totals (always in memory)
            for request in self._buffer:
                self._update_daily_totals(request)

            # Load requests if needed before extending (maintains backward compat)
            # This defers the expensive load from init to first flush
            if not self._requests_loaded:
                self._load_requests()

            self.data["requests"].extend(self._buffer)
            # Keep only last 1000 requests in memory
            if len(self.data["requests"]) > 1000:
                self.data["requests"] = self.data["requests"][-1000:]

            # Save summary file (fast path for future loads)
            self._save_summary()

            # Update JSON file periodically (every 10 flushes = 500 requests)
            # This maintains backward compatibility without killing performance
            if len(self._buffer) >= 500 or not self.costs_jsonl.exists():
                self._save()

            # Clear buffer
            self._buffer.clear()

        except OSError:
            # If JSONL write fails, fallback to immediate JSON save
            for request in self._buffer:
                self._update_daily_totals(request)
            if not self._requests_loaded:
                self._load_requests()
            self.data["requests"].extend(self._buffer)
            self._buffer.clear()
            self._save()
            self._save_summary()

    def _calculate_cost(self, model: str, input_tokens: int, output_tokens: int) -> float:
        """Calculate cost for a request.

        Args:
            model: Model name or tier
            input_tokens: Number of input tokens
            output_tokens: Number of output tokens

        Returns:
            Cost in USD

        """
        pricing = MODEL_PRICING.get(model) or MODEL_PRICING["capable"]
        input_cost = (input_tokens / 1_000_000) * pricing["input"]
        output_cost = (output_tokens / 1_000_000) * pricing["output"]
        return input_cost + output_cost

    def log_request(
        self,
        model: str,
        input_tokens: int,
        output_tokens: int,
        task_type: str = "unknown",
        tier: str | None = None,
    ) -> dict:
        """Log an API request with cost tracking (batched writes).

        **Performance optimized**: Requests are buffered and flushed every
        `batch_size` requests (default: 50) instead of writing to disk
        immediately. This provides a 60x+ performance improvement.

        Args:
            model: Model name used
            input_tokens: Number of input tokens
            output_tokens: Number of output tokens
            task_type: Type of task (summarize, generate_code, etc.)
            tier: Optional tier override (cheap, capable, premium)

        Returns:
            Request record with cost information

        """
        actual_cost = self._calculate_cost(model, input_tokens, output_tokens)
        baseline_cost = self._calculate_cost(BASELINE_MODEL, input_tokens, output_tokens)
        savings = baseline_cost - actual_cost

        request = {
            "timestamp": datetime.now().isoformat(),
            "model": model,
            "tier": tier or self._get_tier(model),
            "task_type": task_type,
            "input_tokens": input_tokens,
            "output_tokens": output_tokens,
            "actual_cost": round(actual_cost, 6),
            "baseline_cost": round(baseline_cost, 6),
            "savings": round(savings, 6),
        }

        # Add to buffer instead of immediate save
        self._buffer.append(request)

        # Flush when buffer reaches batch size
        if len(self._buffer) >= self.batch_size:
            self.flush()

        return request

    def _get_tier(self, model: str) -> str:
        """Determine tier from model name."""
        if "haiku" in model.lower():
            return "cheap"
        if "opus" in model.lower():
            return "premium"
        return "capable"

    def get_summary(self, days: int = 7, include_breakdown: bool = True) -> dict:
        """Get cost summary for recent period (includes buffered requests).

        **Real-time data**: Includes buffered requests that haven't been
        flushed to disk yet, ensuring accurate real-time reporting.

        **Performance optimized**: Main totals computed from pre-aggregated
        daily_totals. Full request history only loaded if include_breakdown=True.

        Args:
            days: Number of days to include
            include_breakdown: If True, include by_tier and by_task breakdown
                (requires loading full request history). Default: True.

        Returns:
            Summary with totals and savings percentage

        """
        cutoff = datetime.now() - timedelta(days=days)
        cutoff_str = cutoff.strftime("%Y-%m-%d")

        totals: dict[str, Any] = {
            "days": days,
            "requests": 0,
            "input_tokens": 0,
            "output_tokens": 0,
            "actual_cost": 0,
            "baseline_cost": 0,
            "savings": 0,
            "by_tier": {"cheap": 0, "capable": 0, "premium": 0},
            "by_task": {},
        }

        # Include daily totals from flushed data (fast - always in memory)
        for date, daily in self.data.get("daily_totals", {}).items():
            if date >= cutoff_str:
                totals["requests"] += daily["requests"]
                totals["input_tokens"] += daily["input_tokens"]
                totals["output_tokens"] += daily["output_tokens"]
                totals["actual_cost"] += daily["actual_cost"]
                totals["baseline_cost"] += daily["baseline_cost"]
                totals["savings"] += daily["savings"]

        # Add buffered request costs to totals (always in memory)
        cutoff_iso = cutoff.isoformat()
        for req in self._buffer:
            if req["timestamp"] >= cutoff_iso:
                totals["requests"] += 1
                totals["input_tokens"] += req["input_tokens"]
                totals["output_tokens"] += req["output_tokens"]
                totals["actual_cost"] += req["actual_cost"]
                totals["baseline_cost"] += req["baseline_cost"]
                totals["savings"] += req["savings"]

        # Include breakdown by tier/task (requires loading full history)
        if include_breakdown:
            # Lazy-load full request history only when needed
            self._load_requests()
            all_requests = list(self.data.get("requests", [])) + self._buffer

            for req in all_requests:
                if req["timestamp"] >= cutoff_iso:
                    tier = req.get("tier", "capable")
                    task = req.get("task_type", "unknown")
                    totals["by_tier"][tier] = totals["by_tier"].get(tier, 0) + 1
                    totals["by_task"][task] = totals["by_task"].get(task, 0) + 1

        # Calculate savings percentage
        if totals["baseline_cost"] > 0:
            totals["savings_percent"] = round(
                (totals["savings"] / totals["baseline_cost"]) * 100,
                1,
            )
        else:
            totals["savings_percent"] = 0

        return totals

    def get_report(self, days: int = 7) -> str:
        """Generate a formatted cost report.

        Args:
            days: Number of days to include

        Returns:
            Formatted report string

        """
        summary = self.get_summary(days)

        lines = [
            "",
            "=" * 60,
            "  COST TRACKING REPORT",
            f"  Last {days} days",
            "=" * 60,
            "",
            "SUMMARY",
            "-" * 40,
            f"  Total requests:      {summary['requests']:,}",
            f"  Input tokens:        {summary['input_tokens']:,}",
            f"  Output tokens:       {summary['output_tokens']:,}",
            "",
            "COSTS",
            "-" * 40,
            f"  Actual cost:         ${summary['actual_cost']:.4f}",
            f"  Baseline (Opus):     ${summary['baseline_cost']:.4f}",
            f"  You saved:           ${summary['savings']:.4f} ({summary['savings_percent']}%)",
            "",
        ]

        # Tier breakdown
        if sum(summary["by_tier"].values()) > 0:
            lines.extend(
                [
                    "BY MODEL TIER",
                    "-" * 40,
                ],
            )
            for tier, count in sorted(summary["by_tier"].items(), key=lambda x: -x[1]):
                if count > 0:
                    pct = (count / summary["requests"]) * 100 if summary["requests"] else 0
                    lines.append(f"  {tier:12} {count:6,} requests ({pct:.1f}%)")
            lines.append("")

        # Task breakdown (top 5)
        if summary["by_task"]:
            lines.extend(
                [
                    "BY TASK TYPE (Top 5)",
                    "-" * 40,
                ],
            )
            top_tasks = heapq.nlargest(5, summary["by_task"].items(), key=lambda x: x[1])
            for task, count in top_tasks:
                lines.append(f"  {task:20} {count:,}")
            lines.append("")

        lines.extend(
            [
                "=" * 60,
                "  Model routing saves money by using cheaper models",
                "  for simple tasks and Opus only when needed.",
                "=" * 60,
                "",
            ],
        )

        return "\n".join(lines)

    def get_today(self) -> dict[str, int | float]:
        """Get today's cost summary (includes buffered requests)."""
        today = datetime.now().strftime("%Y-%m-%d")
        daily_totals = self.data.get("daily_totals", {})
        default: dict[str, int | float] = {
            "requests": 0,
            "input_tokens": 0,
            "output_tokens": 0,
            "actual_cost": 0,
            "baseline_cost": 0,
            "savings": 0,
        }

        # Start with flushed daily totals
        if isinstance(daily_totals, dict) and today in daily_totals:
            result = daily_totals[today]
            totals = result.copy() if isinstance(result, dict) else default.copy()
        else:
            totals = default.copy()

        # Add buffered requests for today (real-time data)
        for req in self._buffer:
            req_date = req["timestamp"][:10]  # Extract YYYY-MM-DD
            if req_date == today:
                totals["requests"] += 1
                totals["input_tokens"] += req["input_tokens"]
                totals["output_tokens"] += req["output_tokens"]
                totals["actual_cost"] = round(totals["actual_cost"] + req["actual_cost"], 6)
                totals["baseline_cost"] = round(totals["baseline_cost"] + req["baseline_cost"], 6)
                totals["savings"] = round(totals["savings"] + req["savings"], 6)

        return totals


def cmd_costs(args):
    """CLI command handler for costs."""
    tracker = CostTracker(storage_dir=getattr(args, "empathy_dir", ".empathy"))
    days = getattr(args, "days", 7)

    if getattr(args, "json", False):
        import json as json_mod

        print(json_mod.dumps(tracker.get_summary(days), indent=2))
    else:
        print(tracker.get_report(days))

    return 0


# Singleton for global tracking
_tracker: CostTracker | None = None


def get_tracker(storage_dir: str = ".empathy") -> CostTracker:
    """Get or create the global cost tracker."""
    global _tracker
    if _tracker is None:
        _tracker = CostTracker(storage_dir)
    return _tracker


def log_request(
    model: str,
    input_tokens: int,
    output_tokens: int,
    task_type: str = "unknown",
    tier: str | None = None,
) -> dict:
    """Convenience function to log a request to the global tracker.

    Usage:
        from attune.cost_tracker import log_request
        log_request("claude-3-haiku-20240307", 1000, 500, "summarize")
    """
    return get_tracker().log_request(model, input_tokens, output_tokens, task_type, tier)
