"""
Session: conversation persistence for goal-oriented execution.

Maintains conversation history across multiple run() calls.
Each exchange (user prompt + assistant response) becomes context for subsequent calls.

Usage:
    session = Session(model="gpt-4")
    session.run("Find the bug", data=logs, cost=5.00)
    session.run("Fix it")  # Model has context from previous exchange
    session.save("session.json")
"""
from __future__ import annotations

import json
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional, Union

from enzu.models import ExecutionReport, RLMExecutionReport


class SessionBudgetExceeded(Exception):
    """Raised when session budget cap (cost or tokens) is exceeded."""

    def __init__(
        self,
        *,
        cost_used: Optional[float] = None,
        cost_cap: Optional[float] = None,
        tokens_used: Optional[int] = None,
        tokens_cap: Optional[int] = None,
    ):
        self.cost_used = cost_used
        self.cost_cap = cost_cap
        self.tokens_used = tokens_used
        self.tokens_cap = tokens_cap

        if cost_cap is not None and cost_used is not None:
            msg = (
                f"Session budget exceeded: ${cost_used:.2f} used of ${cost_cap:.2f} cap. "
                f"Call session.raise_cost_cap(new_amount) to continue."
            )
        elif tokens_cap is not None and tokens_used is not None:
            msg = (
                f"Session token limit exceeded: {tokens_used:,} used of {tokens_cap:,} cap. "
                f"Call session.raise_token_cap(new_amount) to continue."
            )
        else:
            msg = "Session budget exceeded."
        super().__init__(msg)


@dataclass
class Exchange:
    """Single conversation turn: user prompt + assistant response."""
    user: str
    assistant: str
    data_snippet: Optional[str] = None  # First 500 chars of data, for context
    timestamp: str = field(default_factory=lambda: datetime.now().isoformat())
    cost_usd: Optional[float] = None

    def to_dict(self) -> Dict[str, Any]:
        return {
            "user": self.user,
            "assistant": self.assistant,
            "data_snippet": self.data_snippet,
            "timestamp": self.timestamp,
            "cost_usd": self.cost_usd,
        }

    @classmethod
    def from_dict(cls, d: Dict[str, Any]) -> "Exchange":
        return cls(
            user=d["user"],
            assistant=d["assistant"],
            data_snippet=d.get("data_snippet"),
            timestamp=d.get("timestamp", datetime.now().isoformat()),
            cost_usd=d.get("cost_usd"),
        )


def _format_history(exchanges: List[Exchange], max_chars: int = 10000) -> str:
    """
    Format conversation history as text for model context.

    Returns empty string if no history.
    Truncates oldest exchanges if over max_chars.
    """
    if not exchanges:
        return ""

    parts = []
    total = 0

    # Build from most recent, then reverse
    for ex in reversed(exchanges):
        part = f"User: {ex.user}\n"
        if ex.data_snippet:
            part += f"[Data provided: {len(ex.data_snippet)} chars]\n"
        part += f"Assistant: {ex.assistant}\n"

        if total + len(part) > max_chars:
            break
        parts.append(part)
        total += len(part)

    if not parts:
        return ""

    parts.reverse()
    return "== Previous Conversation ==\n" + "\n".join(parts) + "\n== End Previous ==\n"


