"""Agent harness abstractions for different agent backends."""

import asyncio
import json
import os
import shutil
from dataclasses import dataclass, field
from pathlib import Path
from typing import TYPE_CHECKING, Any, Protocol, TextIO, runtime_checkable

from rich.console import Console

from .docker_env import TaskEnvironment
from .log_formatter import InstanceLogWriter, StreamEventFormatter

if TYPE_CHECKING:
    from .config import MCPServerConfig


@dataclass
class AgentResult:
    """Result from an agent run."""

    patch: str
    success: bool
    error: str | None = None
    tokens_input: int = 0
    tokens_output: int = 0
    iterations: int = 0
    tool_calls: int = 0
    tool_usage: dict[str, int] = field(default_factory=dict)
    messages: list[dict[str, Any]] = field(default_factory=list)
    stdout: str = ""
    stderr: str = ""


@runtime_checkable
class AgentHarness(Protocol):
    """Protocol for agent harnesses that solve SWE-bench tasks.

    To add a new harness:
    1. Create a class implementing this protocol
    2. Add it to HARNESS_REGISTRY
    3. Add the harness name to VALID_HARNESSES in config.py
    """

    async def solve(
        self,
        task: dict[str, Any],
        workdir: str,
        timeout: int = 300,
        verbose: bool = False,
        task_id: str | None = None,
        env: TaskEnvironment | None = None,
    ) -> AgentResult:
        """Solve a SWE-bench task and return the patch.

        Args:
            task: SWE-bench task dictionary with problem_statement, etc.
            workdir: Path to the repository working directory.
            timeout: Timeout in seconds.
            verbose: If True, stream output to console.
            task_id: Task identifier for prefixing output.
            env: Optional Docker environment to run inside.

        Returns:
            AgentResult with the generated patch.
        """
        ...


async def _run_cli_command(
    command: list[str],
    workdir: str,
    timeout: int,
    env: dict[str, str] | None = None,
    input_text: str | None = None,
) -> tuple[int, str, str]:
    """Run a CLI command asynchronously.

    Args:
        command: Command and arguments to run.
        workdir: Working directory.
        timeout: Timeout in seconds.
        env: Optional environment variables.
        input_text: Optional text to pass to stdin.

    Returns:
        Tuple of (exit_code, stdout, stderr).
    """
    full_env = {**os.environ}
    if env:
        full_env.update(env)

    process = await asyncio.create_subprocess_exec(
        *command,
        cwd=workdir,
        env=full_env,
        stdin=asyncio.subprocess.PIPE if input_text else asyncio.subprocess.DEVNULL,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE,
    )

    try:
        input_bytes = input_text.encode("utf-8") if input_text else None
        stdout, stderr = await asyncio.wait_for(
            process.communicate(input=input_bytes),
            timeout=timeout,
        )
        return (
            process.returncode or 0,
            stdout.decode("utf-8", errors="replace"),
            stderr.decode("utf-8", errors="replace"),
        )
    except asyncio.TimeoutError:
        process.kill()
        return -1, "", "Command timed out"


