"""
Drop-in replacement for Claude Agent SDK with automatic code caching.

Simply swap your import:
    # Before
    from claude_agent_sdk import ClaudeSDKClient, ClaudeAgentOptions

    # After
    from raysurfer import ClaudeSDKClient, ClaudeAgentOptions

Everything else works exactly the same. Raysurfer automatically:
1. Retrieves relevant cached code before each query
2. Uploads new code after successful task completion
"""

import atexit
import logging
import os
import shutil
import tempfile
from typing import Any, AsyncIterator

# Isolate temp directory to avoid file watcher conflicts when running
# nested inside another Claude Code session.
_ISOLATED_TMPDIR: str | None = None


def _setup_isolated_env() -> str:
    """Set up isolated temp directory for nested Claude Code execution."""
    global _ISOLATED_TMPDIR

    if _ISOLATED_TMPDIR is None:
        _ISOLATED_TMPDIR = tempfile.mkdtemp(prefix="raysurfer_sdk_")
        os.environ["TMPDIR"] = _ISOLATED_TMPDIR
        os.environ["TEMP"] = _ISOLATED_TMPDIR
        os.environ["TMP"] = _ISOLATED_TMPDIR
        os.environ["CLAUDE_AGENT_SDK_SKIP_VERSION_CHECK"] = "1"
        atexit.register(_cleanup_isolated_env)

    return _ISOLATED_TMPDIR


def _cleanup_isolated_env() -> None:
    """Clean up isolated temp directory on exit."""
    global _ISOLATED_TMPDIR
    if _ISOLATED_TMPDIR and os.path.exists(_ISOLATED_TMPDIR):
        try:
            shutil.rmtree(_ISOLATED_TMPDIR)
        except Exception:
            pass
        _ISOLATED_TMPDIR = None


_setup_isolated_env()

# Import and re-export everything from Claude Agent SDK
from claude_agent_sdk import ClaudeAgentOptions as _BaseClaudeAgentOptions
from claude_agent_sdk import ClaudeSDKClient as _BaseClaudeSDKClient
from claude_agent_sdk import (
    AgentDefinition,
    AssistantMessage,
    HookMatcher,
    Message,
    ResultMessage,
    SystemMessage,
    TextBlock,
    ThinkingBlock,
    ToolResultBlock,
    ToolUseBlock,
    UserMessage,
)

from raysurfer.client import AsyncRaySurfer
from raysurfer.types import FileWritten

logger = logging.getLogger(__name__)

# Re-export all SDK types
__all__ = [
    "ClaudeSDKClient",
    "ClaudeAgentOptions",
    "AgentDefinition",
    "HookMatcher",
    "Message",
    "UserMessage",
    "AssistantMessage",
    "SystemMessage",
    "ResultMessage",
    "TextBlock",
    "ThinkingBlock",
    "ToolUseBlock",
    "ToolResultBlock",
]

# Default backend URL
DEFAULT_RAYSURFER_URL = "https://web-production-3d338.up.railway.app"


class ClaudeAgentOptions(_BaseClaudeAgentOptions):
    """
    Drop-in replacement for ClaudeAgentOptions with Raysurfer caching.

    All standard ClaudeAgentOptions fields work exactly the same.
    Raysurfer caching is enabled automatically when RAYSURFER_API_KEY is set.
    """
    pass


