monitor.py

"""Cost monitoring and usage tracking system."""

import sqlite3
import json
import csv
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple, Callable
from dataclasses import dataclass, asdict, field
from datetime import datetime, timedelta
from enum import Enum
import threading

from .utils import estimate_cost, get_timestamp


class AlertLevel(Enum):
    """Alert severity levels."""
    INFO = "info"
    WARNING = "warning"
    CRITICAL = "critical"


@dataclass
class UsageRecord:
    """Single API usage record."""
    id: Optional[int]
    timestamp: str
    model: str
    provider: str
    input_tokens: int
    output_tokens: int
    total_tokens: int
    cost: float
    cached: bool = False
    request_type: str = "chat"
    metadata: Dict[str, Any] = field(default_factory=dict)


@dataclass
class UsageSummary:
    """Aggregated usage summary."""
    period: str
    start_date: str
    end_date: str
    total_requests: int
    total_input_tokens: int
    total_output_tokens: int
    total_tokens: int
    total_cost: float
    cached_requests: int
    cost_saved_from_cache: float
    by_model: Dict[str, Dict[str, Any]] = field(default_factory=dict)
    by_provider: Dict[str, Dict[str, Any]] = field(default_factory=dict)


@dataclass
class BudgetConfig:
    """Budget configuration."""
    daily_limit: Optional[float] = None
    weekly_limit: Optional[float] = None
    monthly_limit: Optional[float] = None
    per_request_limit: Optional[float] = None
    alert_threshold: float = 0.8  # Alert at 80% of limit
    hard_cap: bool = False  # If True, block requests over limit


@dataclass
class Alert:
    """Budget alert."""
    level: AlertLevel
    message: str
    timestamp: str
    current_usage: float
    limit: float
    period: str