async def _run_cli_streaming(
    command: list[str],
    workdir: str,
    timeout: int,
    env: dict[str, str] | None = None,
    prefix: str | None = None,
    console: Console | None = None,
    formatter: StreamEventFormatter | None = None,
) -> tuple[int, str, str]:
    """Run a CLI command with real-time streaming output.

    Args:
        command: Command and arguments to run.
        workdir: Working directory.
        timeout: Timeout in seconds.
        env: Optional environment variables.
        prefix: Optional prefix for each output line.
        console: Rich console for output (if None and no formatter, uses print).
        formatter: Optional StreamEventFormatter for structured JSON output.

    Returns:
        Tuple of (exit_code, stdout, stderr).
    """
    full_env = {**os.environ}
    if env:
        full_env.update(env)

    process = await asyncio.create_subprocess_exec(
        *command,
        cwd=workdir,
        env=full_env,
        stdin=asyncio.subprocess.DEVNULL,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE,
    )

    stdout_lines: list[str] = []
    stderr_lines: list[str] = []

    async def read_stream(
        stream: asyncio.StreamReader | None,
        lines: list[str],
        is_stderr: bool = False,
    ) -> None:
        if stream is None:
            return
        while True:
            line = await stream.readline()
            if not line:
                break
            decoded = line.decode("utf-8", errors="replace").rstrip()
            lines.append(decoded)

            if is_stderr:
                if console:
                    console.print(f"[dim red]{decoded}[/dim red]")
            elif formatter:
                formatter.format_line(decoded, prefix)
            elif console:
                prefix_str = f"[{prefix}] " if prefix else ""
                console.print(f"[dim]{prefix_str}{decoded}[/dim]")

    try:
        await asyncio.wait_for(
            asyncio.gather(
                read_stream(process.stdout, stdout_lines),
                read_stream(process.stderr, stderr_lines, is_stderr=True),
            ),
            timeout=timeout,
        )
        await process.wait()
        return (
            process.returncode or 0,
            "\n".join(stdout_lines),
            "\n".join(stderr_lines),
        )
    except asyncio.TimeoutError:
        process.kill()
        return -1, "\n".join(stdout_lines), "Command timed out"


async def _get_git_diff(workdir: str) -> str:
    """Get the git diff of changes, including only modifications to existing files.

    Stages all changes but only includes modifications (not new files) in the diff.
    This excludes agent-created debug scripts, test files, and other artifacts.

    Args:
        workdir: Path to the git repository.

    Returns:
        Git diff string with only modified file changes.
    """
    await _run_cli_command(["git", "add", "-A"], workdir, timeout=30)

    exit_code, stdout, stderr = await _run_cli_command(
        [
            "git",
            "diff",
            "--cached",
            "HEAD",
            "--diff-filter=M",
        ],
        workdir,
        timeout=30,
    )
    if exit_code == 0:
        return stdout
    return ""


async def _get_git_diff_in_docker(env: TaskEnvironment) -> str:
    """Get the git diff of changes from inside Docker container.

    Args:
        env: Docker task environment.

    Returns:
        Git diff string with only modified file changes.
    """
    workdir = env.workdir

    await env.exec_command(
        f"git config --global --add safe.directory {workdir}",
        timeout=10,
        workdir=workdir,
    )

    _, status_out, _ = await env.exec_command(
        "git status --short",
        timeout=30,
        workdir=workdir,
    )

    await env.exec_command("git add -A", timeout=30, workdir=workdir)

    exit_code, stdout, stderr = await env.exec_command(
        "git diff --cached HEAD --diff-filter=M",
        timeout=30,
        workdir=workdir,
    )
    if exit_code == 0 and stdout.strip():
        return stdout

    exit_code, stdout, stderr = await env.exec_command(
        "git diff --cached HEAD",
        timeout=30,
        workdir=workdir,
    )
    if exit_code == 0 and stdout.strip():
        return stdout

    _, unstaged, _ = await env.exec_command(
        "git diff",
        timeout=30,
        workdir=workdir,
    )
    if unstaged.strip():
        await env.exec_command("git add -A", timeout=30, workdir=workdir)
        _, stdout, _ = await env.exec_command(
            "git diff --cached HEAD",
            timeout=30,
            workdir=workdir,
        )
        if stdout.strip():
            return stdout

    return ""


async def _write_prompt_file(workdir: str, problem_statement: str) -> str:
    """Write the problem statement to a temporary file.

    Args:
        workdir: Working directory.
        problem_statement: The problem statement text.

    Returns:
        Path to the prompt file.
    """
    prompt_path = os.path.join(workdir, ".swebench_prompt.md")
    Path(prompt_path).write_text(f"# Problem Statement\n\n{problem_statement}")
    return prompt_path


