"""
Python equivalent of ts-agent's AgentSdkRunner.

Calls the Claude Agent SDK (claude_code_sdk) with a pre-built prompt
and writes the generated tests into a pre-created scaffold file.
"""

import asyncio
import glob
import logging
import os
import platform
import shutil
import stat
from collections.abc import AsyncIterable, AsyncIterator
from typing import Any

from claude_code_sdk import AssistantMessage, ClaudeCodeOptions, HookMatcher, ResultMessage, query
from claude_code_sdk._internal.transport.subprocess_cli import SubprocessCLITransport

from py_agent.services.test_prune_hook import create_post_tool_hook

logger = logging.getLogger(__name__)

SDK_MODEL = "claude-sonnet-4-6"


def _find_claude_cli() -> str | None:
    """
    Resolve the 'claude' CLI binary path.

    Checks in order:
    0. CLAUDE_CLI_PATH env var (set by @earlyai/cli to the bundled cli.js)
    1. Standard PATH (covers npm global installs, nvm, homebrew, etc.)
    2. Platform-specific well-known locations
    3. Common local node_modules paths
    """
    # 0. Explicit path from early-cli (points to bundled cli.js)
    explicit = os.environ.get("CLAUDE_CLI_PATH")
    if explicit and os.path.isfile(explicit):
        return _make_executable(explicit)

    # 1. Whatever is on PATH right now
    found = shutil.which("claude")
    if found:
        return found

    is_windows = platform.system() == "Windows"

    if is_windows:
        # 2a. Windows: npm global installs live under %APPDATA%\npm
        appdata = os.environ.get("APPDATA", "")
        if appdata:
            for name in ("claude.cmd", "claude.exe", "claude"):
                candidate = os.path.join(appdata, "npm", name)
                if os.path.isfile(candidate):
                    return candidate
    else:
        # 2b. macOS: Claude desktop app bundle
        desktop_pattern = os.path.expanduser(
            "~/Library/Application Support/Claude/claude-code/*/claude"
        )
        matches = sorted(glob.glob(desktop_pattern), reverse=True)
        if matches:
            return matches[0]

    # 3. Local node_modules (project or user-level)
    bin_name = "claude.cmd" if is_windows else "claude"
    for candidate in [
        os.path.expanduser(f"~/node_modules/.bin/{bin_name}"),
        os.path.join(os.getcwd(), "node_modules", ".bin", bin_name),
    ]:
        if os.path.isfile(candidate):
            return candidate

    if not is_windows:
        # 4. nvm-managed Node installs (macOS/Linux)
        nvm_dir = os.environ.get("NVM_DIR", os.path.expanduser("~/.nvm"))
        nvm_versions = os.path.join(nvm_dir, "versions", "node")
        if os.path.isdir(nvm_versions):
            # Check all nvm node versions, newest first
            for candidate in sorted(
                glob.glob(os.path.join(nvm_versions, "*/bin/claude")),
                reverse=True,
            ):
                if os.path.isfile(candidate):
                    return candidate
            # Also check for cli.js bundled with @earlyai/cli
            for candidate in sorted(
                glob.glob(
                    os.path.join(
                        nvm_versions,
                        "*/lib/node_modules/@earlyai/cli/node_modules/@anthropic-ai/claude-agent-sdk/cli.js",
                    )
                ),
                reverse=True,
            ):
                if os.path.isfile(candidate):
                    return _make_executable(candidate)

    return None


def _make_executable(cli_js_path: str) -> str:
    """
    Ensure a cli.js file is directly executable.

    On Unix the #!/usr/bin/env node shebang works as-is.
    On Windows we return the path as-is; the transport layer
    handles running it via node (see _NodeCLITransport).
    """
    if platform.system() == "Windows":
        return cli_js_path

    # Ensure the execute bit is set (it should already be, but just in case)
    st = os.stat(cli_js_path)
    if not (st.st_mode & stat.S_IXUSR):
        os.chmod(cli_js_path, st.st_mode | stat.S_IXUSR)
    return cli_js_path


def _resolve_cli_js_for_windows(cli_path: str) -> str | None:
    """Resolve the .js entry point from a CLI path on Windows.

    npm .cmd wrappers delegate to a .js file via node.  Since
    SubprocessCLITransport uses shell=False (via anyio.open_process),
    .cmd files can't be executed directly.  This resolves to the
    underlying .js so we can run it as ``[node, cli.js, ...]``.
    """
    if cli_path.endswith(".js"):
        return cli_path

    # npm global/local: claude.cmd sits next to node_modules/
    parent = os.path.dirname(cli_path)
    candidate = os.path.join(parent, "node_modules", "@anthropic-ai", "claude-code", "cli.js")
    if os.path.isfile(candidate):
        return candidate

    return None


class _HookAwareTransport(SubprocessCLITransport):
    """Transport that keeps stdin open after prompt delivery.

    The base ``SubprocessCLITransport.end_input()`` closes stdin once the
    prompt stream is exhausted.  However, when hooks are registered the CLI
    sends ``hook_callback`` control requests and expects responses through
    stdin.  Closing stdin prematurely causes "Stream closed" errors.

    This subclass turns ``end_input()`` into a no-op so stdin stays open
    for the lifetime of the process.  The stream is still closed normally
    when ``close()`` is called at the end of the query.
    """

    async def end_input(self) -> None:  # noqa: D102
        # Intentionally keep stdin open for hook callback responses.
        pass


