"""
Shared data models for LLM execution.

These classes are the canonical definitions; they are re-exported by
markitect.prompts.execution.models for backward compatibility.
"""

import threading
from dataclasses import dataclass, field
from typing import Dict, Any, Optional

from llm_connect.exceptions import LLMBudgetExceededError


class BudgetTracker:
    """Shared token budget for a call or delegation chain.

    Thread-safe. Tracks cumulative token spend across multiple adapter
    calls. Raises ``LLMBudgetExceededError`` when the cap is exceeded.

    Example::

        tracker = BudgetTracker(total=4000)
        config = RunConfig(budget_tracker=tracker)
        # All adapter calls sharing this config will consume from the same cap.
    """

    def __init__(self, total: int) -> None:
        if total <= 0:
            raise ValueError(f"BudgetTracker total must be positive, got {total}")
        self.total = total
        self.spent = 0
        self._lock = threading.Lock()

    def remaining(self) -> int:
        """Return tokens remaining in the budget."""
        return max(0, self.total - self.spent)

    def consume(self, tokens: int) -> None:
        """Record *tokens* as spent. Raises ``LLMBudgetExceededError`` if cap exceeded."""
        with self._lock:
            new_spent = self.spent + tokens
            if new_spent > self.total:
                raise LLMBudgetExceededError(
                    f"Token budget exceeded: {new_spent} tokens used, cap is {self.total}",
                    total=self.total,
                    spent=self.spent,
                    requested=tokens,
                )
            self.spent = new_spent

    def __repr__(self) -> str:
        return f"BudgetTracker(total={self.total}, spent={self.spent}, remaining={self.remaining()})"


@dataclass
class RunConfig:
    """
    Configuration for prompt execution.

    Attributes:
        model_name: LLM model to use
        temperature: Model temperature (0.0-1.0)
        max_tokens: Maximum tokens to generate
        model_params: Additional model parameters
        max_depth: Maximum generation depth for nested runs
        skip_if_exists: Skip if identical InputBundleHash exists
        timeout_seconds: Execution timeout
    """
    model_name: str = "gpt-4"
    temperature: float = 0.7
    max_tokens: int = 2000
    model_params: Dict[str, Any] = field(default_factory=dict)
    max_depth: int = 3
    skip_if_exists: bool = True
    timeout_seconds: int = 300
    budget_tracker: Optional["BudgetTracker"] = field(default=None, repr=False)

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary. ``budget_tracker`` is excluded (runtime object)."""
        return {
            "model_name": self.model_name,
            "temperature": self.temperature,
            "max_tokens": self.max_tokens,
            "model_params": self.model_params,
            "max_depth": self.max_depth,
            "skip_if_exists": self.skip_if_exists,
            "timeout_seconds": self.timeout_seconds,
        }

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "RunConfig":
        """Create from dictionary."""
        return cls(
            model_name=data.get("model_name", "gpt-4"),
            temperature=data.get("temperature", 0.7),
            max_tokens=data.get("max_tokens", 2000),
            model_params=data.get("model_params", {}),
            max_depth=data.get("max_depth", 3),
            skip_if_exists=data.get("skip_if_exists", True),
            timeout_seconds=data.get("timeout_seconds", 300),
        )


@dataclass
class LLMResponse:
    """
    Response from LLM execution.

    Attributes:
        content: Generated content
        model: Model used
        usage: Token usage statistics
        finish_reason: Why generation stopped
        metadata: Additional response metadata
    """
    content: str
    model: str
    usage: Dict[str, int] = field(default_factory=dict)
    finish_reason: str = "stop"
    metadata: Dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary."""
        return {
            "content": self.content,
            "model": self.model,
            "usage": self.usage,
            "finish_reason": self.finish_reason,
            "metadata": self.metadata,
        }