def _parse_tool_usage_from_stream(
    stdout: str,
) -> tuple[int, dict[str, int], int, int, int, str | None]:
    """Parse tool usage and metadata from stream-json output.

    Args:
        stdout: Raw stdout from Claude Code CLI in stream-json format.

    Returns:
        Tuple of (total_tool_calls, tool_usage_dict, num_turns, tokens_in, tokens_out, result_subtype).
    """
    tool_usage: dict[str, int] = {}
    total_tool_calls = 0
    num_turns = 0
    tokens_in = 0
    tokens_out = 0
    result_tokens_in = 0
    result_tokens_out = 0
    has_result = False
    result_subtype: str | None = None

    for line in stdout.split("\n"):
        if not line.strip():
            continue
        try:
            event = json.loads(line)

            if event.get("type") == "assistant":
                message = event.get("message", {})
                content = message.get("content", [])
                if isinstance(content, list):
                    for block in content:
                        if isinstance(block, dict) and block.get("type") == "tool_use":
                            tool_name = block.get("name", "unknown")
                            tool_usage[tool_name] = tool_usage.get(tool_name, 0) + 1
                            total_tool_calls += 1

                usage = message.get("usage", {})
                tokens_in += usage.get("input_tokens", 0)
                tokens_out += usage.get("output_tokens", 0)

            if event.get("type") == "result":
                has_result = True
                num_turns = event.get("num_turns", 1)
                result_subtype = event.get("subtype")
                usage = event.get("usage", {})
                result_tokens_in = usage.get("input_tokens", 0)
                result_tokens_out = usage.get("output_tokens", 0)

        except json.JSONDecodeError:
            continue

    if has_result and (result_tokens_in > 0 or result_tokens_out > 0):
        tokens_in = result_tokens_in
        tokens_out = result_tokens_out

    return total_tool_calls, tool_usage, num_turns, tokens_in, tokens_out, result_subtype


DEFAULT_PROMPT = (
    "Fix the following bug in this repository:\n\n"
    "{problem_statement}\n\n"
    "IMPORTANT CONSTRAINTS:\n"
    "- Only modify the minimum files necessary to fix the bug\n"
    "- Do NOT create new test files\n"
    "- Do NOT create documentation files\n"
    "- Do NOT create reproduction scripts\n"
    "- Focus solely on the fix in existing source files"
)

MCP_PROMPT_SUFFIX = (
    "\n\nYou have access to an MCP server with additional tools. "
    "Consider using the MCP tools (prefixed with mcp__) when they would "
    "help you understand or navigate the codebase more effectively."
)