class Session:
    """
    Conversation session with history persistence.

    Each run() call:
    1. Prepends conversation history to data
    2. Executes the task
    3. Appends the exchange to history

    Example:
        session = Session(model="gpt-4")
        session.run("Find the bug", data=logs, cost=5.00)
        session.run("Fix it")  # Has context from previous
        session.save("debug_session.json")
    """

    def __init__(
        self,
        model: str,
        *,
        provider: Optional[str] = None,
        api_key: Optional[str] = None,
        history_max_chars: int = 10000,
        max_cost_usd: Optional[float] = None,
        max_tokens: Optional[int] = None,
    ):
        self.model = model
        self.provider = provider or "openrouter"
        self.api_key = api_key
        self.history_max_chars = history_max_chars
        # Budget caps: None = unlimited.
        self.max_cost_usd = max_cost_usd
        self.max_tokens = max_tokens
        self.exchanges: List[Exchange] = []
        self.total_cost_usd: float = 0.0
        self.total_tokens: int = 0
        self.created_at: str = datetime.now().isoformat()

    def run(
        self,
        task: str,
        *,
        data: Optional[str] = None,
        tokens: Optional[int] = None,
        seconds: Optional[float] = None,
        cost: Optional[float] = None,
        contains: Optional[List[str]] = None,
        matches: Optional[List[str]] = None,
        min_words: Optional[int] = None,
        goal: Optional[str] = None,
        temperature: Optional[float] = None,
        max_steps: Optional[int] = None,
        on_progress: Optional[Callable[[str], None]] = None,
        return_report: bool = False,
    ) -> Union[str, ExecutionReport, RLMExecutionReport]:
        """
        Run a task with conversation history as context.

        History is prepended to data, so the model sees previous exchanges.
        After completion, this exchange is added to history.

        Raises:
            SessionBudgetExceeded: If session budget cap (cost or tokens) would be exceeded.
        """
        # Check budget caps before execution
        if self.max_cost_usd is not None and self.total_cost_usd >= self.max_cost_usd:
            raise SessionBudgetExceeded(
                cost_used=self.total_cost_usd,
                cost_cap=self.max_cost_usd,
            )
        if self.max_tokens is not None and self.total_tokens >= self.max_tokens:
            raise SessionBudgetExceeded(
                tokens_used=self.total_tokens,
                tokens_cap=self.max_tokens,
            )

        # Import here to avoid circular dependency
        from enzu.api import run as enzu_run

        # Format history as text
        history_text = _format_history(self.exchanges, self.history_max_chars)

        # Combine history with new data
        if history_text:
            if data:
                combined_data = history_text + "\n== Current Data ==\n" + data
            else:
                combined_data = history_text
        else:
            combined_data = data

        # Execute
        report = enzu_run(
            task,
            model=self.model,
            provider=self.provider,
            data=combined_data,
            tokens=tokens,
            seconds=seconds,
            cost=cost,
            contains=contains,
            matches=matches,
            min_words=min_words,
            goal=goal,
            temperature=temperature,
            api_key=self.api_key,
            max_steps=max_steps,
            on_progress=on_progress,
            return_report=True,
        )

        # Extract answer and usage
        if isinstance(report, RLMExecutionReport):
            answer = report.answer or ""
            run_cost = report.budget_usage.cost_usd
            run_tokens = report.budget_usage.total_tokens
        else:
            answer = report.output_text or ""
            run_cost = report.budget_usage.cost_usd
            run_tokens = report.budget_usage.total_tokens

        # Track usage
        if run_cost:
            self.total_cost_usd += run_cost
        if run_tokens:
            self.total_tokens += run_tokens

        # Append exchange to history
        data_snippet = data[:500] if data else None
        self.exchanges.append(Exchange(
            user=task,
            assistant=answer,
            data_snippet=data_snippet,
            cost_usd=run_cost,
        ))

        if return_report:
            return report
        return answer

    def clear(self) -> None:
        """Clear conversation history."""
        self.exchanges = []

    def raise_cost_cap(self, new_cap: float) -> None:
        """Raise the cost budget cap. Must be higher than current cap."""
        if self.max_cost_usd is not None and new_cap <= self.max_cost_usd:
            raise ValueError(f"New cap ${new_cap:.2f} must be higher than current ${self.max_cost_usd:.2f}")
        self.max_cost_usd = new_cap

    # Alias for backward compat
    raise_cap = raise_cost_cap

    def raise_token_cap(self, new_cap: int) -> None:
        """Raise the token budget cap. Must be higher than current cap."""
        if self.max_tokens is not None and new_cap <= self.max_tokens:
            raise ValueError(f"New cap {new_cap:,} must be higher than current {self.max_tokens:,}")
        self.max_tokens = new_cap

    @property
    def remaining_budget(self) -> Optional[float]:
        """Remaining cost budget. None if unlimited."""
        if self.max_cost_usd is None:
            return None
        return max(0, self.max_cost_usd - self.total_cost_usd)

    @property
    def remaining_tokens(self) -> Optional[int]:
        """Remaining token budget. None if unlimited."""
        if self.max_tokens is None:
            return None
        return max(0, self.max_tokens - self.total_tokens)

    @property
    def history(self) -> List[Dict[str, Any]]:
        """Get conversation history as list of dicts."""
        return [ex.to_dict() for ex in self.exchanges]

    def to_dict(self) -> Dict[str, Any]:
        """Serialize session for persistence."""
        return {
            "model": self.model,
            "provider": self.provider,
            "created_at": self.created_at,
            "total_cost_usd": self.total_cost_usd,
            "total_tokens": self.total_tokens,
            "max_cost_usd": self.max_cost_usd,
            "max_tokens": self.max_tokens,
            "history_max_chars": self.history_max_chars,
            "exchanges": [ex.to_dict() for ex in self.exchanges],
        }

    @classmethod
    def from_dict(cls, d: Dict[str, Any], api_key: Optional[str] = None) -> "Session":
        """Load session from dict."""
        session = cls(
            model=d["model"],
            provider=d.get("provider", "openrouter"),
            api_key=api_key,
            history_max_chars=d.get("history_max_chars", 10000),
            max_cost_usd=d.get("max_cost_usd"),
            max_tokens=d.get("max_tokens"),
        )
        session.created_at = d.get("created_at", datetime.now().isoformat())
        session.total_cost_usd = d.get("total_cost_usd", 0.0)
        session.total_tokens = d.get("total_tokens", 0)
        session.exchanges = [Exchange.from_dict(ex) for ex in d.get("exchanges", [])]
        return session

    def save(self, path: str) -> None:
        """Save session to JSON file."""
        Path(path).write_text(json.dumps(self.to_dict(), indent=2))

    @classmethod
    def load(cls, path: str, api_key: Optional[str] = None) -> "Session":
        """Load session from JSON file."""
        data = json.loads(Path(path).read_text())
        return cls.from_dict(data, api_key=api_key)

    def __len__(self) -> int:
        """Number of exchanges in history."""
        return len(self.exchanges)

    def __repr__(self) -> str:
        return f"Session(model={self.model!r}, exchanges={len(self.exchanges)}, cost=${self.total_cost_usd:.4f})"
