"""Anthropic OAuth provider - direct API calls with Bearer auth.

This provider bypasses litellm to properly handle OAuth tokens
which require Authorization: Bearer header instead of x-api-key.
"""

import json
from typing import Any

import httpx
from loguru import logger

from nanobot.providers.base import LLMProvider, LLMResponse, ToolCallRequest
from nanobot.providers.oauth_utils import get_auth_headers, get_claude_code_system_prefix


class AnthropicOAuthProvider(LLMProvider):
    """
    Anthropic provider using OAuth token authentication.

    Unlike the LiteLLM provider, this calls the Anthropic API directly
    with proper Bearer token authentication for Claude Max/Pro subscriptions.
    """

    ANTHROPIC_API_URL = "https://api.anthropic.com/v1/messages"

    def __init__(
        self,
        oauth_token: str,
        default_model: str = "claude-opus-4-5",
        api_base: str | None = None,
        thinking_budget: int = 0,
    ):
        super().__init__(api_key=None, api_base=api_base)
        self.oauth_token = oauth_token
        self.default_model = default_model
        self.thinking_budget = thinking_budget
        self._client: httpx.AsyncClient | None = None

    def _get_headers(self) -> dict[str, str]:
        """Get request headers with Bearer auth."""
        return get_auth_headers(self.oauth_token, is_oauth=True)

    def _get_api_url(self) -> str:
        """Get API endpoint URL."""
        if self.api_base:
            return f"{self.api_base.rstrip('/')}/v1/messages"
        return self.ANTHROPIC_API_URL

    @staticmethod
    def _normalize_model(model: str) -> str:
        """Normalize model name for the Anthropic API.

        Anthropic model IDs use hyphens (claude-sonnet-4-5), but users often
        write dots (claude-sonnet-4.5). Normalize so both work.
        """
        return model.replace(".", "-")

    async def _get_client(self) -> httpx.AsyncClient:
        """Get or create async HTTP client."""
        if self._client is None:
            self._client = httpx.AsyncClient(
                timeout=httpx.Timeout(300.0, pool=30.0),
            )
        return self._client

    async def _reset_client(self) -> None:
        """Destroy and recreate the HTTP client after connection errors."""
        old = self._client
        self._client = None
        if old:
            try:
                await old.aclose()
            except Exception:
                pass
        logger.warning("Reset httpx client (pool recycled)")

    async def _diagnose_connectivity(self) -> None:
        """Run diagnostics when ConnectTimeout occurs to understand why."""
        import socket
        import asyncio

        # 1. Raw socket test (bypasses httpx entirely)
        try:
            t0 = __import__('time').monotonic()
            s = socket.create_connection(('api.anthropic.com', 443), timeout=10)
            elapsed = __import__('time').monotonic() - t0
            s.close()
            logger.warning(f"DIAG: raw socket connect OK in {elapsed:.3f}s")
        except Exception as e:
            logger.error(f"DIAG: raw socket connect FAILED: {e}")

        # 2. asyncio connect test (same event loop)
        try:
            t0 = __import__('time').monotonic()
            reader, writer = await asyncio.wait_for(
                asyncio.open_connection('api.anthropic.com', 443),
                timeout=10.0,
            )
            elapsed = __import__('time').monotonic() - t0
            writer.close()
            await writer.wait_closed()
            logger.warning(f"DIAG: asyncio connect OK in {elapsed:.3f}s")
        except Exception as e:
            logger.error(f"DIAG: asyncio connect FAILED: {e}")

        # 3. Fresh httpx client test (new pool)
        try:
            t0 = __import__('time').monotonic()
            async with httpx.AsyncClient(timeout=10.0) as fresh:
                r = await fresh.get('https://api.anthropic.com/')
                elapsed = __import__('time').monotonic() - t0
                logger.warning(f"DIAG: fresh httpx OK in {elapsed:.3f}s (status={r.status_code})")
        except Exception as e:
            logger.error(f"DIAG: fresh httpx FAILED: {e}")

        # 4. DNS resolution
        try:
            ips = socket.getaddrinfo('api.anthropic.com', 443)
            logger.warning(f"DIAG: DNS resolved to {len(ips)} entries, first={ips[0][4][0]}")
        except Exception as e:
            logger.error(f"DIAG: DNS FAILED: {e}")

        # 5. Connection pool state of the broken client
        if self._client:
            transport = self._client._transport
            if hasattr(transport, '_pool'):
                pool = transport._pool
                conns = getattr(pool, '_connections', [])
                reqs = getattr(pool, '_requests', [])
                logger.warning(
                    f"DIAG: pool state: {len(conns)} connections, "
                    f"{len(reqs)} pending requests"
                )
                for i, conn in enumerate(conns[:5]):
                    state = getattr(conn, '_state', 'unknown')
                    logger.warning(f"DIAG: conn[{i}] state={state}")

    def _prepare_messages(
        self,
        messages: list[dict[str, Any]]
    ) -> tuple[str | None, list[dict[str, Any]]]:
        """Prepare messages: extract system prompt and convert OpenAI format to Anthropic.

        The agent loop produces messages in OpenAI format:
          - assistant msgs with tool_calls [{type:"function", function:{name, arguments}}]
          - tool role msgs with tool_call_id, name, content

        Anthropic API expects:
          - assistant msgs with content blocks [{type:"tool_use", id, name, input}]
          - user msgs with content blocks [{type:"tool_result", tool_use_id, content}]

        Returns (system_prompt, anthropic_messages)
        """
        system_parts = []
        converted: list[dict[str, Any]] = []

        for msg in messages:
            role = msg.get("role")

            if role == "system":
                system_parts.append(msg.get("content", ""))
                continue

            if role == "assistant" and msg.get("tool_calls"):
                # Convert OpenAI tool_calls to Anthropic content blocks
                content_blocks: list[dict[str, Any]] = []
                # Preserve thinking blocks (list=raw API blocks with signatures, str=legacy)
                rc = msg.get("reasoning_content")
                if isinstance(rc, list):
                    content_blocks.extend(rc)
                elif isinstance(rc, str) and rc:
                    content_blocks.append({"type": "thinking", "thinking": rc})
                text = msg.get("content")
                if text:
                    content_blocks.append({"type": "text", "text": text})
                for tc in msg["tool_calls"]:
                    func = tc.get("function", {})
                    args = func.get("arguments", "{}")
                    if isinstance(args, str):
                        try:
                            args = json.loads(args)
                        except (json.JSONDecodeError, TypeError):
                            args = {}
                    content_blocks.append({
                        "type": "tool_use",
                        "id": tc.get("id", ""),
                        "name": func.get("name", ""),
                        "input": args,
                    })
                converted.append({"role": "assistant", "content": content_blocks})
                continue

            if role == "assistant" and msg.get("reasoning_content"):
                # Plain assistant message with thinking (no tool calls)
                rc = msg["reasoning_content"]
                if isinstance(rc, list):
                    content_blocks = list(rc)
                else:
                    content_blocks = [{"type": "thinking", "thinking": rc}]
                text = msg.get("content")
                if text:
                    content_blocks.append({"type": "text", "text": text})
                converted.append({"role": "assistant", "content": content_blocks})
                continue

            if role == "tool":
                # Convert tool result to Anthropic user message with tool_result block
                tool_result_block = {
                    "type": "tool_result",
                    "tool_use_id": msg.get("tool_call_id", ""),
                    "content": msg.get("content", ""),
                }
                # Merge into previous user message if it already has tool_result blocks
                if converted and converted[-1].get("role") == "user":
                    prev_content = converted[-1].get("content")
                    if isinstance(prev_content, list):
                        prev_content.append(tool_result_block)
                        continue
                converted.append({"role": "user", "content": [tool_result_block]})
                continue

            if role == "user":
                content = msg.get("content", "")
                # Convert OpenAI image_url blocks to Anthropic image blocks
                if isinstance(content, list):
                    content = self._convert_image_blocks(content)
                # Merge text into previous user message if it has tool_result blocks
                # (handles the "Reflect on the results" interleaved message)
                if converted and converted[-1].get("role") == "user":
                    prev_content = converted[-1].get("content")
                    if isinstance(prev_content, list):
                        if isinstance(content, str):
                            prev_content.append({"type": "text", "text": content})
                        elif isinstance(content, list):
                            prev_content.extend(content)
                        continue
                converted.append({"role": role, "content": content})
                continue

            # Pass through other messages (assistant without tool_calls, etc.)
            converted.append(msg)

        system_prompt = "\n\n".join(system_parts)
        return system_prompt, converted

    @staticmethod
    def _convert_image_blocks(content: list[dict[str, Any]]) -> list[dict[str, Any]]:
        """Convert OpenAI image_url blocks to Anthropic image blocks.

        OpenAI format:  {"type": "image_url", "image_url": {"url": "data:mime;base64,DATA"}}
        Anthropic format: {"type": "image", "source": {"type": "base64", "media_type": "mime", "data": "DATA"}}
        """
        converted = []
        for block in content:
            if block.get("type") == "image_url":
                url = block.get("image_url", {}).get("url", "")
                if url.startswith("data:") and ";base64," in url:
                    header, data = url.split(";base64,", 1)
                    media_type = header.removeprefix("data:")
                    converted.append({
                        "type": "image",
                        "source": {"type": "base64", "media_type": media_type, "data": data},
                    })
                else:
                    converted.append({
                        "type": "image",
                        "source": {"type": "url", "url": url},
                    })
            else:
                converted.append(block)
        return converted

    def _convert_tools_to_anthropic(
        self,
        tools: list[dict[str, Any]] | list[Any] | None
    ) -> list[dict[str, Any]] | None:
        """Convert tools to Anthropic API format.

        Supports both function tools (custom) and native tools (Anthropic).
        Function tools are converted to Anthropic format.
        Native tools are passed through unchanged.
        Tool objects (with to_params/to_schema methods) are converted to dicts.
        """
        if not tools:
            return None

        anthropic_tools = []
        for tool in tools:
            # Convert tool objects to dicts first
            if hasattr(tool, 'to_params'):  # Native Anthropic tool
                tool_dict = tool.to_params()
            elif hasattr(tool, 'to_schema'):  # Function tool
                tool_dict = tool.to_schema()
            else:
                tool_dict = tool  # Already a dict

            # Now process the dict
            if tool_dict.get("type") == "function":
                # Convert function tool format
                func = tool_dict["function"]
                anthropic_tools.append({
                    "name": func["name"],
                    "description": func.get("description", ""),
                    "input_schema": func.get("parameters", {"type": "object", "properties": {}})
                })
            else:
                # Pass through native tool format as-is
                # (bash_20250124, text_editor_20250728, computer_20251124, etc.)
                anthropic_tools.append(tool_dict)

        return anthropic_tools if anthropic_tools else None

    async def _make_request(
        self,
        messages: list[dict[str, Any]],
        system: str | None = None,
        model: str = "claude-opus-4-5",
        max_tokens: int = 4096,
        temperature: float = 0.7,
        tools: list[dict[str, Any]] | None = None,
        thinking_budget_override: int | None = None,
        context_management: dict[str, Any] | None = None,
        beta_flags: set[str] | None = None,
    ) -> dict[str, Any]:
        """Make request to Anthropic API."""
        client = await self._get_client()

        # Add cache breakpoints on the last TWO user messages (4-breakpoint strategy):
        # BP3: Second-to-last user message (stable history from previous turn)
        # BP4: Last user message (current turn, will become BP3 next turn)
        # This allows BP3 to reuse what BP4 cached last turn.
        user_indices = [i for i, m in enumerate(messages) if m.get("role") == "user"]

        if len(user_indices) >= 2:
            # BP3: Second-to-last user message
            idx = user_indices[-2]
            msg = messages[idx]
            content = msg["content"]
            if isinstance(content, str):
                messages[idx] = {**msg, "content": [{"type": "text", "text": content, "cache_control": {"type": "ephemeral"}}]}
            elif isinstance(content, list) and content:
                new_content = list(content)
                new_content[-1] = {**new_content[-1], "cache_control": {"type": "ephemeral"}}
                messages[idx] = {**msg, "content": new_content}

        if len(user_indices) >= 1:
            # BP4: Last user message
            idx = user_indices[-1]
            msg = messages[idx]
            content = msg["content"]
            if isinstance(content, str):
                messages[idx] = {**msg, "content": [{"type": "text", "text": content, "cache_control": {"type": "ephemeral"}}]}
            elif isinstance(content, list) and content:
                new_content = list(content)
                new_content[-1] = {**new_content[-1], "cache_control": {"type": "ephemeral"}}
                messages[idx] = {**msg, "content": new_content}

        payload: dict[str, Any] = {
            "model": model,
            "messages": messages,
            "max_tokens": max_tokens,
        }

        # Extended thinking: temperature must be 1 when enabled
        effective_thinking = thinking_budget_override if thinking_budget_override is not None else self.thinking_budget
        if effective_thinking > 0:
            payload["temperature"] = 1
            # max_tokens must exceed budget_tokens
            if max_tokens <= effective_thinking:
                payload["max_tokens"] = effective_thinking + 4096
            payload["thinking"] = {
                "type": "enabled",
                "budget_tokens": effective_thinking,
            }
        else:
            payload["temperature"] = temperature

        if system:
            payload["system"] = [
                {"type": "text", "text": get_claude_code_system_prefix()},
                {"type": "text", "text": system, "cache_control": {"type": "ephemeral", "ttl": "1h"}},
            ]
        else:
            payload["system"] = [
                {"type": "text", "text": get_claude_code_system_prefix()},
            ]

        if tools:
            cached_tools = list(tools)
            cached_tools[-1] = {**cached_tools[-1], "cache_control": {"type": "ephemeral", "ttl": "1h"}}
            payload["tools"] = cached_tools

        if context_management:
            payload["context_management"] = context_management

        edit_types = [e.get("type") for e in (context_management or {}).get("edits", [])]

        # Build headers with beta flags if provided
        headers = self._get_headers()
        if beta_flags:
            # Merge with existing beta header (from OAuth hardcoded flags)
            existing_beta = headers.get("anthropic-beta", "")
            existing_flags = set(existing_beta.split(",")) if existing_beta else set()
            all_flags = existing_flags | beta_flags
            headers["anthropic-beta"] = ",".join(sorted(all_flags))

        logger.info(
            "Anthropic request: model={} max_tokens={} thinking={} tools={} context_mgmt={} beta={}",
            payload.get("model"), payload.get("max_tokens"),
            payload.get("thinking", "disabled"),
            len(payload.get("tools", [])),
            edit_types or "none",
            headers.get("anthropic-beta", "none"),
        )

        # Debug: Log tool names for diagnostic purposes
        if payload.get("tools"):
            tool_names = [t.get("name", "unnamed") for t in payload["tools"]]
            logger.debug(f"Tool names in request: {tool_names}")

        # Debug: Log message structure to diagnose orphaned tool_result errors
        for idx, m in enumerate(payload.get("messages", [])):
            role = m.get("role", "?")
            content = m.get("content", "")
            if isinstance(content, list):
                block_types = [b.get("type", "?") for b in content]
                logger.debug(f"  msg[{idx}] role={role} blocks={block_types}")
            else:
                logger.debug(f"  msg[{idx}] role={role} text={str(content)[:80]}")

        import asyncio
        import time as _time

        max_retries = 3
        base_delay = 2.0  # seconds

        for attempt in range(max_retries + 1):
            _t0 = _time.monotonic()
            try:
                response = await client.post(
                    self._get_api_url(),
                    headers=headers,
                    json=payload,
                )
            except httpx.ConnectTimeout:
                elapsed = _time.monotonic() - _t0
                logger.error(f"ConnectTimeout after {elapsed:.1f}s (attempt {attempt+1}/{max_retries+1})")
                if attempt == 0:
                    await self._diagnose_connectivity()
                await self._reset_client()
                if attempt < max_retries:
                    delay = base_delay * (2 ** attempt)
                    logger.info(f"Retrying in {delay:.1f}s...")
                    await asyncio.sleep(delay)
                    continue
                raise
            except httpx.PoolTimeout:
                elapsed = _time.monotonic() - _t0
                logger.error(f"PoolTimeout after {elapsed:.1f}s (attempt {attempt+1}/{max_retries+1})")
                await self._reset_client()
                if attempt < max_retries:
                    delay = base_delay * (2 ** attempt)
                    logger.info(f"Retrying in {delay:.1f}s...")
                    await asyncio.sleep(delay)
                    continue
                raise
            except (httpx.ConnectError, httpx.TimeoutException) as e:
                elapsed = _time.monotonic() - _t0
                logger.error(f"{type(e).__name__} after {elapsed:.1f}s (attempt {attempt+1}/{max_retries+1})")
                if attempt < max_retries:
                    delay = base_delay * (2 ** attempt)
                    logger.info(f"Retrying in {delay:.1f}s...")
                    await asyncio.sleep(delay)
                    continue
                raise
            elapsed = _time.monotonic() - _t0
            if elapsed > 30:
                logger.warning(f"Anthropic API slow response: {elapsed:.1f}s")

            # Dump rate limit headers for analysis
            try:
                import datetime
                import os
                header_dump = {
                    "timestamp": datetime.datetime.now(datetime.UTC).isoformat(),
                    "status_code": response.status_code,
                    "model": payload.get("model"),
                    "headers": dict(response.headers),
                }
                dump_path = "/root/.nanobot/workspace/api_headers.jsonl"
                with open(dump_path, "a") as f:
                    f.write(json.dumps(header_dump) + "\n")
                # Capture rate limit state for quota-based model switching
                hdrs = response.headers
                rate_limit_state = {
                    "updated_at": datetime.datetime.utcnow().isoformat(),
                    "model": payload.get("model"),
                    "weekly_all_models": float(hdrs["anthropic-ratelimit-unified-7d-utilization"]) if hdrs.get("anthropic-ratelimit-unified-7d-utilization") else None,
                    "weekly_sonnet": float(hdrs["anthropic-ratelimit-unified-7d_sonnet-utilization"]) if hdrs.get("anthropic-ratelimit-unified-7d_sonnet-utilization") else None,
                    "session_5h": float(hdrs["anthropic-ratelimit-unified-5h-utilization"]) if hdrs.get("anthropic-ratelimit-unified-5h-utilization") else None,
                    "weekly_reset": int(hdrs["anthropic-ratelimit-unified-7d-reset"]) if hdrs.get("anthropic-ratelimit-unified-7d-reset") else None,
                    "session_reset": int(hdrs["anthropic-ratelimit-unified-5h-reset"]) if hdrs.get("anthropic-ratelimit-unified-5h-reset") else None,
                    "binding_limit": hdrs.get("anthropic-ratelimit-unified-representative-claim"),
                    "sonnet_fallback": hdrs.get("anthropic-ratelimit-unified-fallback"),
                }
                state_path = "/root/.nanobot/workspace/memory/rate_limits.json"
                os.makedirs(os.path.dirname(state_path), exist_ok=True)
                with open(state_path, "w") as f:
                    json.dump(rate_limit_state, f, indent=2)
            except Exception as e:
                logger.warning("Rate limit header capture failed: {}", e)

            # Retry on 5xx server errors and 429 rate limits
            if response.status_code >= 500 or response.status_code == 429:
                error_text = response.text
                logger.warning(f"Anthropic API {response.status_code} (attempt {attempt+1}/{max_retries+1}): {error_text[:200]}")
                if attempt < max_retries:
                    if response.status_code == 429:
                        retry_after = response.headers.get("retry-after")
                        delay = float(retry_after) if retry_after else base_delay * (2 ** attempt)
                    else:
                        delay = base_delay * (2 ** attempt)
                    logger.info(f"Retrying in {delay:.1f}s...")
                    await asyncio.sleep(delay)
                    continue
                raise Exception(f"Anthropic API error {response.status_code}: {error_text}")

            if response.status_code != 200:
                error_text = response.text
                raise Exception(f"Anthropic API error {response.status_code}: {error_text}")

            return response.json()

        # Should not reach here, but just in case
        raise Exception("Exhausted all retry attempts")

    async def chat(
        self,
        messages: list[dict[str, Any]],
        tools: list[dict[str, Any]] | list[Any] | None = None,
        model: str | None = None,
        max_tokens: int = 4096,
        temperature: float = 0.7,
        thinking_budget: int | None = None,
        context_management: dict[str, Any] | None = None,
    ) -> LLMResponse:
        """Send chat completion request to Anthropic API."""
        model = model or self.default_model

        # Strip provider prefix if present (e.g. "anthropic/claude-opus-4-5" -> "claude-opus-4-5")
        if "/" in model:
            model = model.split("/")[-1]

        # Normalize dots to hyphens (claude-sonnet-4.5 -> claude-sonnet-4-5)
        model = self._normalize_model(model)

        system, prepared_messages = self._prepare_messages(messages)

        # Collect beta flags from native tools BEFORE conversion
        beta_flags: set[str] = set()
        if tools:
            for tool in tools:
                if hasattr(tool, 'beta_flag') and tool.beta_flag:
                    beta_flags.add(tool.beta_flag)

        logger.debug(f"Beta flags collected: {beta_flags} (from {len(tools) if tools else 0} tools)")

        # Convert tools to API format
        anthropic_tools = self._convert_tools_to_anthropic(tools)

        # Per-call thinking override (None = use instance default)
        effective_thinking = self.thinking_budget if thinking_budget is None else thinking_budget

        try:
            response = await self._make_request(
                messages=prepared_messages,
                system=system,
                model=model,
                max_tokens=max_tokens,
                temperature=temperature,
                tools=anthropic_tools,
                thinking_budget_override=effective_thinking,
                context_management=context_management,
                beta_flags=beta_flags,
            )
            return self._parse_response(response)
        except Exception as e:
            logger.exception("Exception in chat():")
            error_msg = f"{type(e).__name__}: {str(e)}" if str(e) else f"{type(e).__name__} (no message)"
            return LLMResponse(
                content=f"Error calling LLM: {error_msg}",
                finish_reason="error",
            )

    def _parse_response(self, response: dict[str, Any]) -> LLMResponse:
        """Parse Anthropic API response."""
        content_blocks = response.get("content", [])

        text_content = ""
        thinking_blocks: list[dict[str, Any]] = []
        tool_calls = []

        for block in content_blocks:
            if block.get("type") == "thinking":
                # Preserve full block including signature for multi-turn replay
                thinking_blocks.append(block)
            elif block.get("type") == "text":
                text_content += block.get("text", "")
            elif block.get("type") == "tool_use":
                tool_calls.append(ToolCallRequest(
                    id=block.get("id", ""),
                    name=block.get("name", ""),
                    arguments=block.get("input", {}),
                ))

        usage = {}
        if "usage" in response:
            usage = {
                "prompt_tokens": response["usage"].get("input_tokens", 0),
                "completion_tokens": response["usage"].get("output_tokens", 0),
                "total_tokens": (
                    response["usage"].get("input_tokens", 0) +
                    response["usage"].get("output_tokens", 0)
                ),
            }

        stop_reason = response.get("stop_reason", "end_turn")
        thinking_chars = sum(len(b.get("thinking", "")) for b in thinking_blocks) if thinking_blocks else 0
        raw_usage = response.get("usage", {})
        cache_write = raw_usage.get("cache_creation_input_tokens", 0)
        cache_read = raw_usage.get("cache_read_input_tokens", 0)
        logger.info(
            "Anthropic response: stop={} tool_calls={} thinking={} chars, "
            "input={} output={} cache_write={} cache_read={} tokens",
            stop_reason, len(tool_calls), thinking_chars,
            usage.get("prompt_tokens", 0), usage.get("completion_tokens", 0),
            cache_write, cache_read,
        )

        # Log context editing activity if any edits were applied
        if applied_edits := response.get("context_management", {}).get("applied_edits"):
            for edit in applied_edits:
                edit_type = edit.get("type", "?")
                cleared_tokens = edit.get("cleared_input_tokens", 0)
                if edit_type == "clear_tool_uses_20250919":
                    logger.info(
                        "Context edit: cleared {} tool uses ({} tokens)",
                        edit.get("cleared_tool_uses", 0), cleared_tokens,
                    )
                elif edit_type == "clear_thinking_20251015":
                    logger.info(
                        "Context edit: cleared {} thinking turns ({} tokens)",
                        edit.get("cleared_thinking_turns", 0), cleared_tokens,
                    )

        return LLMResponse(
            content=text_content or None,
            tool_calls=tool_calls,
            finish_reason=stop_reason,
            usage=usage,
            reasoning_content=thinking_blocks or None,
        )

    def get_default_model(self) -> str:
        """Get the default model."""
        return self.default_model

    async def close(self):
        """Close the HTTP client."""
        if self._client:
            await self._client.aclose()
            self._client = None