class ClaudeCodeHarness:
    """Harness that uses Claude Code CLI (claude) for solving tasks."""

    def __init__(
        self,
        model: str | None = None,
        mcp_server: "MCPServerConfig | None" = None,
        prompt: str | None = None,
        max_iterations: int = 10,
        verbosity: int = 1,
        log_file: TextIO | InstanceLogWriter | None = None,
    ) -> None:
        """Initialize Claude Code harness.

        Args:
            model: Optional model override.
            mcp_server: MCP server configuration to use.
            prompt: Custom prompt template. Use {problem_statement} placeholder.
            max_iterations: Maximum number of agentic turns.
            verbosity: Verbosity level (0=silent, 1=summary, 2=detailed).
            log_file: Optional file handle for writing raw JSON logs.
        """
        self.model = model
        self.mcp_server = mcp_server
        self.prompt_template = prompt or DEFAULT_PROMPT
        if mcp_server:
            self.prompt_template += MCP_PROMPT_SUFFIX
        self.max_iterations = max_iterations
        self.verbosity = verbosity
        self.log_file = log_file
        self._console = Console()

    async def solve(
        self,
        task: dict[str, Any],
        workdir: str,
        timeout: int = 300,
        verbose: bool = False,
        task_id: str | None = None,
        env: TaskEnvironment | None = None,
    ) -> AgentResult:
        """Solve task using Claude Code CLI.

        If env is provided and has claude_cli_installed=True, runs inside Docker.
        Otherwise runs locally on the host.
        """
        if env and env.claude_cli_installed:
            return await self._solve_in_docker(task, env, timeout, verbose, task_id)

        return await self._solve_locally(task, workdir, timeout, verbose, task_id)

    async def _solve_locally(
        self,
        task: dict[str, Any],
        workdir: str,
        timeout: int,
        verbose: bool,
        task_id: str | None,
    ) -> AgentResult:
        """Solve task using local Claude Code CLI."""
        if not shutil.which("claude"):
            return AgentResult(
                patch="",
                success=False,
                error="Claude Code CLI (claude) not found in PATH",
            )

        problem_statement = task.get("problem_statement", "")
        prompt = self.prompt_template.format(problem_statement=problem_statement)
        instance_id = task_id or task.get("instance_id", "unknown")

        mcp_server_name = None
        if self.mcp_server:
            mcp_server_name = self.mcp_server.name
            args = self.mcp_server.get_args_for_workdir(workdir)
            mcp_env = self.mcp_server.get_expanded_env()
            add_cmd = [
                "claude",
                "mcp",
                "add",
                mcp_server_name,
                "--",
                self.mcp_server.command,
            ] + args
            exit_code, stdout, stderr = await _run_cli_command(
                add_cmd, workdir, timeout=30, env=mcp_env
            )
            if exit_code != 0:
                self._console.print(
                    f"[yellow]Warning: MCP server add failed (exit {exit_code}): {stderr or stdout}[/yellow]"
                )

        try:
            command = [
                "claude",
                "--print",
                "--verbose",
                "--dangerously-skip-permissions",
                "--output-format",
                "stream-json",
                "--max-turns",
                str(self.max_iterations),
            ]

            if self.model:
                command.extend(["--model", self.model])

            command.append(prompt)

            if verbose:
                from .log_formatter import FormatterConfig

                config = FormatterConfig(
                    verbosity=self.verbosity,
                    workdir=workdir,
                    show_timestamps=True,
                )
                formatter = StreamEventFormatter(self._console, config, self.log_file)

                run_type = "mcp" if self.mcp_server else "baseline"
                formatter.print_run_start(instance_id, run_type)

                exit_code, stdout, stderr = await _run_cli_streaming(
                    command,
                    workdir,
                    timeout,
                    prefix=instance_id,
                    console=self._console,
                    formatter=formatter,
                )
            else:
                exit_code, stdout, stderr = await _run_cli_command(
                    command,
                    workdir,
                    timeout,
                )

            total_tool_calls, tool_usage, num_turns, tokens_in, tokens_out, result_subtype = (
                _parse_tool_usage_from_stream(stdout)
            )

            if result_subtype == "error_max_turns" and num_turns > self.max_iterations:
                num_turns = self.max_iterations

            if exit_code != 0:
                error_msg = stderr or "Unknown error"
                if mcp_server_name:
                    await _run_cli_command(
                        ["claude", "mcp", "remove", mcp_server_name],
                        workdir,
                        timeout=10,
                    )
                return AgentResult(
                    patch="",
                    success=False,
                    error=f"Claude Code failed (exit {exit_code}): {error_msg}",
                    stdout=stdout,
                    stderr=stderr,
                    tokens_input=tokens_in,
                    tokens_output=tokens_out,
                    iterations=num_turns,
                    tool_calls=total_tool_calls,
                    tool_usage=tool_usage,
                )

            if mcp_server_name:
                await _run_cli_command(
                    ["claude", "mcp", "remove", mcp_server_name],
                    workdir,
                    timeout=10,
                )

            patch = await _get_git_diff(workdir)

            return AgentResult(
                patch=patch,
                success=bool(patch),
                error=None if patch else "No changes made by Claude Code",
                iterations=num_turns or 1,
                stdout=stdout,
                stderr=stderr,
                tokens_input=tokens_in,
                tokens_output=tokens_out,
                tool_calls=total_tool_calls,
                tool_usage=tool_usage,
            )
        except Exception:
            if mcp_server_name:
                await _run_cli_command(
                    ["claude", "mcp", "remove", mcp_server_name],
                    workdir,
                    timeout=10,
                )
            raise

    async def _solve_in_docker(
        self,
        task: dict[str, Any],
        env: TaskEnvironment,
        timeout: int,
        verbose: bool,
        task_id: str | None,
    ) -> AgentResult:
        """Solve task using Claude Code CLI inside Docker container."""
        problem_statement = task.get("problem_statement", "")
        prompt = self.prompt_template.format(problem_statement=problem_statement)
        instance_id = task_id or task.get("instance_id", "unknown")

        api_key = os.environ.get("ANTHROPIC_API_KEY", "")
        if not api_key:
            return AgentResult(
                patch="",
                success=False,
                error="ANTHROPIC_API_KEY environment variable not set",
            )

        docker_env = {
            "ANTHROPIC_API_KEY": api_key,
            "HOME": "/home/mcpbr",
            "USER": "mcpbr",
        }

        prompt_file = "/tmp/.mcpbr_prompt.txt"
        await env.exec_command(
            f"cat > {prompt_file} << 'MCPBR_PROMPT_EOF'\n{prompt}\nMCPBR_PROMPT_EOF",
            timeout=10,
        )
        await env.exec_command(f"chown mcpbr:mcpbr {prompt_file}", timeout=5)

        env_file = "/tmp/.mcpbr_env.sh"
        env_exports = f"export ANTHROPIC_API_KEY='{api_key}'\nexport HOME='/home/mcpbr'\n"

        # Add MCP server env vars
        if self.mcp_server:
            for key, value in self.mcp_server.get_expanded_env().items():
                env_exports += f"export {key}='{value}'\n"

        await env.exec_command(
            f"cat > {env_file} << 'MCPBR_ENV_EOF'\n{env_exports}MCPBR_ENV_EOF",
            timeout=10,
        )
        await env.exec_command(f"chown mcpbr:mcpbr {env_file}", timeout=5)

        # Build MCP server add command prefix (runs in same shell session as claude)
        mcp_server_name = None
        mcp_add_prefix = ""
        if self.mcp_server:
            mcp_server_name = self.mcp_server.name
            args = self.mcp_server.get_args_for_workdir(env.workdir)
            args_str = " ".join(args)
            mcp_add_prefix = (
                f"claude mcp add {mcp_server_name} -- {self.mcp_server.command} {args_str} && "
            )

        try:
            claude_args = [
                "--print",
                "--verbose",
                "--dangerously-skip-permissions",
                "--output-format",
                "stream-json",
                "--max-turns",
                str(self.max_iterations),
            ]

            if self.model:
                claude_args.extend(["--model", self.model])

            claude_args_str = " ".join(claude_args)
            claude_cmd = f'claude {claude_args_str} "$(cat {prompt_file})"'

            command = [
                "/bin/bash",
                "-c",
                f"cd {env.workdir} && su - mcpbr -c 'source {env_file} && cd {env.workdir} && {mcp_add_prefix}{claude_cmd}'",
            ]

            if verbose:
                from .log_formatter import FormatterConfig

                config = FormatterConfig(
                    verbosity=self.verbosity,
                    workdir=env.workdir,
                    show_timestamps=True,
                )
                formatter = StreamEventFormatter(self._console, config, self.log_file)

                run_type = "mcp" if self.mcp_server else "baseline"
                formatter.print_run_start(instance_id, run_type)

                def on_stdout(line: str) -> None:
                    formatter.format_line(line, instance_id)

                def on_stderr(line: str) -> None:
                    self._console.print(f"[dim red]{line}[/dim red]")

                exit_code, stdout, stderr = await env.exec_command_streaming(
                    command,
                    workdir=env.workdir,
                    environment=docker_env,
                    timeout=timeout,
                    on_stdout=on_stdout,
                    on_stderr=on_stderr,
                )
            else:
                exit_code, stdout, stderr = await env.exec_command(
                    command,
                    timeout=timeout,
                    environment=docker_env,
                )

            total_tool_calls, tool_usage, num_turns, tokens_in, tokens_out, result_subtype = (
                _parse_tool_usage_from_stream(stdout)
            )

            if result_subtype == "error_max_turns" and num_turns > self.max_iterations:
                num_turns = self.max_iterations

            if exit_code != 0:
                error_msg = stderr or "Unknown error"
                if mcp_server_name:
                    await env.exec_command(
                        f"su - mcpbr -c 'source {env_file} && claude mcp remove {mcp_server_name}'",
                        timeout=10,
                        environment=docker_env,
                    )
                return AgentResult(
                    patch="",
                    success=False,
                    error=f"Claude Code failed (exit {exit_code}): {error_msg}",
                    stdout=stdout,
                    stderr=stderr,
                    tokens_input=tokens_in,
                    tokens_output=tokens_out,
                    iterations=num_turns,
                    tool_calls=total_tool_calls,
                    tool_usage=tool_usage,
                )

            if mcp_server_name:
                await env.exec_command(
                    f"su - mcpbr -c 'source {env_file} && claude mcp remove {mcp_server_name}'",
                    timeout=10,
                    environment=docker_env,
                )

            _, git_status, _ = await env.exec_command(
                "git status --short",
                timeout=30,
            )

            # Debug: check if the buggy line still exists (should show nothing if fixed)
            _, sep_check, _ = await env.exec_command(
                "grep -n '= 1$' /workspace/astropy/modeling/separable.py | grep 'cright'",
                timeout=10,
            )

            # Also check file modification time
            _, file_info, _ = await env.exec_command(
                "stat -c '%Y %n' /workspace/astropy/modeling/separable.py",
                timeout=10,
            )

            patch = await _get_git_diff_in_docker(env)

            error_msg = None
            if not patch:
                buggy_line = sep_check.strip()
                if git_status.strip():
                    error_msg = f"Files changed but no valid patch: {git_status.strip()}"
                else:
                    # If buggy_line is empty, the fix was applied but git didn't detect it
                    if not buggy_line:
                        error_msg = "Edit applied but git shows no changes (file unchanged?)"
                    else:
                        error_msg = f"Buggy line still present: {buggy_line}"

            return AgentResult(
                patch=patch,
                success=bool(patch),
                error=error_msg,
                iterations=num_turns or 1,
                stdout=stdout,
                stderr=stderr,
                tokens_input=tokens_in,
                tokens_output=tokens_out,
                tool_calls=total_tool_calls,
                tool_usage=tool_usage,
            )
        except Exception:
            if mcp_server_name:
                try:
                    await env.exec_command(
                        f"su - mcpbr -c 'source {env_file} && claude mcp remove {mcp_server_name}'",
                        timeout=10,
                        environment=docker_env,
                    )
                except Exception:
                    pass
            raise
        finally:
            await env.exec_command(f"rm -f {prompt_file} {env_file}", timeout=5)


