"""Cost tracking module for OCLAWMA.

This module provides comprehensive cost tracking for jobs including:
- API costs (LLM calls, storage, compute)
- Compute time tracking
- Storage usage tracking
- Per-project and per-tenant budget management
- Cost attribution reports (daily/weekly/monthly)
- Export to CSV/JSON

Example:
    >>> from oclawma.costs import CostTracker, JobCost
    >>> tracker = CostTracker()
    >>> tracker.start_job("job-123", project_id="proj-1")
    >>> tracker.record_api_cost("job-123", "openai", "gpt-4", 0.05)
    >>> tracker.end_job("job-123")
    >>> cost = tracker.get_job_cost("job-123")
"""

from __future__ import annotations

import csv
import json
import sqlite3
import threading
import time
from collections.abc import Iterator
from contextlib import contextmanager, suppress
from dataclasses import dataclass, field
from datetime import datetime, timedelta
from enum import Enum
from pathlib import Path
from typing import Any, Callable

from pydantic import BaseModel, Field, field_validator


class AlertLevel(str, Enum):
    """Budget alert levels."""

    INFO = "info"
    WARNING = "warning"
    CRITICAL = "critical"


class ReportPeriod(str, Enum):
    """Report period options."""

    DAILY = "daily"
    WEEKLY = "weekly"
    MONTHLY = "monthly"


class BudgetExceededError(Exception):
    """Raised when a budget is exceeded."""

    pass


class CostTrackingError(Exception):
    """Raised when there's a cost tracking error."""

    pass


