"""Claude SDK wrapper for Promethian Agent."""

from __future__ import annotations

import json
import os
import re
import time
from dataclasses import dataclass, field
from typing import Any, TYPE_CHECKING

import anthropic

if TYPE_CHECKING:
    from promethian.utils.logging import VerboseLogger


# Model pricing per 1M tokens (as of 2025)
# cache_write: cost to write to cache (first time)
# cache_read: cost to read from cache (subsequent times, ~90% cheaper)
MODEL_PRICING = {
    "claude-opus-4-5-20251101": {
        "input": 15.00, "output": 75.00,
        "cache_write": 18.75, "cache_read": 1.50,  # 1.25x write, 0.1x read
    },
    "claude-opus-4-20250514": {
        "input": 15.00, "output": 75.00,
        "cache_write": 18.75, "cache_read": 1.50,
    },
    "claude-sonnet-4-20250514": {
        "input": 3.00, "output": 15.00,
        "cache_write": 3.75, "cache_read": 0.30,
    },
    "claude-haiku-3-5-20241022": {
        "input": 0.80, "output": 4.00,
        "cache_write": 1.00, "cache_read": 0.08,
    },
    # Fallback for unknown models
    "default": {"input": 3.00, "output": 15.00, "cache_write": 3.75, "cache_read": 0.30},
}


@dataclass
class ModelUsage:
    """Token usage for a specific model."""
    model: str
    input_tokens: int = 0
    output_tokens: int = 0
    api_calls: int = 0
    # Cache statistics
    cache_creation_input_tokens: int = 0  # Tokens written to cache
    cache_read_input_tokens: int = 0  # Tokens read from cache

    @property
    def total_tokens(self) -> int:
        return self.input_tokens + self.output_tokens

    @property
    def cache_hit_rate(self) -> float:
        """Percentage of input tokens served from cache."""
        total_cacheable = self.cache_creation_input_tokens + self.cache_read_input_tokens
        if total_cacheable == 0:
            return 0.0
        return self.cache_read_input_tokens / total_cacheable

    @property
    def estimated_cost(self) -> float:
        """Estimate cost in USD based on token usage, accounting for caching."""
        pricing = MODEL_PRICING.get(self.model, MODEL_PRICING["default"])

        # Regular input tokens (non-cached)
        regular_input = self.input_tokens - self.cache_creation_input_tokens - self.cache_read_input_tokens
        input_cost = (max(0, regular_input) / 1_000_000) * pricing["input"]

        # Cache write cost (1.25x input price)
        cache_write_cost = (self.cache_creation_input_tokens / 1_000_000) * pricing.get("cache_write", pricing["input"])

        # Cache read cost (0.1x input price - the big savings!)
        cache_read_cost = (self.cache_read_input_tokens / 1_000_000) * pricing.get("cache_read", pricing["input"])

        # Output cost
        output_cost = (self.output_tokens / 1_000_000) * pricing["output"]

        return input_cost + cache_write_cost + cache_read_cost + output_cost

    @property
    def estimated_savings(self) -> float:
        """Estimated savings from caching vs no caching."""
        if self.cache_read_input_tokens == 0:
            return 0.0
        pricing = MODEL_PRICING.get(self.model, MODEL_PRICING["default"])
        # What we would have paid without caching
        full_price = (self.cache_read_input_tokens / 1_000_000) * pricing["input"]
        # What we actually paid
        cached_price = (self.cache_read_input_tokens / 1_000_000) * pricing.get("cache_read", pricing["input"])
        return full_price - cached_price