class CostMonitor:
    """Monitors API costs and usage."""

    def __init__(
        self,
        db_path: str = "usage.db",
        budget: Optional[BudgetConfig] = None,
        alert_callback: Optional[Callable[[Alert], None]] = None
    ):
        """
        Initialize the cost monitor.

        Args:
            db_path: Path to SQLite database
            budget: Budget configuration
            alert_callback: Function to call when alerts are triggered
        """
        self.db_path = Path(db_path)
        self.budget = budget or BudgetConfig()
        self.alert_callback = alert_callback
        self._lock = threading.Lock()
        self._alerts: List[Alert] = []

        self._init_db()

    def _init_db(self):
        """Initialize SQLite database."""
        with sqlite3.connect(self.db_path) as conn:
            conn.execute("""
                CREATE TABLE IF NOT EXISTS usage_records (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    timestamp TEXT NOT NULL,
                    model TEXT NOT NULL,
                    provider TEXT NOT NULL,
                    input_tokens INTEGER NOT NULL,
                    output_tokens INTEGER NOT NULL,
                    total_tokens INTEGER NOT NULL,
                    cost REAL NOT NULL,
                    cached INTEGER DEFAULT 0,
                    request_type TEXT DEFAULT 'chat',
                    metadata TEXT
                )
            """)
            conn.execute("""
                CREATE INDEX IF NOT EXISTS idx_timestamp 
                ON usage_records(timestamp)
            """)
            conn.execute("""
                CREATE INDEX IF NOT EXISTS idx_model 
                ON usage_records(model)
            """)
            conn.execute("""
                CREATE TABLE IF NOT EXISTS alerts (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    level TEXT NOT NULL,
                    message TEXT NOT NULL,
                    timestamp TEXT NOT NULL,
                    current_usage REAL NOT NULL,
                    limit_value REAL NOT NULL,
                    period TEXT NOT NULL
                )
            """)
            conn.commit()

    def record_usage(
        self,
        model: str,
        provider: str,
        input_tokens: int,
        output_tokens: int,
        cost: Optional[float] = None,
        cached: bool = False,
        request_type: str = "chat",
        metadata: Optional[Dict[str, Any]] = None
    ) -> UsageRecord:
        """
        Record API usage.

        Args:
            model: Model name
            provider: Provider name
            input_tokens: Number of input tokens
            output_tokens: Number of output tokens
            cost: Cost (calculated if not provided)
            cached: Whether response was from cache
            request_type: Type of request
            metadata: Additional metadata

        Returns:
            The created usage record
        """
        with self._lock:
            timestamp = get_timestamp()
            total_tokens = input_tokens + output_tokens

            if cost is None:
                cost = estimate_cost(input_tokens, output_tokens, model)

            record = UsageRecord(
                id=None,
                timestamp=timestamp,
                model=model,
                provider=provider,
                input_tokens=input_tokens,
                output_tokens=output_tokens,
                total_tokens=total_tokens,
                cost=cost,
                cached=cached,
                request_type=request_type,
                metadata=metadata or {}
            )

            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.execute("""
                    INSERT INTO usage_records
                    (timestamp, model, provider, input_tokens, output_tokens,
                     total_tokens, cost, cached, request_type, metadata)
                    VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                """, (
                    timestamp, model, provider, input_tokens, output_tokens,
                    total_tokens, cost, int(cached), request_type,
                    json.dumps(metadata or {})
                ))
                record.id = cursor.lastrowid
                conn.commit()

            # Check budget alerts
            self._check_budget_alerts()

            return record

    def _check_budget_alerts(self):
        """Check and trigger budget alerts."""
        now = datetime.utcnow()

        checks = [
            ("daily", self.budget.daily_limit, self.get_daily_cost()),
            ("weekly", self.budget.weekly_limit, self.get_weekly_cost()),
            ("monthly", self.budget.monthly_limit, self.get_monthly_cost()),
        ]

        for period, limit, current in checks:
            if limit is None:
                continue

            ratio = current / limit

            if ratio >= 1.0:
                self._trigger_alert(
                    AlertLevel.CRITICAL,
                    f"{period.capitalize()} budget exceeded: ${current:.4f} / ${limit:.4f}",
                    current, limit, period
                )
            elif ratio >= self.budget.alert_threshold:
                self._trigger_alert(
                    AlertLevel.WARNING,
                    f"{period.capitalize()} budget at {ratio*100:.1f}%: ${current:.4f} / ${limit:.4f}",
                    current, limit, period
                )

    def _trigger_alert(self, level: AlertLevel, message: str, 
                       current: float, limit: float, period: str):
        """Trigger a budget alert."""
        alert = Alert(
            level=level,
            message=message,
            timestamp=get_timestamp(),
            current_usage=current,
            limit=limit,
            period=period
        )

        self._alerts.append(alert)

        # Save to database
        with sqlite3.connect(self.db_path) as conn:
            conn.execute("""
                INSERT INTO alerts (level, message, timestamp, current_usage, limit_value, period)
                VALUES (?, ?, ?, ?, ?, ?)
            """, (level.value, message, alert.timestamp, current, limit, period))
            conn.commit()

        # Call callback if set
        if self.alert_callback:
            try:
                self.alert_callback(alert)
            except Exception:
                pass

    def check_budget(self, estimated_cost: float) -> Tuple[bool, Optional[str]]:
        """
        Check if a request is within budget.

        Args:
            estimated_cost: Estimated cost of the request

        Returns:
            Tuple of (allowed, reason)
        """
        if not self.budget.hard_cap:
            return (True, None)

        # Check per-request limit
        if self.budget.per_request_limit and estimated_cost > self.budget.per_request_limit:
            return (False, f"Request cost ${estimated_cost:.4f} exceeds per-request limit ${self.budget.per_request_limit:.4f}")

        # Check daily limit
        if self.budget.daily_limit:
            daily = self.get_daily_cost()
            if daily + estimated_cost > self.budget.daily_limit:
                return (False, f"Would exceed daily budget: ${daily + estimated_cost:.4f} > ${self.budget.daily_limit:.4f}")

        # Check weekly limit
        if self.budget.weekly_limit:
            weekly = self.get_weekly_cost()
            if weekly + estimated_cost > self.budget.weekly_limit:
                return (False, f"Would exceed weekly budget: ${weekly + estimated_cost:.4f} > ${self.budget.weekly_limit:.4f}")

        # Check monthly limit
        if self.budget.monthly_limit:
            monthly = self.get_monthly_cost()
            if monthly + estimated_cost > self.budget.monthly_limit:
                return (False, f"Would exceed monthly budget: ${monthly + estimated_cost:.4f} > ${self.budget.monthly_limit:.4f}")

        return (True, None)

    def get_daily_cost(self, date: Optional[datetime] = None) -> float:
        """Get total cost for a specific day."""
        date = date or datetime.utcnow()
        start = date.replace(hour=0, minute=0, second=0, microsecond=0)
        end = start + timedelta(days=1)
        return self._get_cost_in_range(start, end)

    def get_weekly_cost(self, date: Optional[datetime] = None) -> float:
        """Get total cost for the week containing the date."""
        date = date or datetime.utcnow()
        start = date - timedelta(days=date.weekday())
        start = start.replace(hour=0, minute=0, second=0, microsecond=0)
        end = start + timedelta(days=7)
        return self._get_cost_in_range(start, end)

    def get_monthly_cost(self, date: Optional[datetime] = None) -> float:
        """Get total cost for the month containing the date."""
        date = date or datetime.utcnow()
        start = date.replace(day=1, hour=0, minute=0, second=0, microsecond=0)
        if date.month == 12:
            end = start.replace(year=date.year + 1, month=1)
        else:
            end = start.replace(month=date.month + 1)
        return self._get_cost_in_range(start, end)

    def _get_cost_in_range(self, start: datetime, end: datetime) -> float:
        """Get total cost in a date range."""
        with sqlite3.connect(self.db_path) as conn:
            result = conn.execute("""
                SELECT COALESCE(SUM(cost), 0)
                FROM usage_records
                WHERE timestamp >= ? AND timestamp < ?
            """, (start.isoformat(), end.isoformat())).fetchone()
            return result[0] if result else 0.0

    def get_summary(
        self,
        start_date: Optional[datetime] = None,
        end_date: Optional[datetime] = None,
        period: str = "custom"
    ) -> UsageSummary:
        """
        Get usage summary for a period.

        Args:
            start_date: Start of period (defaults to 30 days ago)
            end_date: End of period (defaults to now)
            period: Period label

        Returns:
            UsageSummary with aggregated data
        """
        end_date = end_date or datetime.utcnow()
        start_date = start_date or (end_date - timedelta(days=30))

        with sqlite3.connect(self.db_path) as conn:
            # Overall totals
            totals = conn.execute("""
                SELECT 
                    COUNT(*) as total_requests,
                    COALESCE(SUM(input_tokens), 0) as total_input,
                    COALESCE(SUM(output_tokens), 0) as total_output,
                    COALESCE(SUM(total_tokens), 0) as total_tokens,
                    COALESCE(SUM(cost), 0) as total_cost,
                    COALESCE(SUM(CASE WHEN cached = 1 THEN 1 ELSE 0 END), 0) as cached_requests
                FROM usage_records
                WHERE timestamp >= ? AND timestamp < ?
            """, (start_date.isoformat(), end_date.isoformat())).fetchone()

            # By model
            by_model = {}
            model_rows = conn.execute("""
                SELECT 
                    model,
                    COUNT(*) as requests,
                    SUM(total_tokens) as tokens,
                    SUM(cost) as cost
                FROM usage_records
                WHERE timestamp >= ? AND timestamp < ?
                GROUP BY model
            """, (start_date.isoformat(), end_date.isoformat())).fetchall()

            for row in model_rows:
                by_model[row[0]] = {
                    "requests": row[1],
                    "tokens": row[2],
                    "cost": row[3]
                }

            # By provider
            by_provider = {}
            provider_rows = conn.execute("""
                SELECT 
                    provider,
                    COUNT(*) as requests,
                    SUM(total_tokens) as tokens,
                    SUM(cost) as cost
                FROM usage_records
                WHERE timestamp >= ? AND timestamp < ?
                GROUP BY provider
            """, (start_date.isoformat(), end_date.isoformat())).fetchall()

            for row in provider_rows:
                by_provider[row[0]] = {
                    "requests": row[1],
                    "tokens": row[2],
                    "cost": row[3]
                }

        return UsageSummary(
            period=period,
            start_date=start_date.isoformat(),
            end_date=end_date.isoformat(),
            total_requests=totals[0],
            total_input_tokens=totals[1],
            total_output_tokens=totals[2],
            total_tokens=totals[3],
            total_cost=totals[4],
            cached_requests=totals[5],
            cost_saved_from_cache=0.0,  # Would need cache integration
            by_model=by_model,
            by_provider=by_provider
        )

    def get_daily_summary(self, date: Optional[datetime] = None) -> UsageSummary:
        """Get summary for a specific day."""
        date = date or datetime.utcnow()
        start = date.replace(hour=0, minute=0, second=0, microsecond=0)
        end = start + timedelta(days=1)
        return self.get_summary(start, end, "daily")

    def get_weekly_summary(self, date: Optional[datetime] = None) -> UsageSummary:
        """Get summary for the week containing the date."""
        date = date or datetime.utcnow()
        start = date - timedelta(days=date.weekday())
        start = start.replace(hour=0, minute=0, second=0, microsecond=0)
        end = start + timedelta(days=7)
        return self.get_summary(start, end, "weekly")

    def get_monthly_summary(self, date: Optional[datetime] = None) -> UsageSummary:
        """Get summary for the month containing the date."""
        date = date or datetime.utcnow()
        start = date.replace(day=1, hour=0, minute=0, second=0, microsecond=0)
        if date.month == 12:
            end = start.replace(year=date.year + 1, month=1)
        else:
            end = start.replace(month=date.month + 1)
        return self.get_summary(start, end, "monthly")

    def export_json(self, filepath: str, start_date: Optional[datetime] = None,
                    end_date: Optional[datetime] = None) -> str:
        """Export usage data to JSON file."""
        end_date = end_date or datetime.utcnow()
        start_date = start_date or (end_date - timedelta(days=30))

        with sqlite3.connect(self.db_path) as conn:
            conn.row_factory = sqlite3.Row
            rows = conn.execute("""
                SELECT * FROM usage_records
                WHERE timestamp >= ? AND timestamp < ?
                ORDER BY timestamp
            """, (start_date.isoformat(), end_date.isoformat())).fetchall()

            records = [dict(row) for row in rows]
            for record in records:
                record["metadata"] = json.loads(record.get("metadata", "{}"))

        summary = self.get_summary(start_date, end_date)

        export_data = {
            "export_date": get_timestamp(),
            "period": {
                "start": start_date.isoformat(),
                "end": end_date.isoformat()
            },
            "summary": asdict(summary),
            "records": records
        }

        with open(filepath, "w") as f:
            json.dump(export_data, f, indent=2)

        return filepath

    def export_csv(self, filepath: str, start_date: Optional[datetime] = None,
                   end_date: Optional[datetime] = None) -> str:
        """Export usage data to CSV file."""
        end_date = end_date or datetime.utcnow()
        start_date = start_date or (end_date - timedelta(days=30))

        with sqlite3.connect(self.db_path) as conn:
            conn.row_factory = sqlite3.Row
            rows = conn.execute("""
                SELECT id, timestamp, model, provider, input_tokens, output_tokens,
                       total_tokens, cost, cached, request_type
                FROM usage_records
                WHERE timestamp >= ? AND timestamp < ?
                ORDER BY timestamp
            """, (start_date.isoformat(), end_date.isoformat())).fetchall()

        with open(filepath, "w", newline="") as f:
            writer = csv.writer(f)
            writer.writerow([
                "id", "timestamp", "model", "provider", "input_tokens",
                "output_tokens", "total_tokens", "cost", "cached", "request_type"
            ])
            for row in rows:
                writer.writerow(list(row))

        return filepath

    def get_recent_alerts(self, limit: int = 10) -> List[Alert]:
        """Get recent alerts."""
        with sqlite3.connect(self.db_path) as conn:
            rows = conn.execute("""
                SELECT level, message, timestamp, current_usage, limit_value, period
                FROM alerts
                ORDER BY timestamp DESC
                LIMIT ?
            """, (limit,)).fetchall()

            return [
                Alert(
                    level=AlertLevel(row[0]),
                    message=row[1],
                    timestamp=row[2],
                    current_usage=row[3],
                    limit=row[4],
                    period=row[5]
                )
                for row in rows
            ]

    def set_budget(self, budget: BudgetConfig):
        """Update budget configuration."""
        self.budget = budget

    def clear_history(self, before_date: Optional[datetime] = None):
        """Clear usage history."""
        with sqlite3.connect(self.db_path) as conn:
            if before_date:
                conn.execute(
                    "DELETE FROM usage_records WHERE timestamp < ?",
                    (before_date.isoformat(),)
                )
            else:
                conn.execute("DELETE FROM usage_records")
            conn.commit()