HARNESS_REGISTRY: dict[str, type] = {
    "claude-code": ClaudeCodeHarness,
}


def create_harness(
    harness_name: str,
    model: str | None = None,
    mcp_server: "MCPServerConfig | None" = None,
    prompt: str | None = None,
    max_iterations: int = 10,
    verbosity: int = 1,
    log_file: TextIO | InstanceLogWriter | None = None,
) -> AgentHarness:
    """Factory function to create an agent harness.

    Args:
        harness_name: Name of the harness (currently only 'claude-code').
        model: Optional model override.
        mcp_server: MCP server configuration (used by claude-code harness).
        prompt: Custom prompt template. Use {problem_statement} placeholder.
        max_iterations: Maximum agent iterations (used by claude-code harness).
        verbosity: Verbosity level for logging (0=silent, 1=summary, 2=detailed).
        log_file: Optional file handle for writing raw JSON logs.

    Returns:
        AgentHarness instance.

    Raises:
        ValueError: If harness_name is not recognized.
    """
    if harness_name not in HARNESS_REGISTRY:
        raise ValueError(
            f"Unknown harness: {harness_name}. Available harnesses: {list(HARNESS_REGISTRY.keys())}"
        )

    harness_class = HARNESS_REGISTRY[harness_name]

    return harness_class(
        model=model,
        mcp_server=mcp_server,
        prompt=prompt,
        max_iterations=max_iterations,
        verbosity=verbosity,
        log_file=log_file,
    )


def list_available_harnesses() -> list[str]:
    """List available external harness names.

    Returns:
        List of harness names.
    """
    return list(HARNESS_REGISTRY.keys())
