"""CostSession: context manager for tracking LLM costs in a pipeline."""

import uuid
import warnings
from datetime import datetime, timezone
from typing import Any, Optional

from agentbudget.exceptions import BudgetExceededError, BudgetExceededWarning
from agentbudget.pricing import calculate_cost
from agentbudget.providers.base import detect_provider


def _log(verbose: bool, message: str) -> None:
    """Print message only if verbose is True.

    Args:
        verbose: Whether to output the message.
        message: The message to print.
    """
    if verbose:
        print(message)


class CostSession:
    """Context manager for tracking LLM API costs in a pipeline or agent.

    Tracks individual LLM calls (steps), accumulates costs, enforces budget
    limits, and prints a formatted summary at the end of the session.

    Example:
        with CostSession(name="my-pipeline", budget=0.05) as session:
            response = openai_client.chat.completions.create(...)
            session.track(response, label="step-1")
    """

    # Separator characters for summary output
    _LINE_CHAR = "─"
    _SEP_CHAR = "│"

    def __init__(
        self,
        name: str,
        budget: Optional[float] = None,
        strict: bool = False,
        verbose: bool = True,
        provider: str = "auto",
    ) -> None:
        """Initialize a CostSession.

        Args:
            name: Session name, used in the summary header and persisted storage.
            budget: Optional spending limit in USD. If exceeded, triggers a
                warning (default) or exception (if strict=True).
            strict: If True, raise BudgetExceededError when budget is exceeded.
                If False (default), emit a BudgetExceededWarning.
            verbose: If True (default), print the formatted summary when the
                context manager exits. Set to False to suppress console output.
            provider: Provider hint: "openai", "anthropic", "gemini", or "auto" (default).
                With "auto", the provider is detected from each response object.
        """
        self.name = name
        self.budget = budget
        self.strict = strict
        self.verbose = verbose
        self.provider = provider
        self._steps: list[dict] = []

    @property
    def total_cost(self) -> float:
        """Total accumulated cost across all tracked steps, in USD."""
        return sum(step["cost"] for step in self._steps)

    @property
    def total_tokens(self) -> int:
        """Total token count across all tracked steps."""
        return sum(step["input_tokens"] + step["output_tokens"] for step in self._steps)

    def track(self, response: Any, label: str = "", model: Optional[str] = None) -> None:
        """Track an LLM API response and accumulate its cost.

        Detects the provider automatically (or uses the provider set in the
        constructor), extracts token usage, calculates cost, and appends the
        result as a step. After each call, checks the budget and warns or
        raises an exception if it is exceeded.

        Args:
            response: The raw API response object from OpenAI, Anthropic, or Gemini SDK.
            label: Optional human-readable label for this step (e.g., "generate",
                "summarize"). Defaults to "step-N" if not provided.
            model: Optional model name override. Use this when the response object
                does not include the model name (e.g., older google-generativeai SDK).
                Overrides any model name extracted from the response.

        Raises:
            ValueError: If the provider cannot be detected from the response.
            BudgetExceededError: If budget is set, strict=True, and cost exceeds budget.
        """
        # Detect provider
        if self.provider == "auto":
            provider_adapter = detect_provider(response)
        elif self.provider == "openai":
            from agentbudget.providers.openai import OpenAIProvider
            provider_adapter = OpenAIProvider()
        elif self.provider == "anthropic":
            from agentbudget.providers.anthropic import AnthropicProvider
            provider_adapter = AnthropicProvider()
        elif self.provider == "gemini":
            from agentbudget.providers.gemini import GeminiProvider
            provider_adapter = GeminiProvider()
        else:
            raise ValueError(
                f"Unknown provider: {self.provider!r}. "
                "Use 'auto', 'openai', 'anthropic', or 'gemini'."
            )

        usage = provider_adapter.extract_usage(response)
        effective_model = model if model is not None else usage.model
        cost = calculate_cost(effective_model, usage.input_tokens, usage.output_tokens)

        step_label = label or f"step-{len(self._steps) + 1}"

        self._steps.append({
            "label": step_label,
            "model": effective_model,
            "provider": usage.provider,
            "input_tokens": usage.input_tokens,
            "output_tokens": usage.output_tokens,
            "cost": cost,
        })

        # Budget enforcement
        if self.budget is not None and self.total_cost > self.budget:
            msg = (
                f"[agentbudget] Budget of ${self.budget:.2f} exceeded. "
                f"Current cost: ${self.total_cost:.4f}"
            )
            if self.strict:
                raise BudgetExceededError(msg)
            else:
                warnings.warn(msg, BudgetExceededWarning, stacklevel=2)

    def summary(self) -> str:
        """Return the formatted session summary as a string.

        Useful for logging the summary to a file or external system instead
        of (or in addition to) printing it to the console.

        Returns:
            A multi-line string with the formatted session summary table.
        """
        line_char = self._LINE_CHAR
        sep = self._SEP_CHAR
        lines = []

        # Compute dynamic column widths based on actual content
        label_width = max(
            (len(step["label"]) for step in self._steps),
            default=5,
        )
        label_width = max(label_width, len("TOTAL"), len("Budget"), 8) + 1

        model_width = max(
            (len(step["model"]) for step in self._steps),
            default=5,
        )
        model_width = max(model_width, 5) + 1

        # Total line width: space + label + space + sep + space + model + space + sep + tok_col(11) + sep + cost_col(9) + sep + pct_col(4)
        # tok_col = " XXXXXX tok " = 1+6+5 = 12 chars, cost_col = " $X.XXXX  " = 10 chars, pct_col = " XX%" = 4 chars
        # Each column has 1 leading space
        # label_col = 1+label_width chars, model_col = 1+model_width chars
        line_width = (1 + label_width) + 1 + (1 + model_width) + 1 + 12 + 1 + 10 + 1 + 4
        line = line_char * line_width

        lines.append(line)
        lines.append(f" Session: {self.name}")
        lines.append(line)

        total = self.total_cost

        for step in self._steps:
            tokens = step["input_tokens"] + step["output_tokens"]
            cost = step["cost"]
            pct = round((cost / total) * 100) if total > 0 else 0

            label_col = f" {step['label']:<{label_width}}"
            model_col = f" {step['model']:<{model_width}}"
            tok_col = f" {tokens:>6,} tok "
            cost_col = f" ${cost:.4f}  "
            pct_col = f" {pct:>2}%"

            lines.append(f"{label_col}{sep}{model_col}{sep}{tok_col}{sep}{cost_col}{sep}{pct_col}")

        lines.append(line)

        total_tokens = self.total_tokens
        total_col = f" {'TOTAL':<{label_width}}"
        empty_model_col = f" {'':<{model_width}}"
        total_tok_col = f" {total_tokens:>6,} tok "
        total_cost_col = f" ${total:.4f}"
        lines.append(f"{total_col}{sep}{empty_model_col}{sep}{total_tok_col}{sep}{total_cost_col}")

        if self.budget is not None:
            if total <= self.budget:
                status = "✓ Within budget"
            else:
                status = "✗ Budget exceeded"
            budget_col = f" {'Budget':<{label_width}}"
            budget_val_col = f" ${self.budget:<{model_width}.2f}"
            lines.append(f"{budget_col}{sep}{budget_val_col}{sep} {status}")

        lines.append(line)

        return "\n".join(lines)

    def __enter__(self) -> "CostSession":
        """Enter the context manager."""
        return self

    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        """Exit the context manager: persist session and print summary."""
        if self._steps:
            # Persist to storage
            try:
                from agentbudget.tracker import CostTracker
                tracker = CostTracker()
                tracker.save_session(self._to_dict())
            except OSError as e:
                _log(self.verbose, f"[agentbudget] Warning: could not persist session: {e}")

            # Print summary
            _log(self.verbose, self.summary())

    def _to_dict(self) -> dict:
        """Serialize the session to a dict for JSON storage."""
        return {
            "id": str(uuid.uuid4()),
            "name": self.name,
            "timestamp": datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"),
            "budget": self.budget,
            "total_cost": self.total_cost,
            "total_tokens": self.total_tokens,
            "steps": list(self._steps),
        }
