"""
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
import sys
import time
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 claude_code_sdk.types import TextBlock, ToolUseBlock

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,
    request_id: str,
    timeout: int = 240,
    testable_name: 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,
        "ANTHROPIC_CUSTOM_HEADERS": f"x-request-id: {request_id}",
    }

    done_event = asyncio.Event()
    test_counts: dict[str, int] = {"passed": 0, "failed": 0}
    post_hook = create_post_tool_hook(done_event, test_counts)

    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,
        )

    tag = f"[agent] [{testable_name}]" if testable_name else "[agent]"
    logger.info(f"Running Claude Agent SDK — test file: {test_file_path}")
    print(f"  {tag} Prompt size: {len(user_prompt)} chars", file=sys.stderr)

    cost_usd: float | None = None
    usage: dict[str, Any] | None = None
    terminated = False
    grace_deadline: float | None = None
    GRACE_SECONDS = 10
    turn = 0
    try:
        async with asyncio.timeout(timeout):
            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 after {turn} turns. Cost: ${cost_usd:.4f}")
                    else:
                        logger.info(f"Agent finished after {turn} turns.")
                    # Got cost data during grace period — terminate now
                    if grace_deadline is not None and not terminated:
                        terminated = True
                        logging.getLogger("claude_code_sdk._internal.query").setLevel(
                            logging.CRITICAL
                        )
                        if transport._process and transport._process.returncode is None:
                            transport._process.terminate()
                elif isinstance(message, AssistantMessage):
                    turn += 1
                    for block in message.content:
                        if isinstance(block, TextBlock):
                            preview = block.text[:200].replace("\n", " ")
                            print(f"  {tag} Turn {turn} text: {preview}", file=sys.stderr)
                            # Fail fast on known fatal errors (no tools used yet = first response)
                            if turn == 1 and not any(
                                isinstance(b, ToolUseBlock) for b in message.content
                            ):
                                lower = block.text.lower()
                                if "credit" in lower or "balance" in lower or "billing" in lower:
                                    logger.warning(f"Agent error: {block.text.strip()}")
                                    if transport._process and transport._process.returncode is None:
                                        transport._process.terminate()
                                    return {"code": "", "cost_usd": None, "usage": None}
                        elif isinstance(block, ToolUseBlock):
                            input_preview = str(block.input)[:150]
                            print(
                                f"  {tag} Turn {turn} tool: {block.name}({input_preview})",
                                file=sys.stderr,
                            )

                # PostToolUse hook detected all tests passed — start grace period
                # to receive the ResultMessage with cost/usage before terminating.
                if done_event.is_set() and grace_deadline is None:
                    grace_deadline = time.monotonic() + GRACE_SECONDS
                    logger.info("All tests passed — waiting for cost data before terminating")

                # Grace period expired without ResultMessage — force terminate
                if (
                    grace_deadline is not None
                    and time.monotonic() > grace_deadline
                    and not terminated
                ):
                    terminated = True
                    logger.info("Grace period expired — terminating agent session")
                    logging.getLogger("claude_code_sdk._internal.query").setLevel(logging.CRITICAL)
                    if transport._process and transport._process.returncode is None:
                        transport._process.terminate()
    except TimeoutError:
        logger.warning(f"Agent timed out after {timeout}s — returning partial result")
        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,
        "passed": test_counts["passed"],
        "failed": test_counts["failed"],
    }