class _NodeCLITransport(_HookAwareTransport):
    """Transport that runs ``node <cli.js> ...`` instead of ``<cli> ...``.

    On Windows, .js and .cmd files cannot be spawned directly with
    shell=False (which anyio.open_process uses).  This subclass prepends
    the node executable so the command becomes ``[node, cli.js, --args...]``.
    """

    def __init__(
        self,
        prompt: str | AsyncIterable[dict[str, Any]],
        options: ClaudeCodeOptions,
        cli_js_path: str,
        node_path: str,
    ):
        self._node_path = node_path
        super().__init__(prompt=prompt, options=options, cli_path=cli_js_path)

    def _build_command(self) -> list[str]:
        cmd = super()._build_command()
        # cmd is [cli_js_path, ...args]; prepend node to execute it
        return [self._node_path, *cmd]


SDK_ALLOWED_TOOLS = ["Read", "Write", "Edit", "Bash", "Glob", "Grep"]
SDK_PERMISSION_MODE = "bypassPermissions"
SDK_MAX_TURNS = 50


async def _prompt_as_stream(prompt: str) -> AsyncIterator[dict[str, Any]]:
    """Wrap a string prompt as a streaming input message.

    The Claude Code SDK passes string prompts as CLI arguments which hits
    the ~32K command-line length limit on Windows.  Yielding the prompt as
    an async iterable of dicts makes the SDK send it via stdin instead.
    """
    yield {
        "type": "user",
        "message": {"role": "user", "content": prompt},
        "parent_tool_use_id": None,
        "session_id": None,
    }


async def run_agent(
    user_prompt: str,
    system_prompt: str,
    test_file_path: str,
    working_directory: str,
    jwt_token: str,
    api_base_url: str,
) -> dict[str, Any]:
    """
    Runs the Claude Agent SDK to generate tests into test_file_path.

    Returns a dict with keys: code (str), cost_usd (float | None).
    """
    anthropic_proxy_url = api_base_url.rstrip("/") + "/anthropic"

    env = {
        **os.environ,
        "ANTHROPIC_BASE_URL": anthropic_proxy_url,
        "ANTHROPIC_AUTH_TOKEN": jwt_token,
    }

    done_event = asyncio.Event()
    post_hook = create_post_tool_hook(done_event)

    options = ClaudeCodeOptions(
        model=SDK_MODEL,
        system_prompt=system_prompt,
        allowed_tools=SDK_ALLOWED_TOOLS,
        permission_mode="bypassPermissions",
        cwd=working_directory,
        max_turns=SDK_MAX_TURNS,
        env=env,
        hooks={
            "PostToolUse": [
                HookMatcher(matcher="Bash", hooks=[post_hook]),
            ],
        },
    )

    cli_path = _find_claude_cli()
    if cli_path:
        logger.debug(f"Using Claude CLI at: {cli_path}")
    else:
        raise RuntimeError(
            "Claude CLI ('claude') not found. "
            "Install it with: npm install -g @anthropic-ai/claude-code"
        )

    prompt_stream = _prompt_as_stream(user_prompt)

    # Use _HookAwareTransport to keep stdin open for hook callbacks.
    # On Windows, _NodeCLITransport (extends _HookAwareTransport) runs
    # [node, cli.js, ...] since .cmd/.js can't be launched with shell=False.
    transport: SubprocessCLITransport
    if platform.system() == "Windows":
        cli_js = _resolve_cli_js_for_windows(cli_path)
        node_path = shutil.which("node")
        if cli_js and node_path:
            logger.debug(f"Windows: using node ({node_path}) to run {cli_js}")
            transport = _NodeCLITransport(
                prompt=prompt_stream,
                options=options,
                cli_js_path=cli_js,
                node_path=node_path,
            )
        else:
            transport = _HookAwareTransport(
                prompt=prompt_stream,
                options=options,
                cli_path=cli_path,
            )
    else:
        transport = _HookAwareTransport(
            prompt=prompt_stream,
            options=options,
            cli_path=cli_path,
        )

    logger.info(f"Running Claude Agent SDK — test file: {test_file_path}")

    cost_usd: float | None = None
    usage: dict[str, Any] | None = None
    terminated = False
    try:
        async for message in query(prompt=prompt_stream, options=options, transport=transport):
            if isinstance(message, ResultMessage):
                cost_usd = message.total_cost_usd
                usage = message.usage
                if cost_usd is not None:
                    logger.info(f"Agent finished. Cost: ${cost_usd:.4f}")
                else:
                    logger.info("Agent finished.")
            elif isinstance(message, AssistantMessage):
                logger.debug("Agent message received")

            # PostToolUse hook detected all tests passed — kill the subprocess.
            # Don't break (causes anyio cancel-scope errors); instead terminate
            # the process so the stream EOFs and the loop ends naturally.
            if done_event.is_set() and not terminated:
                terminated = True
                logger.info("All tests passed — terminating agent session")
                # Suppress the SDK's "Fatal error in message reader" log that
                # fires when the process exits with SIGTERM (143).
                logging.getLogger("claude_code_sdk._internal.query").setLevel(logging.CRITICAL)
                if transport._process and transport._process.returncode is None:
                    transport._process.terminate()
    except Exception:
        if not terminated:
            raise
        # Exit code 143 (SIGTERM) is expected when we kill the process after
        # all tests pass.  The test file is already on disk.
        logger.info("Agent process terminated (expected after all tests passed)")

    code = ""
    if os.path.exists(test_file_path):
        with open(test_file_path, encoding="utf-8") as f:
            code = f.read()

    return {"code": code, "cost_usd": cost_usd, "usage": usage}