@dataclass
class APICallCost:
    """Cost of a single API call.

    Attributes:
        provider: API provider name (e.g., "openai", "anthropic")
        model: Model name (e.g., "gpt-4", "claude-3")
        cost: Cost in USD
        tokens_input: Number of input tokens
        tokens_output: Number of output tokens
        timestamp: When the call was made
    """

    provider: str
    model: str
    cost: float
    tokens_input: int = 0
    tokens_output: int = 0
    timestamp: datetime = field(default_factory=datetime.utcnow)

    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary."""
        return {
            "provider": self.provider,
            "model": self.model,
            "cost": self.cost,
            "tokens_input": self.tokens_input,
            "tokens_output": self.tokens_output,
            "timestamp": self.timestamp.isoformat(),
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> APICallCost:
        """Create from dictionary."""
        return cls(
            provider=data["provider"],
            model=data["model"],
            cost=data["cost"],
            tokens_input=data.get("tokens_input", 0),
            tokens_output=data.get("tokens_output", 0),
            timestamp=datetime.fromisoformat(data["timestamp"]),
        )


class JobCost(BaseModel):
    """Cost information for a job.

    Attributes:
        job_id: Unique job identifier
        project_id: Project identifier for grouping
        tenant_id: Tenant identifier for multi-tenant support
        compute_time_seconds: Total compute time in seconds
        api_costs: List of API call costs
        storage_bytes: Storage used in bytes
        total_cost: Total cost in USD (computed)
        created_at: When the job started
        completed_at: When the job completed
        metadata: Additional metadata
    """

    job_id: str = Field(..., description="Unique job identifier")
    project_id: str | None = Field(default=None, description="Project identifier")
    tenant_id: str | None = Field(default=None, description="Tenant identifier")
    compute_time_seconds: float = Field(default=0.0, ge=0, description="Compute time in seconds")
    api_costs: list[APICallCost] = Field(default_factory=list, description="API call costs")
    storage_bytes: int = Field(default=0, ge=0, description="Storage used in bytes")
    created_at: datetime = Field(default_factory=datetime.utcnow, description="Job start time")
    completed_at: datetime | None = Field(default=None, description="Job completion time")
    metadata: dict[str, Any] = Field(default_factory=dict, description="Additional metadata")

    @property
    def total_api_cost(self) -> float:
        """Total cost from API calls."""
        return sum(call.cost for call in self.api_costs)

    @property
    def total_cost(self) -> float:
        """Total cost including API costs and compute.

        Compute cost is estimated at $0.0001 per second (approx $0.36/hour).
        """
        compute_cost = self.compute_time_seconds * 0.0001
        storage_cost = (self.storage_bytes / 1e9) * 0.023 / 30  # S3 standard rate per day
        return self.total_api_cost + compute_cost + storage_cost

    @property
    def total_tokens(self) -> int:
        """Total tokens used across all API calls."""
        return sum(call.tokens_input + call.tokens_output for call in self.api_costs)

    @property
    def duration_seconds(self) -> float:
        """Job duration in seconds."""
        if self.completed_at:
            return (self.completed_at - self.created_at).total_seconds()
        return 0.0

    def add_api_cost(
        self,
        provider: str,
        model: str,
        cost: float,
        tokens_input: int = 0,
        tokens_output: int = 0,
    ) -> None:
        """Add an API call cost."""
        self.api_costs.append(
            APICallCost(
                provider=provider,
                model=model,
                cost=cost,
                tokens_input=tokens_input,
                tokens_output=tokens_output,
            )
        )

    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary."""
        return {
            "job_id": self.job_id,
            "project_id": self.project_id,
            "tenant_id": self.tenant_id,
            "compute_time_seconds": self.compute_time_seconds,
            "api_costs": [call.to_dict() for call in self.api_costs],
            "storage_bytes": self.storage_bytes,
            "total_api_cost": self.total_api_cost,
            "total_cost": self.total_cost,
            "total_tokens": self.total_tokens,
            "created_at": self.created_at.isoformat(),
            "completed_at": self.completed_at.isoformat() if self.completed_at else None,
            "metadata": self.metadata,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> JobCost:
        """Create from dictionary."""
        job_cost = cls(
            job_id=data["job_id"],
            project_id=data.get("project_id"),
            tenant_id=data.get("tenant_id"),
            compute_time_seconds=data.get("compute_time_seconds", 0.0),
            api_costs=[APICallCost.from_dict(c) for c in data.get("api_costs", [])],
            storage_bytes=data.get("storage_bytes", 0),
            metadata=data.get("metadata", {}),
        )
        if data.get("created_at"):
            job_cost.created_at = datetime.fromisoformat(data["created_at"])
        if data.get("completed_at"):
            job_cost.completed_at = datetime.fromisoformat(data["completed_at"])
        return job_cost


@dataclass
class BudgetAlert:
    """Budget alert notification.

    Attributes:
        level: Alert severity level
        message: Alert message
        budget_id: ID of the budget that triggered the alert
        current_usage: Current usage amount
        budget_limit: Budget limit
        triggered_at: When the alert was triggered
    """

    level: AlertLevel
    message: str
    budget_id: str
    current_usage: float
    budget_limit: float
    triggered_at: datetime = field(default_factory=datetime.utcnow)

    @property
    def usage_percent(self) -> float:
        """Usage as percentage of budget."""
        if self.budget_limit <= 0:
            return 0.0
        return (self.current_usage / self.budget_limit) * 100


class CostBudget(BaseModel):
    """Budget configuration for cost tracking.

    Attributes:
        id: Unique budget identifier
        name: Human-readable budget name
        project_id: Project this budget applies to (None for global)
        tenant_id: Tenant this budget applies to (None for global)
        monthly_limit: Monthly spending limit in USD
        alert_thresholds: List of percentage thresholds for alerts (e.g., [50, 75, 90])
        alert_callback: Optional callback for budget alerts
        created_at: When the budget was created
        updated_at: When the budget was last updated
    """

    id: str = Field(..., description="Unique budget identifier")
    name: str = Field(..., description="Budget name")
    project_id: str | None = Field(default=None, description="Project identifier")
    tenant_id: str | None = Field(default=None, description="Tenant identifier")
    monthly_limit: float = Field(default=100.0, gt=0, description="Monthly limit in USD")
    alert_thresholds: list[int] = Field(
        default_factory=lambda: [50, 75, 90],
        description="Alert threshold percentages",
    )
    created_at: datetime = Field(default_factory=datetime.utcnow)
    updated_at: datetime = Field(default_factory=datetime.utcnow)

    @field_validator("alert_thresholds")
    @classmethod
    def validate_thresholds(cls, v: list[int]) -> list[int]:
        """Validate alert thresholds are between 0 and 100."""
        for threshold in v:
            if not 0 <= threshold <= 100:
                raise ValueError(f"Threshold must be between 0 and 100, got {threshold}")
        return sorted(v)

    def check_threshold(self, current_usage: float) -> list[BudgetAlert]:
        """Check if any alert thresholds have been crossed.

        Args:
            current_usage: Current spending amount

        Returns:
            List of alerts that should be triggered
        """
        alerts = []
        usage_percent = (current_usage / self.monthly_limit) * 100 if self.monthly_limit > 0 else 0

        for threshold in self.alert_thresholds:
            if usage_percent >= threshold:
                level = (
                    AlertLevel.INFO
                    if threshold < 75
                    else (AlertLevel.WARNING if threshold < 90 else AlertLevel.CRITICAL)
                )
                alerts.append(
                    BudgetAlert(
                        level=level,
                        message=f"Budget '{self.name}' has reached {threshold}% of limit",
                        budget_id=self.id,
                        current_usage=current_usage,
                        budget_limit=self.monthly_limit,
                    )
                )

        return alerts

    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary."""
        return {
            "id": self.id,
            "name": self.name,
            "project_id": self.project_id,
            "tenant_id": self.tenant_id,
            "monthly_limit": self.monthly_limit,
            "alert_thresholds": self.alert_thresholds,
            "created_at": self.created_at.isoformat(),
            "updated_at": self.updated_at.isoformat(),
        }


class CostReport(BaseModel):
    """Cost attribution report.

    Attributes:
        period: Report period type
        start_date: Report start date
        end_date: Report end date
        total_cost: Total cost for the period
        total_jobs: Total number of jobs
        costs_by_project: Cost breakdown by project
        costs_by_tenant: Cost breakdown by tenant
        costs_by_provider: API cost breakdown by provider
        top_jobs: Highest cost jobs
        daily_breakdown: Daily cost totals
    """

    period: ReportPeriod
    start_date: datetime
    end_date: datetime
    total_cost: float
    total_jobs: int
    costs_by_project: dict[str, float]
    costs_by_tenant: dict[str, float]
    costs_by_provider: dict[str, float]
    top_jobs: list[dict[str, Any]]
    daily_breakdown: dict[str, float]

    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary."""
        return {
            "period": self.period.value,
            "start_date": self.start_date.isoformat(),
            "end_date": self.end_date.isoformat(),
            "total_cost": self.total_cost,
            "total_jobs": self.total_jobs,
            "costs_by_project": self.costs_by_project,
            "costs_by_tenant": self.costs_by_tenant,
            "costs_by_provider": self.costs_by_provider,
            "top_jobs": self.top_jobs,
            "daily_breakdown": self.daily_breakdown,
        }

    def to_json(self) -> str:
        """Convert to JSON string."""
        return json.dumps(self.to_dict(), indent=2)

    def to_csv(self, filepath: Path | str) -> None:
        """Export report to CSV file.

        Creates multiple sheets:
        - summary.csv: Overall report summary
        - by_project.csv: Cost breakdown by project
        - by_provider.csv: Cost breakdown by provider
        - daily.csv: Daily breakdown
        """
        base_path = Path(filepath)
        base_path.mkdir(parents=True, exist_ok=True)

        # Summary
        with open(base_path / "summary.csv", "w", newline="") as f:
            writer = csv.writer(f)
            writer.writerow(["Period", "Start Date", "End Date", "Total Cost", "Total Jobs"])
            writer.writerow(
                [
                    self.period.value,
                    self.start_date.isoformat(),
                    self.end_date.isoformat(),
                    self.total_cost,
                    self.total_jobs,
                ]
            )

        # By project
        with open(base_path / "by_project.csv", "w", newline="") as f:
            writer = csv.writer(f)
            writer.writerow(["Project ID", "Cost"])
            for project_id, cost in sorted(
                self.costs_by_project.items(), key=lambda x: x[1], reverse=True
            ):
                writer.writerow([project_id, cost])

        # By provider
        with open(base_path / "by_provider.csv", "w", newline="") as f:
            writer = csv.writer(f)
            writer.writerow(["Provider", "Cost"])
            for provider, cost in sorted(
                self.costs_by_provider.items(), key=lambda x: x[1], reverse=True
            ):
                writer.writerow([provider, cost])

        # Daily breakdown
        with open(base_path / "daily.csv", "w", newline="") as f:
            writer = csv.writer(f)
            writer.writerow(["Date", "Cost"])
            for date, cost in sorted(self.daily_breakdown.items()):
                writer.writerow([date, cost])


class CostStore:
    """SQLite-backed storage for cost data.

    Provides persistent storage for job costs and budgets with
    efficient querying for reports and analytics.
    """

    def __init__(self, db_path: Path | str = ":memory:") -> None:
        """Initialize cost store.

        Args:
            db_path: Path to SQLite database (default: in-memory)
        """
        self.db_path = db_path
        self._lock = threading.RLock()
        self._local = threading.local()
        self._init_db()

    def _get_connection(self) -> sqlite3.Connection:
        """Get thread-local database connection."""
        if not hasattr(self._local, "connection") or self._local.connection is None:
            self._local.connection = sqlite3.connect(self.db_path, check_same_thread=False)
            self._local.connection.row_factory = sqlite3.Row
        return self._local.connection

    def _init_db(self) -> None:
        """Initialize database schema."""
        with self._lock:
            conn = self._get_connection()
            conn.executescript(
                """
                CREATE TABLE IF NOT EXISTS job_costs (
                    job_id TEXT PRIMARY KEY,
                    project_id TEXT,
                    tenant_id TEXT,
                    compute_time_seconds REAL DEFAULT 0,
                    api_costs_json TEXT DEFAULT '[]',
                    storage_bytes INTEGER DEFAULT 0,
                    total_cost REAL DEFAULT 0,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    completed_at TIMESTAMP,
                    metadata_json TEXT DEFAULT '{}'
                );

                CREATE INDEX IF NOT EXISTS idx_job_costs_project
                    ON job_costs(project_id);
                CREATE INDEX IF NOT EXISTS idx_job_costs_tenant
                    ON job_costs(tenant_id);
                CREATE INDEX IF NOT EXISTS idx_job_costs_created
                    ON job_costs(created_at);

                CREATE TABLE IF NOT EXISTS budgets (
                    id TEXT PRIMARY KEY,
                    name TEXT NOT NULL,
                    project_id TEXT,
                    tenant_id TEXT,
                    monthly_limit REAL NOT NULL,
                    alert_thresholds_json TEXT DEFAULT '[50, 75, 90]',
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                );

                CREATE INDEX IF NOT EXISTS idx_budgets_project
                    ON budgets(project_id);
                CREATE INDEX IF NOT EXISTS idx_budgets_tenant
                    ON budgets(tenant_id);

                CREATE TABLE IF NOT EXISTS budget_alerts (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    budget_id TEXT NOT NULL,
                    level TEXT NOT NULL,
                    message TEXT NOT NULL,
                    current_usage REAL NOT NULL,
                    budget_limit REAL NOT NULL,
                    triggered_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    FOREIGN KEY (budget_id) REFERENCES budgets(id)
                );

                CREATE INDEX IF NOT EXISTS idx_budget_alerts_budget
                    ON budget_alerts(budget_id);
            """
            )
            conn.commit()

    def save_job_cost(self, job_cost: JobCost) -> None:
        """Save or update a job cost record."""
        with self._lock:
            conn = self._get_connection()
            conn.execute(
                """
                INSERT OR REPLACE INTO job_costs (
                    job_id, project_id, tenant_id, compute_time_seconds,
                    api_costs_json, storage_bytes, total_cost, created_at,
                    completed_at, metadata_json
                ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                """,
                (
                    job_cost.job_id,
                    job_cost.project_id,
                    job_cost.tenant_id,
                    job_cost.compute_time_seconds,
                    json.dumps([call.to_dict() for call in job_cost.api_costs]),
                    job_cost.storage_bytes,
                    job_cost.total_cost,
                    job_cost.created_at.isoformat(),
                    job_cost.completed_at.isoformat() if job_cost.completed_at else None,
                    json.dumps(job_cost.metadata),
                ),
            )
            conn.commit()

    def get_job_cost(self, job_id: str) -> JobCost | None:
        """Get cost information for a specific job."""
        with self._lock:
            conn = self._get_connection()
            row = conn.execute(
                "SELECT * FROM job_costs WHERE job_id = ?",
                (job_id,),
            ).fetchone()

            if not row:
                return None

            return self._row_to_job_cost(row)

    def get_job_costs(
        self,
        project_id: str | None = None,
        tenant_id: str | None = None,
        start_date: datetime | None = None,
        end_date: datetime | None = None,
        limit: int = 1000,
    ) -> list[JobCost]:
        """Query job costs with filters."""
        with self._lock:
            conn = self._get_connection()
            query = "SELECT * FROM job_costs WHERE 1=1"
            params = []

            if project_id:
                query += " AND project_id = ?"
                params.append(project_id)
            if tenant_id:
                query += " AND tenant_id = ?"
                params.append(tenant_id)
            if start_date:
                query += " AND created_at >= ?"
                params.append(start_date.isoformat())
            if end_date:
                query += " AND created_at <= ?"
                params.append(end_date.isoformat())

            query += " ORDER BY created_at DESC LIMIT ?"
            params.append(limit)

            rows = conn.execute(query, params).fetchall()
            return [self._row_to_job_cost(row) for row in rows]

    def _row_to_job_cost(self, row: sqlite3.Row) -> JobCost:
        """Convert database row to JobCost."""
        api_costs = json.loads(row["api_costs_json"] or "[]")
        metadata = json.loads(row["metadata_json"] or "{}")

        job_cost = JobCost(
            job_id=row["job_id"],
            project_id=row["project_id"],
            tenant_id=row["tenant_id"],
            compute_time_seconds=row["compute_time_seconds"],
            api_costs=[APICallCost.from_dict(c) for c in api_costs],
            storage_bytes=row["storage_bytes"],
            metadata=metadata,
        )
        job_cost.created_at = datetime.fromisoformat(row["created_at"])
        if row["completed_at"]:
            job_cost.completed_at = datetime.fromisoformat(row["completed_at"])

        return job_cost

    def save_budget(self, budget: CostBudget) -> None:
        """Save or update a budget."""
        with self._lock:
            conn = self._get_connection()
            budget.updated_at = datetime.utcnow()
            conn.execute(
                """
                INSERT OR REPLACE INTO budgets (
                    id, name, project_id, tenant_id, monthly_limit,
                    alert_thresholds_json, created_at, updated_at
                ) VALUES (?, ?, ?, ?, ?, ?, ?, ?)
                """,
                (
                    budget.id,
                    budget.name,
                    budget.project_id,
                    budget.tenant_id,
                    budget.monthly_limit,
                    json.dumps(budget.alert_thresholds),
                    budget.created_at.isoformat(),
                    budget.updated_at.isoformat(),
                ),
            )
            conn.commit()

    def get_budget(self, budget_id: str) -> CostBudget | None:
        """Get a budget by ID."""
        with self._lock:
            conn = self._get_connection()
            row = conn.execute(
                "SELECT * FROM budgets WHERE id = ?",
                (budget_id,),
            ).fetchone()

            if not row:
                return None

            return self._row_to_budget(row)

    def get_budgets(
        self,
        project_id: str | None = None,
        tenant_id: str | None = None,
    ) -> list[CostBudget]:
        """Get budgets with optional filters."""
        with self._lock:
            conn = self._get_connection()
            query = "SELECT * FROM budgets WHERE 1=1"
            params = []

            if project_id:
                query += " AND project_id = ?"
                params.append(project_id)
            if tenant_id:
                query += " AND tenant_id = ?"
                params.append(tenant_id)

            rows = conn.execute(query, params).fetchall()
            return [self._row_to_budget(row) for row in rows]

    def _row_to_budget(self, row: sqlite3.Row) -> CostBudget:
        """Convert database row to CostBudget."""
        budget = CostBudget(
            id=row["id"],
            name=row["name"],
            project_id=row["project_id"],
            tenant_id=row["tenant_id"],
            monthly_limit=row["monthly_limit"],
            alert_thresholds=json.loads(row["alert_thresholds_json"]),
        )
        budget.created_at = datetime.fromisoformat(row["created_at"])
        budget.updated_at = datetime.fromisoformat(row["updated_at"])
        return budget

    def delete_budget(self, budget_id: str) -> bool:
        """Delete a budget."""
        with self._lock:
            conn = self._get_connection()
            cursor = conn.execute("DELETE FROM budgets WHERE id = ?", (budget_id,))
            conn.commit()
            return cursor.rowcount > 0

    def save_alert(self, alert: BudgetAlert) -> None:
        """Save a budget alert."""
        with self._lock:
            conn = self._get_connection()
            conn.execute(
                """
                INSERT INTO budget_alerts (
                    budget_id, level, message, current_usage, budget_limit, triggered_at
                ) VALUES (?, ?, ?, ?, ?, ?)
                """,
                (
                    alert.budget_id,
                    alert.level.value,
                    alert.message,
                    alert.current_usage,
                    alert.budget_limit,
                    alert.triggered_at.isoformat(),
                ),
            )
            conn.commit()

    def get_monthly_usage(
        self,
        project_id: str | None = None,
        tenant_id: str | None = None,
        year: int | None = None,
        month: int | None = None,
    ) -> float:
        """Get total usage for a month."""
        now = datetime.utcnow()
        year = year or now.year
        month = month or now.month

        start_date = datetime(year, month, 1)
        end_date = datetime(year + 1, 1, 1) if month == 12 else datetime(year, month + 1, 1)

        with self._lock:
            conn = self._get_connection()
            query = """
                SELECT COALESCE(SUM(total_cost), 0) as total
                FROM job_costs
                WHERE created_at >= ? AND created_at < ?
            """
            params = [start_date.isoformat(), end_date.isoformat()]

            if project_id:
                query += " AND project_id = ?"
                params.append(project_id)
            if tenant_id:
                query += " AND tenant_id = ?"
                params.append(tenant_id)

            row = conn.execute(query, params).fetchone()
            return row["total"] if row else 0.0

    def close(self) -> None:
        """Close database connection."""
        with self._lock:
            if hasattr(self._local, "connection") and self._local.connection:
                self._local.connection.close()
                self._local.connection = None


class CostTracker:
    """Tracks costs during job execution.

    Provides a context manager for tracking job costs including:
    - Compute time
    - API costs
    - Storage usage

    Example:
        >>> tracker = CostTracker()
        >>> with tracker.track_job("job-123", project_id="proj-1") as job:
        ...     job.add_api_cost("openai", "gpt-4", 0.05, 100, 50)
        ...     # ... do work ...
        >>> cost = tracker.get_job_cost("job-123")
    """

    def __init__(self, store: CostStore | None = None) -> None:
        """Initialize cost tracker.

        Args:
            store: CostStore instance (creates in-memory store if None)
        """
        self.store = store or CostStore()
        self._active_jobs: dict[str, JobCost] = {}
        self._job_start_times: dict[str, float] = {}
        self._lock = threading.RLock()

    def start_job(
        self,
        job_id: str,
        project_id: str | None = None,
        tenant_id: str | None = None,
        metadata: dict[str, Any] | None = None,
    ) -> JobCost:
        """Start tracking a new job.

        Args:
            job_id: Unique job identifier
            project_id: Optional project identifier
            tenant_id: Optional tenant identifier
            metadata: Optional metadata dictionary

        Returns:
            JobCost instance for the job
        """
        with self._lock:
            if job_id in self._active_jobs:
                raise CostTrackingError(f"Job {job_id} is already being tracked")

            job_cost = JobCost(
                job_id=job_id,
                project_id=project_id,
                tenant_id=tenant_id,
                metadata=metadata or {},
            )
            self._active_jobs[job_id] = job_cost
            self._job_start_times[job_id] = time.perf_counter()
            return job_cost

    def end_job(self, job_id: str) -> JobCost:
        """End tracking a job and save the cost data.

        Args:
            job_id: Job identifier

        Returns:
            Completed JobCost

        Raises:
            CostTrackingError: If job is not being tracked
        """
        with self._lock:
            if job_id not in self._active_jobs:
                raise CostTrackingError(f"Job {job_id} is not being tracked")

            job_cost = self._active_jobs.pop(job_id)
            start_time = self._job_start_times.pop(job_id, None)

            if start_time:
                job_cost.compute_time_seconds = time.perf_counter() - start_time

            job_cost.completed_at = datetime.utcnow()
            self.store.save_job_cost(job_cost)
            return job_cost

    def record_api_cost(
        self,
        job_id: str,
        provider: str,
        model: str,
        cost: float,
        tokens_input: int = 0,
        tokens_output: int = 0,
    ) -> None:
        """Record an API call cost for a job.

        Args:
            job_id: Job identifier
            provider: API provider name
            model: Model name
            cost: Cost in USD
            tokens_input: Input token count
            tokens_output: Output token count
        """
        with self._lock:
            if job_id not in self._active_jobs:
                raise CostTrackingError(f"Job {job_id} is not being tracked")

            self._active_jobs[job_id].add_api_cost(
                provider=provider,
                model=model,
                cost=cost,
                tokens_input=tokens_input,
                tokens_output=tokens_output,
            )

    def record_storage(self, job_id: str, bytes_used: int) -> None:
        """Record storage usage for a job.

        Args:
            job_id: Job identifier
            bytes_used: Storage used in bytes
        """
        with self._lock:
            if job_id not in self._active_jobs:
                raise CostTrackingError(f"Job {job_id} is not being tracked")

            self._active_jobs[job_id].storage_bytes += bytes_used

    def get_job_cost(self, job_id: str) -> JobCost | None:
        """Get cost information for a job.

        Checks active jobs first, then database.
        """
        with self._lock:
            if job_id in self._active_jobs:
                return self._active_jobs[job_id]
        return self.store.get_job_cost(job_id)

    @contextmanager
    def track_job(
        self,
        job_id: str,
        project_id: str | None = None,
        tenant_id: str | None = None,
        metadata: dict[str, Any] | None = None,
    ) -> Iterator[JobCost]:
        """Context manager for tracking a job.

        Automatically starts and ends job tracking.

        Example:
            >>> with tracker.track_job("job-123", project_id="proj-1") as job:
            ...     job.add_api_cost("openai", "gpt-4", 0.05)
            ...     # ... do work ...
        """
        job = self.start_job(job_id, project_id, tenant_id, metadata)
        try:
            yield job
        finally:
            self.end_job(job_id)

    def generate_report(
        self,
        period: ReportPeriod,
        project_id: str | None = None,
        tenant_id: str | None = None,
    ) -> CostReport:
        """Generate a cost attribution report.

        Args:
            period: Report period (daily, weekly, monthly)
            project_id: Optional project filter
            tenant_id: Optional tenant filter

        Returns:
            CostReport with aggregated data
        """
        now = datetime.utcnow()

        if period == ReportPeriod.DAILY:
            start_date = now - timedelta(days=1)
        elif period == ReportPeriod.WEEKLY:
            start_date = now - timedelta(weeks=1)
        else:  # MONTHLY
            start_date = now - timedelta(days=30)

        end_date = now

        jobs = self.store.get_job_costs(
            project_id=project_id,
            tenant_id=tenant_id,
            start_date=start_date,
            end_date=end_date,
            limit=10000,
        )

        # Aggregate data
        total_cost = sum(job.total_cost for job in jobs)
        costs_by_project: dict[str, float] = {}
        costs_by_tenant: dict[str, float] = {}
        costs_by_provider: dict[str, float] = {}
        daily_breakdown: dict[str, float] = {}

        for job in jobs:
            # By project
            proj = job.project_id or "default"
            costs_by_project[proj] = costs_by_project.get(proj, 0) + job.total_cost

            # By tenant
            ten = job.tenant_id or "default"
            costs_by_tenant[ten] = costs_by_tenant.get(ten, 0) + job.total_cost

            # By provider
            for call in job.api_costs:
                costs_by_provider[call.provider] = (
                    costs_by_provider.get(call.provider, 0) + call.cost
                )

            # Daily breakdown
            date_key = job.created_at.strftime("%Y-%m-%d")
            daily_breakdown[date_key] = daily_breakdown.get(date_key, 0) + job.total_cost

        # Top jobs
        top_jobs = sorted(jobs, key=lambda j: j.total_cost, reverse=True)[:10]
        top_jobs_data = [
            {
                "job_id": job.job_id,
                "project_id": job.project_id,
                "total_cost": job.total_cost,
                "api_cost": job.total_api_cost,
                "compute_time": job.compute_time_seconds,
            }
            for job in top_jobs
        ]

        return CostReport(
            period=period,
            start_date=start_date,
            end_date=end_date,
            total_cost=total_cost,
            total_jobs=len(jobs),
            costs_by_project=costs_by_project,
            costs_by_tenant=costs_by_tenant,
            costs_by_provider=costs_by_provider,
            top_jobs=top_jobs_data,
            daily_breakdown=daily_breakdown,
        )

    def close(self) -> None:
        """Close the cost tracker and store."""
        self.store.close()


class BudgetManager:
    """Manages cost budgets and alerts.

    Provides budget CRUD operations and automatic alert generation
    when budgets are exceeded.

    Example:
        >>> manager = BudgetManager(cost_store)
        >>> budget = manager.create_budget(
        ...     name="Project A",
        ...     project_id="proj-a",
        ...     monthly_limit=100.0
        ... )
        >>> alerts = manager.check_budget(budget.id)
    """

    def __init__(
        self,
        store: CostStore,
        alert_callback: Callable[[BudgetAlert], None] | None = None,
    ) -> None:
        """Initialize budget manager.

        Args:
            store: CostStore instance
            alert_callback: Optional callback for budget alerts
        """
        self.store = store
        self.alert_callback = alert_callback
        self._lock = threading.RLock()

    def create_budget(
        self,
        name: str,
        project_id: str | None = None,
        tenant_id: str | None = None,
        monthly_limit: float = 100.0,
        alert_thresholds: list[int] | None = None,
        budget_id: str | None = None,
    ) -> CostBudget:
        """Create a new budget.

        Args:
            name: Budget name
            project_id: Optional project identifier
            tenant_id: Optional tenant identifier
            monthly_limit: Monthly spending limit
            alert_thresholds: List of alert percentage thresholds
            budget_id: Optional custom budget ID

        Returns:
            Created CostBudget
        """
        budget = CostBudget(
            id=budget_id or f"budget-{int(time.time() * 1000)}",
            name=name,
            project_id=project_id,
            tenant_id=tenant_id,
            monthly_limit=monthly_limit,
            alert_thresholds=alert_thresholds or [50, 75, 90],
        )
        self.store.save_budget(budget)
        return budget

    def get_budget(self, budget_id: str) -> CostBudget | None:
        """Get a budget by ID."""
        return self.store.get_budget(budget_id)

    def list_budgets(
        self,
        project_id: str | None = None,
        tenant_id: str | None = None,
    ) -> list[CostBudget]:
        """List budgets with optional filters."""
        return self.store.get_budgets(project_id=project_id, tenant_id=tenant_id)

    def update_budget(
        self,
        budget_id: str,
        name: str | None = None,
        monthly_limit: float | None = None,
        alert_thresholds: list[int] | None = None,
    ) -> CostBudget | None:
        """Update a budget.

        Args:
            budget_id: Budget ID
            name: New name (optional)
            monthly_limit: New limit (optional)
            alert_thresholds: New thresholds (optional)

        Returns:
            Updated budget or None if not found
        """
        budget = self.store.get_budget(budget_id)
        if not budget:
            return None

        if name is not None:
            budget.name = name
        if monthly_limit is not None:
            budget.monthly_limit = monthly_limit
        if alert_thresholds is not None:
            budget.alert_thresholds = alert_thresholds

        budget.updated_at = datetime.utcnow()
        self.store.save_budget(budget)
        return budget

    def delete_budget(self, budget_id: str) -> bool:
        """Delete a budget."""
        return self.store.delete_budget(budget_id)

    def check_budget(self, budget_id: str) -> list[BudgetAlert]:
        """Check budget status and generate alerts.

        Args:
            budget_id: Budget ID to check

        Returns:
            List of triggered alerts
        """
        with self._lock:
            budget = self.store.get_budget(budget_id)
            if not budget:
                return []

            # Get current usage
            usage = self.store.get_monthly_usage(
                project_id=budget.project_id,
                tenant_id=budget.tenant_id,
            )

            # Check thresholds
            alerts = budget.check_threshold(usage)

            # Save alerts and trigger callbacks
            for alert in alerts:
                self.store.save_alert(alert)
                if self.alert_callback:
                    with suppress(Exception):
                        self.alert_callback(alert)

            return alerts

    def check_all_budgets(self) -> dict[str, list[BudgetAlert]]:
        """Check all budgets and return alerts.

        Returns:
            Dictionary mapping budget IDs to their alerts
        """
        budgets = self.store.get_budgets()
        results = {}
        for budget in budgets:
            alerts = self.check_budget(budget.id)
            if alerts:
                results[budget.id] = alerts
        return results

    def get_budget_usage(self, budget_id: str) -> dict[str, float] | None:
        """Get current usage for a budget.

        Returns:
            Dictionary with usage data or None if budget not found
        """
        budget = self.store.get_budget(budget_id)
        if not budget:
            return None

        usage = self.store.get_monthly_usage(
            project_id=budget.project_id,
            tenant_id=budget.tenant_id,
        )

        return {
            "current_usage": usage,
            "monthly_limit": budget.monthly_limit,
            "remaining": max(0, budget.monthly_limit - usage),
            "usage_percent": (
                (usage / budget.monthly_limit * 100) if budget.monthly_limit > 0 else 0
            ),
        }