class ClaudeSDKClient:
    """
    Drop-in replacement for ClaudeSDKClient with automatic Raysurfer caching.

    Usage is identical to the original ClaudeSDKClient:

        from raysurfer import ClaudeSDKClient, ClaudeAgentOptions

        options = ClaudeAgentOptions(
            allowed_tools=["Read", "Write", "Bash"],
            system_prompt="You are a helpful assistant.",
        )

        async with ClaudeSDKClient(options=options) as client:
            await client.query("Fetch data from GitHub API")
            async for msg in client.receive_response():
                print(msg)

    Set RAYSURFER_API_KEY environment variable to enable caching.
    """

    def __init__(self, options: ClaudeAgentOptions | None = None):
        self._options = options or ClaudeAgentOptions()
        self._base_client: _BaseClaudeSDKClient | None = None
        self._raysurfer: AsyncRaySurfer | None = None
        self._current_query: str | None = None
        self._generated_files: list[FileWritten] = []
        self._task_succeeded: bool = False
        self._cache_enabled: bool = False
        self._cached_code_blocks: list[dict[str, Any]] = []  # Track retrieved code blocks

    async def __aenter__(self) -> "ClaudeSDKClient":
        """Initialize the client."""
        # Check if caching should be enabled
        api_key = os.environ.get("RAYSURFER_API_KEY")
        base_url = os.environ.get("RAYSURFER_BASE_URL", DEFAULT_RAYSURFER_URL)

        if api_key:
            self._cache_enabled = True
            self._raysurfer = AsyncRaySurfer(api_key=api_key, base_url=base_url)
            await self._raysurfer.__aenter__()

        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
        """Clean up resources."""
        if self._base_client:
            await self._base_client.__aexit__(exc_type, exc_val, exc_tb)
        if self._raysurfer:
            await self._raysurfer.__aexit__(exc_type, exc_val, exc_tb)

    async def query(self, prompt: str) -> None:
        """
        Send a query to Claude.

        If RAYSURFER_API_KEY is set, automatically retrieves relevant cached
        code and injects it into the system prompt.
        """
        self._current_query = prompt
        self._generated_files = []
        self._task_succeeded = False
        self._cached_code_blocks = []

        # Retrieve cached code if caching is enabled
        augmented_options = await self._augment_options_with_cache(prompt)

        # Initialize and query the base client
        self._base_client = _BaseClaudeSDKClient(options=augmented_options)
        await self._base_client.__aenter__()
        await self._base_client.query(prompt)

    async def receive_response(self) -> AsyncIterator[Message]:
        """
        Receive and yield response messages from Claude.

        After successful task completion, automatically uploads any
        generated code to the Raysurfer cache.
        """
        if not self._base_client:
            raise RuntimeError("Must call query() before receive_response()")

        async for message in self._base_client.receive_response():
            # Track Write tool calls
            if isinstance(message, AssistantMessage):
                for block in message.content:
                    if isinstance(block, ToolUseBlock) and block.name == "Write":
                        self._track_file(block.input)

            # Check for successful completion
            if isinstance(message, ResultMessage):
                if message.subtype == "success":
                    self._task_succeeded = True

            yield message

        # Upload generated code if task succeeded
        if self._cache_enabled and self._task_succeeded and self._generated_files:
            await self._upload_to_cache()

        # Submit votes for cached code blocks that were used
        if self._cache_enabled and self._task_succeeded and self._cached_code_blocks:
            await self._submit_votes()

    def _track_file(self, tool_input: dict[str, Any]) -> None:
        """Track a file written by the Write tool."""
        file_path = tool_input.get("file_path", "")
        content = tool_input.get("content", "")
        if file_path and content:
            self._generated_files.append(FileWritten(path=file_path, content=content))

    async def _augment_options_with_cache(self, task: str) -> _BaseClaudeAgentOptions:
        """Retrieve cached code and augment system prompt."""
        if not self._cache_enabled or not self._raysurfer:
            return self._options

        try:
            response = await self._raysurfer.get_code_files(
                task=task,
                top_k=5,
                min_verdict_score=0.3,
                prefer_complete=True,
            )

            if not response.files:
                return self._options

            # Track retrieved code blocks for voting
            self._cached_code_blocks = [
                {
                    "code_block_id": f.code_block_id,
                    "filename": f.filename,
                    "description": f.description,
                }
                for f in response.files
            ]

            # Build augmented system prompt
            snippets = self._format_code_snippets(response.files)
            base_prompt = self._options.system_prompt

            if isinstance(base_prompt, dict) and base_prompt.get("type") == "preset":
                augmented_prompt = {
                    **base_prompt,
                    "append": base_prompt.get("append", "") + snippets,
                }
            else:
                augmented_prompt = (base_prompt or "") + snippets

            # Create new options with augmented prompt
            return _BaseClaudeAgentOptions(
                allowed_tools=self._options.allowed_tools,
                disallowed_tools=self._options.disallowed_tools,
                permission_mode=self._options.permission_mode,
                system_prompt=augmented_prompt,
                cwd=self._options.cwd,
                add_dirs=self._options.add_dirs,
                max_turns=self._options.max_turns,
                model=self._options.model,
                env=self._options.env,
                mcp_servers=self._options.mcp_servers,
                hooks=self._options.hooks,
                can_use_tool=self._options.can_use_tool,
                setting_sources=self._options.setting_sources,
                include_partial_messages=self._options.include_partial_messages,
                fork_session=self._options.fork_session,
                continue_conversation=self._options.continue_conversation,
                resume=self._options.resume,
                agents=self._options.agents,
                plugins=self._options.plugins,
                enable_file_checkpointing=self._options.enable_file_checkpointing,
                output_format=self._options.output_format,
                sandbox=self._options.sandbox,
                extra_args=self._options.extra_args,
                max_buffer_size=self._options.max_buffer_size,
                stderr=self._options.stderr,
                user=self._options.user,
                settings=self._options.settings,
                permission_prompt_tool_name=self._options.permission_prompt_tool_name,
            )
        except Exception as e:
            logger.debug(f"Failed to retrieve cached code: {e}")
            return self._options

    def _format_code_snippets(self, files) -> str:
        """Format cached code files as markdown for system prompt."""
        snippets = "\n\n## Cached Code (from Raysurfer)\n\n"
        snippets += "The following pre-validated code is available for this task. "
        snippets += "Use it directly or adapt it as needed.\n\n"

        for f in files:
            snippets += f"### {f.filename}\n"
            snippets += f"**Description**: {f.description}\n"
            snippets += f"**Entrypoint**: `{f.entrypoint}`\n"
            snippets += f"**Confidence**: {f.verdict_score:.0%}\n\n"
            snippets += f"```{f.language}\n{f.source}\n```\n\n"

        return snippets

    async def _upload_to_cache(self) -> None:
        """Upload generated code to the Raysurfer cache."""
        if not self._raysurfer or not self._current_query:
            return

        try:
            result = await self._raysurfer.submit_execution_result(
                task=self._current_query,
                files_written=self._generated_files,
                succeeded=self._task_succeeded,
            )
            if result.code_blocks_stored > 0:
                logger.info(f"Cached {result.code_blocks_stored} code blocks")
        except Exception as e:
            logger.debug(f"Failed to upload to cache: {e}")

    async def _submit_votes(self) -> None:
        """Submit thumbs up votes for cached code blocks that helped complete the task."""
        if not self._raysurfer or not self._current_query:
            return

        for block in self._cached_code_blocks:
            try:
                await self._raysurfer.record_cache_usage(
                    task=self._current_query,
                    code_block_id=block["code_block_id"],
                    code_block_name=block["filename"],
                    code_block_description=block["description"],
                    succeeded=self._task_succeeded,
                )
                logger.debug(f"Submitted vote for {block['filename']}")
            except Exception as e:
                logger.debug(f"Failed to submit vote for {block['filename']}: {e}")