@dataclass
class UsageStats:
    """Aggregated usage statistics across all models."""
    by_model: dict[str, ModelUsage] = field(default_factory=dict)

    @property
    def total_input_tokens(self) -> int:
        return sum(m.input_tokens for m in self.by_model.values())

    @property
    def total_output_tokens(self) -> int:
        return sum(m.output_tokens for m in self.by_model.values())

    @property
    def total_tokens(self) -> int:
        return self.total_input_tokens + self.total_output_tokens

    @property
    def total_api_calls(self) -> int:
        return sum(m.api_calls for m in self.by_model.values())

    @property
    def estimated_cost(self) -> float:
        """Total estimated cost in USD."""
        return sum(m.estimated_cost for m in self.by_model.values())

    @property
    def total_cache_read_tokens(self) -> int:
        return sum(m.cache_read_input_tokens for m in self.by_model.values())

    @property
    def total_cache_write_tokens(self) -> int:
        return sum(m.cache_creation_input_tokens for m in self.by_model.values())

    @property
    def total_savings(self) -> float:
        return sum(m.estimated_savings for m in self.by_model.values())

    def record(
        self,
        model: str,
        input_tokens: int,
        output_tokens: int,
        cache_creation_input_tokens: int = 0,
        cache_read_input_tokens: int = 0,
    ) -> None:
        """Record usage for a model."""
        if model not in self.by_model:
            self.by_model[model] = ModelUsage(model=model)
        self.by_model[model].input_tokens += input_tokens
        self.by_model[model].output_tokens += output_tokens
        self.by_model[model].api_calls += 1
        self.by_model[model].cache_creation_input_tokens += cache_creation_input_tokens
        self.by_model[model].cache_read_input_tokens += cache_read_input_tokens

    def merge(self, other: "UsageStats") -> None:
        """Merge another UsageStats into this one."""
        for model, usage in other.by_model.items():
            if model not in self.by_model:
                self.by_model[model] = ModelUsage(model=model)
            self.by_model[model].input_tokens += usage.input_tokens
            self.by_model[model].output_tokens += usage.output_tokens
            self.by_model[model].api_calls += usage.api_calls
            self.by_model[model].cache_creation_input_tokens += usage.cache_creation_input_tokens
            self.by_model[model].cache_read_input_tokens += usage.cache_read_input_tokens

    def to_dict(self) -> dict:
        """Convert to dictionary for serialization."""
        result = {
            "total_input_tokens": self.total_input_tokens,
            "total_output_tokens": self.total_output_tokens,
            "total_tokens": self.total_tokens,
            "total_api_calls": self.total_api_calls,
            "estimated_cost_usd": round(self.estimated_cost, 4),
            "by_model": {
                model: {
                    "input_tokens": usage.input_tokens,
                    "output_tokens": usage.output_tokens,
                    "api_calls": usage.api_calls,
                    "estimated_cost_usd": round(usage.estimated_cost, 4),
                    "cache_read_tokens": usage.cache_read_input_tokens,
                    "cache_write_tokens": usage.cache_creation_input_tokens,
                }
                for model, usage in self.by_model.items()
            },
        }
        # Add cache savings if any caching occurred
        if self.total_cache_read_tokens > 0:
            result["cache_savings_usd"] = round(self.total_savings, 4)
        return result


