"""
Shell Tools

Run commands via subprocess and manage persistent shell sessions.
Commands execute within the workspace directory.
"""

import asyncio
import logging
import os
import shlex
import signal
import subprocess
import sys
import threading
import time
from typing import Dict, Any, Optional

from .base import BaseTool, ToolMetadata, ToolParameter, ToolPermission, ToolExecutionContext, ToolResult

logger = logging.getLogger(__name__)


class RunCommandTool(BaseTool):
    """Run a shell command and return the output."""

    def get_metadata(self) -> ToolMetadata:
        return ToolMetadata(
            name="RUN_COMMAND",
            description="Execute a shell command in the workspace directory. Returns stdout, stderr, and exit code.",
            category="shell",
            parameters=[
                ToolParameter("command", "string", "The command to execute", required=True),
                ToolParameter("timeout", "integer", "Timeout in seconds", default=300),
                ToolParameter("cwd", "string", "Working directory (defaults to workspace root)"),
                ToolParameter("env", "object", "Additional environment variables"),
            ],
            permissions=[ToolPermission.EXECUTE],
            timeout=300,
        )

    def execute(self, params: Dict[str, Any], context: ToolExecutionContext) -> ToolResult:
        command = params.get("command", "")
        timeout = params.get("timeout", 300)
        cwd = params.get("cwd", context.get_cwd())
        extra_env = params.get("env", {})

        if not command:
            return ToolResult.fail("command is required")

        # Build environment
        env = os.environ.copy()
        env.update(extra_env)

        try:
            is_windows = sys.platform == "win32"
            if is_windows:
                process = subprocess.Popen(
                    command,
                    shell=True,
                    stdout=subprocess.PIPE,
                    stderr=subprocess.PIPE,
                    cwd=cwd,
                    env=env,
                    text=True,
                    errors="replace",
                )
            else:
                process = subprocess.Popen(
                    command,
                    shell=True,
                    stdout=subprocess.PIPE,
                    stderr=subprocess.PIPE,
                    cwd=cwd,
                    env=env,
                    text=True,
                    errors="replace",
                    preexec_fn=os.setsid,
                )

            try:
                stdout, stderr = process.communicate(timeout=timeout)
            except subprocess.TimeoutExpired:
                # Kill process group
                if not is_windows:
                    os.killpg(os.getpgid(process.pid), signal.SIGKILL)
                else:
                    process.kill()
                stdout, stderr = process.communicate()
                return ToolResult.fail(
                    f"Command timed out after {timeout}s",
                    stdout=stdout[:5000] if stdout else "",
                    stderr=stderr[:5000] if stderr else "",
                    exit_code=-1,
                    timed_out=True,
                )

            exit_code = process.returncode

            # Build output
            output_parts = []
            if stdout:
                output_parts.append(stdout.strip())
            if stderr:
                output_parts.append(f"[stderr]\n{stderr.strip()}")

            output = "\n".join(output_parts) or "(no output)"

            # Truncate if too large
            if len(output) > 50_000:
                output = output[:50_000] + f"\n... (truncated, total {len(output)} chars)"

            if exit_code == 0:
                return ToolResult.ok(output, exit_code=0)
            else:
                return ToolResult(
                    success=False,
                    output=output,
                    error=f"Command exited with code {exit_code}",
                    metadata={"exit_code": exit_code},
                )

        except FileNotFoundError:
            return ToolResult.fail(f"Command not found or working directory does not exist: {cwd}")
        except Exception as e:
            return ToolResult.fail(f"Error executing command: {e}")


