"""Shared Claude Agent SDK client wrapper."""

from __future__ import annotations

import asyncio
import logging
import os
import shutil
import tempfile
from dataclasses import dataclass
from importlib import import_module
from pathlib import Path
from typing import Any

logger = logging.getLogger(__name__)


@dataclass
class ClaudeCompletion:
    """Normalized completion result from Claude Agent SDK."""

    text: str
    input_tokens: int = 0
    output_tokens: int = 0
    raw_usage: dict[str, Any] | None = None


def _load_sdk() -> tuple[Any, Any]:
    """Load Claude Agent SDK symbols."""
    try:
        mod = import_module("claude_agent_sdk")
        return mod.query, mod.ClaudeAgentOptions
    except Exception as exc:
        raise ImportError(
            "Claude Agent SDK is not installed. Install 'claude-agent-sdk'."
        ) from exc


class ClaudeSDKClient:
    """Small adapter over Claude Agent SDK query() streaming API."""

    _RATE_LIMIT_RETRY_TIME_MARKS_SECONDS: tuple[float, ...] = (0.2, 1.0, 5.0)
    _SENSITIVE_KEY_TOKENS: tuple[str, ...] = (
        "api_key",
        "apikey",
        "auth",
        "token",
        "secret",
        "password",
        "authorization",
        "cookie",
    )

    def __init__(
        self,
        api_key: str | None = None,
        auth_token: str | None = None,
        enable_debug: bool = False,
    ):
        self._query: Any | None = None
        self._options_cls: Any | None = None
        self._env: dict[str, str] = {}
        self._enable_debug = enable_debug
        if api_key:
            self._env["ANTHROPIC_API_KEY"] = api_key
        if auth_token:
            self._env["ANTHROPIC_AUTH_TOKEN"] = auth_token

    async def complete(
        self,
        prompt: str,
        *,
        model: str | None = None,
        system_prompt: str | None = None,
        max_turns: int = 100,
        cwd: str | None = None,
        env: dict[str, str] | None = None,
        observability: Any | None = None,
        agent_name: str = "claude_sdk",
        task_id: str = "-",
    ) -> ClaudeCompletion:
        """Run a single-turn completion and normalize text/token usage."""
        self._ensure_sdk_loaded()
        self._emit_observability_event(
            observability,
            agent_name,
            task_id,
            "request_start",
            {
                "model": model,
                "max_turns": max_turns,
                "cwd": cwd,
                "prompt_chars": len(prompt),
                "system_prompt_chars": len(system_prompt) if system_prompt else 0,
            },
        )
        options_kwargs: dict[str, Any] = {"max_turns": max_turns}
        if model:
            options_kwargs["model"] = model
        if system_prompt:
            options_kwargs["system_prompt"] = system_prompt
        if cwd:
            options_kwargs["cwd"] = cwd

        effective_env = dict(self._env)
        if env:
            effective_env.update(env)
        effective_env = self._ensure_claude_home_writable(effective_env, cwd=cwd)
        if effective_env:
            options_kwargs["env"] = effective_env
        total_attempts = len(self._RATE_LIMIT_RETRY_TIME_MARKS_SECONDS) + 1
        for attempt in range(total_attempts):
            self._emit_observability_event(
                observability,
                agent_name,
                task_id,
                "attempt_start",
                {"attempt": attempt + 1, "total_attempts": total_attempts},
            )
            debug_stderr = None
            if self._enable_debug:
                debug_stderr = tempfile.TemporaryFile(mode="w+t", encoding="utf-8")
            attempt_options_kwargs = dict(options_kwargs)
            if debug_stderr is not None:
                attempt_options_kwargs["debug_stderr"] = debug_stderr
                attempt_options_kwargs["extra_args"] = {"debug-to-stderr": None}
            options = self._options_cls(**attempt_options_kwargs)
            assistant_parts: list[str] = []
            result_text: str | None = None
            usage: dict[str, Any] | None = None
            error_text: str | None = None
            result_subtype: str | None = None
            session_id: str | None = None
            stderr_detail = ""

            try:
                async for msg in self._query(prompt=prompt, options=options):
                    session_id = self._record_stream_message(
                        msg=msg,
                        observability=observability,
                        agent_name=agent_name,
                        task_id=task_id,
                        current_session_id=session_id,
                    )
                    content = getattr(msg, "content", None)
                    # Only assistant messages contain model output content.
                    if content and hasattr(msg, "model"):
                        for block in content:
                            text = getattr(block, "text", None)
                            if text:
                                assistant_parts.append(text)

                    msg_result = getattr(msg, "result", None)
                    if isinstance(msg_result, str) and msg_result.strip():
                        result_text = msg_result

                    msg_subtype = getattr(msg, "subtype", None)
                    if isinstance(msg_subtype, str):
                        result_subtype = msg_subtype

                    msg_usage = getattr(msg, "usage", None)
                    if isinstance(msg_usage, dict):
                        usage = msg_usage

                    if getattr(msg, "is_error", False):
                        error_text = msg_result if isinstance(msg_result, str) else "Claude SDK error"
                stderr_detail = self._combine_stderr_details(self._read_debug_stderr(debug_stderr))
            except Exception as exc:
                stderr_detail = self._combine_stderr_details(
                    self._read_debug_stderr(debug_stderr),
                    self._extract_exception_stderr(exc),
                )
                error_message = self._format_error(
                    f"Claude SDK query failed: {exc}",
                    stderr_detail,
                    add_hint=True,
                )
                should_retry = await self._should_retry_rate_limit_error(error_message, attempt)
                self._emit_observability_event(
                    observability,
                    agent_name,
                    task_id,
                    "request_error",
                    {
                        "attempt": attempt + 1,
                        "error": self._truncate_text(str(exc)),
                        "retrying": should_retry,
                        "stderr": self._truncate_text(stderr_detail),
                    },
                )
                if should_retry:
                    continue
                raise RuntimeError(error_message) from exc
            finally:
                if debug_stderr is not None:
                    debug_stderr.close()

            if error_text:
                error_message = self._format_error(error_text, stderr_detail, add_hint=True)
                should_retry = await self._should_retry_rate_limit_error(error_message, attempt)
                self._emit_observability_event(
                    observability,
                    agent_name,
                    task_id,
                    "request_error",
                    {
                        "attempt": attempt + 1,
                        "error": self._truncate_text(error_text),
                        "retrying": should_retry,
                        "stderr": self._truncate_text(stderr_detail),
                    },
                )
                if should_retry:
                    continue
                raise RuntimeError(error_message)

            if result_subtype and "error" in result_subtype.lower():
                error_message = self._format_error(
                    f"Claude SDK execution ended with subtype '{result_subtype}'.",
                    stderr_detail,
                    add_hint=True,
                )
                should_retry = await self._should_retry_rate_limit_error(error_message, attempt)
                self._emit_observability_event(
                    observability,
                    agent_name,
                    task_id,
                    "request_error",
                    {
                        "attempt": attempt + 1,
                        "error": f"result subtype={result_subtype}",
                        "retrying": should_retry,
                        "stderr": self._truncate_text(stderr_detail),
                    },
                )
                if should_retry:
                    continue
                raise RuntimeError(error_message)

            text = (result_text or "\n".join(assistant_parts)).strip()
            if not text:
                error_message = self._format_error(
                    "Claude SDK returned empty response",
                    stderr_detail,
                    add_hint=True,
                )
                should_retry = await self._should_retry_rate_limit_error(error_message, attempt)
                self._emit_observability_event(
                    observability,
                    agent_name,
                    task_id,
                    "request_error",
                    {
                        "attempt": attempt + 1,
                        "error": "empty response",
                        "retrying": should_retry,
                        "stderr": self._truncate_text(stderr_detail),
                    },
                )
                if should_retry:
                    continue
                raise RuntimeError(error_message)

            input_tokens, output_tokens = self._extract_token_counts(usage)
            self._emit_observability_event(
                observability,
                agent_name,
                task_id,
                "request_complete",
                {
                    "attempt": attempt + 1,
                    "session_id": session_id,
                    "result_subtype": result_subtype,
                    "result_preview": self._truncate_text(text, max_chars=180),
                    "input_tokens": input_tokens,
                    "output_tokens": output_tokens,
                    "usage": self._sanitize_payload(usage),
                },
            )
            return ClaudeCompletion(
                text=text,
                input_tokens=input_tokens,
                output_tokens=output_tokens,
                raw_usage=usage,
            )

        raise RuntimeError("Claude SDK retry loop exhausted unexpectedly")

    def _ensure_sdk_loaded(self) -> None:
        if self._query is not None and self._options_cls is not None:
            return
        self._query, self._options_cls = _load_sdk()

    def _extract_token_counts(self, usage: dict[str, Any] | None) -> tuple[int, int]:
        """Best-effort token extraction across SDK usage payload variants."""
        if not isinstance(usage, dict):
            return 0, 0

        input_tokens = self._to_int(
            usage.get("input_tokens") or usage.get("inputTokens")
        )
        output_tokens = self._to_int(
            usage.get("output_tokens") or usage.get("outputTokens")
        )

        if input_tokens == 0:
            input_tokens = sum(
                self._to_int(v)
                for k, v in usage.items()
                if "input" in k.lower() and "output" not in k.lower()
            )

        if output_tokens == 0:
            output_tokens = sum(
                self._to_int(v) for k, v in usage.items() if "output" in k.lower()
            )

        return input_tokens, output_tokens

    def _ensure_claude_home_writable(
        self,
        env: dict[str, str],
        *,
        cwd: str | None = None,
    ) -> dict[str, str]:
        """Fallback to a project-local HOME when ~/.claude paths are not writable."""
        effective = dict(env)
        current_home = Path(effective.get("HOME") or str(Path.home())).expanduser()

        if self._claude_home_is_writable(current_home):
            return effective

        fallback_root = Path(cwd or os.getcwd()) / ".app_factory" / "claude_home"
        fallback_home = self._prepare_fallback_claude_home(
            source_home=current_home,
            fallback_home=fallback_root,
        )
        effective["HOME"] = str(fallback_home)
        logger.warning(
            "Claude home '%s' is not writable; using fallback HOME at '%s'.",
            current_home,
            fallback_home,
        )
        return effective

    @staticmethod
    def _claude_home_is_writable(home: Path) -> bool:
        claude_dir = home / ".claude"
        required_dirs = [claude_dir, claude_dir / "todos", claude_dir / "debug"]
        config_file = home / ".claude.json"

        try:
            for directory in required_dirs:
                directory.mkdir(parents=True, exist_ok=True)
                probe = directory / ".app_factory_write_probe"
                probe.write_text("ok", encoding="utf-8")
                probe.unlink()

            config_file.touch(exist_ok=True)
            with open(config_file, "a", encoding="utf-8"):
                pass
        except OSError:
            return False

        return True

    @staticmethod
    def _prepare_fallback_claude_home(source_home: Path, fallback_home: Path) -> Path:
        fallback_home.mkdir(parents=True, exist_ok=True)

        fallback_claude_dir = fallback_home / ".claude"
        (fallback_claude_dir / "todos").mkdir(parents=True, exist_ok=True)
        (fallback_claude_dir / "debug").mkdir(parents=True, exist_ok=True)

        source_claude_dir = source_home / ".claude"
        source_config = source_home / ".claude.json"
        target_config = fallback_home / ".claude.json"

        if source_config.exists() and source_config.is_file() and os.access(source_config, os.R_OK):
            try:
                shutil.copy2(source_config, target_config)
            except OSError:
                pass
        else:
            target_config.touch(exist_ok=True)

        if source_claude_dir.exists() and source_claude_dir.is_dir() and os.access(
            source_claude_dir, os.R_OK
        ):
            try:
                shutil.copytree(source_claude_dir, fallback_claude_dir, dirs_exist_ok=True)
            except OSError:
                pass

        return fallback_home

    @staticmethod
    def _to_int(value: Any) -> int:
        try:
            return int(value)
        except (TypeError, ValueError):
            return 0

    @staticmethod
    def _read_debug_stderr(debug_stderr: Any) -> str:
        if debug_stderr is None:
            return ""
        try:
            debug_stderr.flush()
            debug_stderr.seek(0)
            value = debug_stderr.read()
            if isinstance(value, str):
                return value.strip()
        except Exception:
            pass
        return ""

    @staticmethod
    def _extract_exception_stderr(exc: Exception) -> str:
        stderr = getattr(exc, "stderr", None)
        return stderr.strip() if isinstance(stderr, str) else ""

    @staticmethod
    def _combine_stderr_details(*details: str) -> str:
        merged: list[str] = []
        seen: set[str] = set()
        for detail in details:
            value = detail.strip() if isinstance(detail, str) else ""
            if not value or value in seen:
                continue
            seen.add(value)
            merged.append(value)

        if not merged:
            return ""

        placeholder = "Check stderr output for details"
        non_placeholder = [detail for detail in merged if placeholder not in detail]
        preferred = non_placeholder if non_placeholder else merged
        return "\n\n".join(preferred)

    async def _should_retry_rate_limit_error(self, error_message: str, attempt: int) -> bool:
        if attempt >= len(self._RATE_LIMIT_RETRY_TIME_MARKS_SECONDS):
            return False

        text = error_message.lower()
        retryable_tokens = (
            "rate_limit_event",
            "rate limit",
            "rate-limited",
            "too many requests",
            "status code: 429",
            "status code 429",
        )
        if not any(token in text for token in retryable_tokens):
            return False

        time_marks = self._RATE_LIMIT_RETRY_TIME_MARKS_SECONDS
        target_mark = time_marks[attempt]
        previous_mark = time_marks[attempt - 1] if attempt > 0 else 0.0
        delay = max(target_mark - previous_mark, 0.0)
        logger.warning(
            "Claude SDK rate limit/transient event detected (attempt %d/%d). "
            "Retrying in %.1fs (target %.1fs from first failure).",
            attempt + 1,
            len(self._RATE_LIMIT_RETRY_TIME_MARKS_SECONDS) + 1,
            delay,
            target_mark,
        )
        await asyncio.sleep(delay)
        return True

    @staticmethod
    def _format_error(message: str, stderr_detail: str, add_hint: bool = False) -> str:
        hint = ""
        if add_hint:
            hint = (
                "\nHint: verify Claude auth is available (ANTHROPIC_API_KEY or "
                "ANTHROPIC_AUTH_TOKEN, or a valid Claude Code OAuth session) and that the "
                "process can write ~/.claude and ~/.claude.json."
            )
        if stderr_detail:
            return f"{message}\nSDK stderr:\n{stderr_detail}{hint}"
        return f"{message}{hint}"

    def _record_stream_message(
        self,
        *,
        msg: Any,
        observability: Any | None,
        agent_name: str,
        task_id: str,
        current_session_id: str | None,
    ) -> str | None:
        session_id = getattr(msg, "session_id", None) or current_session_id
        parent_tool_use_id = getattr(msg, "parent_tool_use_id", None)

        stream_event = getattr(msg, "event", None)
        if stream_event is not None and hasattr(msg, "uuid"):
            stream_event_type = None
            if isinstance(stream_event, dict):
                stream_event_type = stream_event.get("type") or stream_event.get("event")
            self._emit_observability_event(
                observability,
                agent_name,
                task_id,
                "stream_event",
                {
                    "session_id": session_id,
                    "stream_event_type": stream_event_type,
                    "parent_tool_use_id": parent_tool_use_id,
                },
            )
            return session_id

        content = getattr(msg, "content", None)
        if content:
            is_assistant_message = hasattr(msg, "model")
            self._emit_observability_event(
                observability,
                agent_name,
                task_id,
                "assistant_message" if is_assistant_message else "user_message",
                {
                    "session_id": session_id,
                    "parent_tool_use_id": parent_tool_use_id,
                    "model": getattr(msg, "model", None),
                    "content_block_count": len(content) if isinstance(content, list) else 1,
                },
            )
            if isinstance(content, list):
                for block in content:
                    self._record_content_block(
                        block=block,
                        observability=observability,
                        agent_name=agent_name,
                        task_id=task_id,
                        session_id=session_id,
                        parent_tool_use_id=parent_tool_use_id,
                    )

        subtype = getattr(msg, "subtype", None)
        if isinstance(subtype, str):
            if hasattr(msg, "duration_ms"):
                self._emit_observability_event(
                    observability,
                    agent_name,
                    task_id,
                    "result_message",
                    {
                        "session_id": session_id,
                        "subtype": subtype,
                        "is_error": bool(getattr(msg, "is_error", False)),
                        "num_turns": getattr(msg, "num_turns", None),
                        "duration_ms": getattr(msg, "duration_ms", None),
                        "duration_api_ms": getattr(msg, "duration_api_ms", None),
                        "total_cost_usd": getattr(msg, "total_cost_usd", None),
                        "usage": self._sanitize_payload(getattr(msg, "usage", None)),
                    },
                )
            elif hasattr(msg, "data"):
                self._emit_observability_event(
                    observability,
                    agent_name,
                    task_id,
                    "system_message",
                    {
                        "session_id": session_id,
                        "subtype": subtype,
                        "data": self._sanitize_payload(getattr(msg, "data", None)),
                    },
                )

        return session_id

    def _record_content_block(
        self,
        *,
        block: Any,
        observability: Any | None,
        agent_name: str,
        task_id: str,
        session_id: str | None,
        parent_tool_use_id: str | None,
    ) -> None:
        block_name = getattr(block, "name", None)
        block_input = getattr(block, "input", None)
        block_id = getattr(block, "id", None)
        if block_name is not None and block_input is not None and block_id is not None:
            self._emit_observability_event(
                observability,
                agent_name,
                task_id,
                "tool_use",
                {
                    "session_id": session_id,
                    "tool_use_id": block_id,
                    "parent_tool_use_id": parent_tool_use_id,
                    "tool_name": str(block_name),
                    "tool_input": self._sanitize_payload(block_input),
                },
            )
            return

        tool_use_id = getattr(block, "tool_use_id", None)
        if tool_use_id is not None:
            content = getattr(block, "content", None)
            self._emit_observability_event(
                observability,
                agent_name,
                task_id,
                "tool_result",
                {
                    "session_id": session_id,
                    "tool_use_id": tool_use_id,
                    "parent_tool_use_id": parent_tool_use_id,
                    "is_error": bool(getattr(block, "is_error", False)),
                    "content": self._sanitize_payload(content),
                },
            )
            return

        text = getattr(block, "text", None)
        if isinstance(text, str) and text:
            self._emit_observability_event(
                observability,
                agent_name,
                task_id,
                "text_block",
                {
                    "session_id": session_id,
                    "chars": len(text),
                    "preview": self._truncate_text(text),
                },
            )
            return

        thinking = getattr(block, "thinking", None)
        if isinstance(thinking, str) and thinking:
            self._emit_observability_event(
                observability,
                agent_name,
                task_id,
                "thinking_block",
                {
                    "session_id": session_id,
                    "chars": len(thinking),
                },
            )

    def _emit_observability_event(
        self,
        observability: Any | None,
        agent_name: str,
        task_id: str,
        event_type: str,
        payload: dict[str, Any] | None = None,
    ) -> None:
        if observability is None:
            return
        log_method = getattr(observability, "log_claude_event", None)
        if not callable(log_method):
            return
        try:
            log_method(
                agent_name=agent_name,
                task_id=task_id,
                event_type=event_type,
                payload=self._sanitize_payload(payload),
            )
        except Exception:
            # Observability should never break execution.
            logger.debug("Failed to emit observability event", exc_info=True)

    @classmethod
    def _is_sensitive_key(cls, key: Any) -> bool:
        if not isinstance(key, str):
            return False
        lowered = key.lower()
        return any(token in lowered for token in cls._SENSITIVE_KEY_TOKENS)

    @classmethod
    def _sanitize_payload(cls, value: Any, *, _depth: int = 0) -> Any:
        if _depth >= 4:
            return "[truncated]"

        if isinstance(value, dict):
            sanitized: dict[str, Any] = {}
            for idx, (k, v) in enumerate(value.items()):
                if idx >= 40:
                    sanitized["__truncated_items__"] = len(value) - 40
                    break
                key = str(k)
                if cls._is_sensitive_key(key):
                    sanitized[key] = "[REDACTED]"
                else:
                    sanitized[key] = cls._sanitize_payload(v, _depth=_depth + 1)
            return sanitized

        if isinstance(value, (list, tuple)):
            items = [cls._sanitize_payload(v, _depth=_depth + 1) for v in value[:40]]
            if len(value) > 40:
                items.append(f"...({len(value) - 40} more)")
            return items

        if isinstance(value, str):
            return cls._truncate_text(value)

        if isinstance(value, (int, float, bool)) or value is None:
            return value

        return cls._truncate_text(str(value))

    @staticmethod
    def _truncate_text(value: str, max_chars: int = 400) -> str:
        if not isinstance(value, str):
            return ""
        trimmed = value.strip()
        if len(trimmed) <= max_chars:
            return trimmed
        return f"{trimmed[:max_chars]}...({len(trimmed) - max_chars} more chars)"