class ClaudeClient:
    """
    Wrapper around the Anthropic Claude SDK.

    Handles completions, structured output, and tool use.
    """

    def __init__(
        self,
        api_key: str | None = None,
        model: str = "claude-opus-4-5-20251101",
        max_tokens: int = 8192,
        track_usage: bool = True,
    ):
        self.api_key = api_key or os.environ.get("ANTHROPIC_API_KEY")
        if not self.api_key:
            raise ValueError(
                "ANTHROPIC_API_KEY is required. "
                "Set it in your environment or pass it to the client."
            )

        self.client = anthropic.Anthropic(api_key=self.api_key)
        self.model = model
        self.max_tokens = max_tokens
        self.track_usage = track_usage
        self._usage = UsageStats()
        self._logger: "VerboseLogger | None" = None

    def set_logger(self, logger: "VerboseLogger") -> None:
        """Set the verbose logger for API call logging."""
        self._logger = logger

    @property
    def usage(self) -> UsageStats:
        """Get current usage statistics."""
        return self._usage

    def reset_usage(self) -> None:
        """Reset usage statistics."""
        self._usage = UsageStats()

    def _record_usage(self, response: Any, model: str, latency_ms: int = 0, purpose: str = "") -> None:
        """Record usage from an API response, including cache stats."""
        if self.track_usage and hasattr(response, "usage"):
            # Extract cache tokens if present
            cache_creation = getattr(response.usage, "cache_creation_input_tokens", 0) or 0
            cache_read = getattr(response.usage, "cache_read_input_tokens", 0) or 0

            self._usage.record(
                model=model,
                input_tokens=response.usage.input_tokens,
                output_tokens=response.usage.output_tokens,
                cache_creation_input_tokens=cache_creation,
                cache_read_input_tokens=cache_read,
            )

            # Log API call if logger is set
            if self._logger:
                # Calculate cost for this call
                pricing = MODEL_PRICING.get(model, MODEL_PRICING["default"])
                regular_input = response.usage.input_tokens - cache_creation - cache_read
                cost = (
                    (max(0, regular_input) / 1_000_000) * pricing["input"]
                    + (cache_creation / 1_000_000) * pricing.get("cache_write", pricing["input"])
                    + (cache_read / 1_000_000) * pricing.get("cache_read", pricing["input"])
                    + (response.usage.output_tokens / 1_000_000) * pricing["output"]
                )

                self._logger.api_call(
                    model=model,
                    input_tokens=response.usage.input_tokens,
                    output_tokens=response.usage.output_tokens,
                    latency_ms=latency_ms,
                    cache_read_tokens=cache_read,
                    cache_write_tokens=cache_creation,
                    cost_usd=cost,
                    purpose=purpose,
                )

    @staticmethod
    def _make_cacheable_system(system: str) -> list[dict]:
        """Convert system prompt to cacheable format."""
        return [{
            "type": "text",
            "text": system,
            "cache_control": {"type": "ephemeral"}
        }]

    @staticmethod
    def _make_cacheable_content(cacheable_text: str, dynamic_text: str = "") -> list[dict]:
        """Create message content with cacheable and dynamic parts."""
        content = [{
            "type": "text",
            "text": cacheable_text,
            "cache_control": {"type": "ephemeral"}
        }]
        if dynamic_text:
            content.append({
                "type": "text",
                "text": dynamic_text
            })
        return content

    async def complete(
        self,
        prompt: str,
        system: str | None = None,
        model: str | None = None,
        max_tokens: int | None = None,
        temperature: float = 0.7,
        cache_system: bool = True,
    ) -> str:
        """
        Get a text completion from Claude.

        Args:
            prompt: The user prompt
            system: Optional system prompt
            model: Model to use (defaults to self.model)
            max_tokens: Max tokens to generate
            temperature: Sampling temperature
            cache_system: Whether to cache the system prompt (default True)
        """
        messages = [{"role": "user", "content": prompt}]
        used_model = model or self.model

        kwargs = {
            "model": used_model,
            "max_tokens": max_tokens or self.max_tokens,
            "messages": messages,
            "temperature": temperature,
        }

        if system:
            # Use cached system prompt format for efficiency
            if cache_system:
                kwargs["system"] = self._make_cacheable_system(system)
            else:
                kwargs["system"] = system

        start_time = time.time()
        response = self.client.messages.create(**kwargs)
        latency_ms = int((time.time() - start_time) * 1000)
        self._record_usage(response, used_model, latency_ms=latency_ms, purpose="complete")

        return response.content[0].text

    async def complete_json(
        self,
        prompt: str,
        system: str | None = None,
        model: str | None = None,
        max_tokens: int | None = None,
    ) -> dict:
        """
        Get a JSON-structured completion from Claude.

        Instructs Claude to return valid JSON and parses the response.
        """
        full_prompt = f"""{prompt}

Respond with valid JSON only. No markdown code blocks, no explanation, just the JSON object."""

        response = await self.complete(
            prompt=full_prompt,
            system=system,
            model=model,
            max_tokens=max_tokens,
            temperature=0.3,  # Lower temperature for structured output
        )

        return self._parse_json(response)

    def _parse_json(self, text: str) -> dict:
        """Parse JSON from response, handling markdown code blocks."""
        json_str = text.strip()

        # Remove markdown code blocks if present
        if json_str.startswith("```"):
            json_str = re.sub(r"^```(?:json)?\n?", "", json_str)
            json_str = re.sub(r"\n?```$", "", json_str)

        return json.loads(json_str)

    async def complete_with_tools(
        self,
        prompt: str,
        tools: list[dict],
        system: str | None = None,
        model: str | None = None,
        max_tokens: int | None = None,
        cache_system: bool = True,
        cacheable_context: str | None = None,
    ) -> tuple[str, list[dict]]:
        """
        Get a completion with tool use from Claude.

        Args:
            prompt: The user prompt (dynamic part)
            tools: List of tool definitions
            system: Optional system prompt
            model: Model to use
            max_tokens: Max tokens to generate
            cache_system: Whether to cache the system prompt
            cacheable_context: Static context to cache (e.g., tool descriptions, skills)
                              This gets cached separately from the dynamic prompt.

        Returns (text_content, tool_calls) where tool_calls is a list of
        {id, name, input} dictionaries.
        """
        used_model = model or self.model

        # Build message content - with caching if cacheable_context provided
        if cacheable_context:
            # Split into cached (static) and uncached (dynamic) parts
            message_content = self._make_cacheable_content(cacheable_context, prompt)
            messages = [{"role": "user", "content": message_content}]
        else:
            messages = [{"role": "user", "content": prompt}]

        kwargs = {
            "model": used_model,
            "max_tokens": max_tokens or self.max_tokens,
            "messages": messages,
            "tools": tools,
        }

        if system:
            if cache_system:
                kwargs["system"] = self._make_cacheable_system(system)
            else:
                kwargs["system"] = system

        start_time = time.time()
        response = self.client.messages.create(**kwargs)
        latency_ms = int((time.time() - start_time) * 1000)
        self._record_usage(response, used_model, latency_ms=latency_ms, purpose="complete_with_tools")

        # Extract text and tool calls
        text_parts = []
        tool_calls = []

        for block in response.content:
            if block.type == "text":
                text_parts.append(block.text)
            elif block.type == "tool_use":
                tool_calls.append({
                    "id": block.id,
                    "name": block.name,
                    "input": block.input,
                })

        return "\n".join(text_parts), tool_calls

    async def continue_with_tool_result(
        self,
        messages: list[dict],
        tool_results: list[dict],
        tools: list[dict],
        system: str | None = None,
        model: str | None = None,
        max_tokens: int | None = None,
        cache_system: bool = True,
    ) -> tuple[str, list[dict]]:
        """
        Continue a conversation after tool execution.

        tool_results should be a list of {tool_use_id, content} dicts.
        """
        # Add tool results as assistant turn
        tool_result_content = [
            {
                "type": "tool_result",
                "tool_use_id": r["tool_use_id"],
                "content": r["content"],
            }
            for r in tool_results
        ]

        messages = messages + [{"role": "user", "content": tool_result_content}]
        used_model = model or self.model

        kwargs = {
            "model": used_model,
            "max_tokens": max_tokens or self.max_tokens,
            "messages": messages,
            "tools": tools,
        }

        if system:
            if cache_system:
                kwargs["system"] = self._make_cacheable_system(system)
            else:
                kwargs["system"] = system

        start_time = time.time()
        response = self.client.messages.create(**kwargs)
        latency_ms = int((time.time() - start_time) * 1000)
        self._record_usage(response, used_model, latency_ms=latency_ms, purpose="continue_with_tool_result")

        # Extract text and tool calls
        text_parts = []
        tool_calls = []

        for block in response.content:
            if block.type == "text":
                text_parts.append(block.text)
            elif block.type == "tool_use":
                tool_calls.append({
                    "id": block.id,
                    "name": block.name,
                    "input": block.input,
                })

        return "\n".join(text_parts), tool_calls