class PersistentShellManager:
    """Manages persistent shell sessions (pty-based)."""

    _sessions: Dict[str, Any] = {}  # session_id → subprocess
    _lock = threading.Lock()

    @classmethod
    def create_session(cls, session_id: str, cwd: str) -> bool:
        with cls._lock:
            if session_id in cls._sessions:
                return True  # Already exists

            try:
                is_windows = sys.platform == "win32"
                shell_cmd = "cmd.exe" if is_windows else "/bin/bash"

                proc = subprocess.Popen(
                    [shell_cmd],
                    stdin=subprocess.PIPE,
                    stdout=subprocess.PIPE,
                    stderr=subprocess.STDOUT,
                    cwd=cwd,
                    text=True,
                    errors="replace",
                    bufsize=0,
                )
                cls._sessions[session_id] = proc
                logger.info(f"Shell session created: {session_id}")
                return True
            except Exception as e:
                logger.error(f"Failed to create shell session {session_id}: {e}")
                return False

    @classmethod
    def send_command(cls, session_id: str, command: str, timeout: int = 60) -> Optional[str]:
        with cls._lock:
            proc = cls._sessions.get(session_id)
            if not proc or proc.poll() is not None:
                return None

        try:
            proc.stdin.write(command + "\n")
            proc.stdin.flush()

            # Read output with timeout
            output_lines = []
            end_time = time.time() + timeout

            while time.time() < end_time:
                # Non-blocking read
                try:
                    line = proc.stdout.readline()
                    if line:
                        output_lines.append(line)
                    else:
                        break
                except Exception:
                    break

            return "".join(output_lines)
        except Exception as e:
            return f"Error: {e}"

    @classmethod
    def close_session(cls, session_id: str):
        with cls._lock:
            proc = cls._sessions.pop(session_id, None)
            if proc:
                proc.terminate()
                try:
                    proc.wait(timeout=5)
                except subprocess.TimeoutExpired:
                    proc.kill()
                logger.info(f"Shell session closed: {session_id}")

    @classmethod
    def list_sessions(cls) -> list:
        with cls._lock:
            return [
                {"session_id": sid, "alive": proc.poll() is None}
                for sid, proc in cls._sessions.items()
            ]


class PersistentShellTool(BaseTool):
    """Manage persistent interactive shell sessions."""

    def get_metadata(self) -> ToolMetadata:
        return ToolMetadata(
            name="PERSISTENT_SHELL",
            description="Create and interact with persistent shell sessions.",
            category="shell",
            parameters=[
                ToolParameter("action", "string", "Action: create, send, close, list", required=True,
                              enum=["create", "send", "close", "list"]),
                ToolParameter("session_id", "string", "Session identifier"),
                ToolParameter("command", "string", "Command to send (for 'send' action)"),
                ToolParameter("timeout", "integer", "Command timeout in seconds", default=60),
            ],
            permissions=[ToolPermission.EXECUTE],
        )

    def execute(self, params: Dict[str, Any], context: ToolExecutionContext) -> ToolResult:
        action = params.get("action", "")
        session_id = params.get("session_id", "default")
        command = params.get("command", "")
        timeout = params.get("timeout", 60)

        if action == "create":
            ok = PersistentShellManager.create_session(session_id, context.get_cwd())
            return ToolResult.ok(f"Session '{session_id}' created") if ok else ToolResult.fail("Failed to create session")

        elif action == "send":
            if not command:
                return ToolResult.fail("command is required for 'send' action")
            output = PersistentShellManager.send_command(session_id, command, timeout)
            if output is None:
                return ToolResult.fail(f"Session '{session_id}' not found or dead")
            return ToolResult.ok(output)

        elif action == "close":
            PersistentShellManager.close_session(session_id)
            return ToolResult.ok(f"Session '{session_id}' closed")

        elif action == "list":
            sessions = PersistentShellManager.list_sessions()
            if not sessions:
                return ToolResult.ok("No active sessions")
            lines = [f"{'🟢' if s['alive'] else '🔴'} {s['session_id']}" for s in sessions]
            return ToolResult.ok("\n".join(lines))

        return ToolResult.fail(f"Unknown action: {action}")


def get_all_tools():
    return [RunCommandTool, PersistentShellTool]
