"""PentestGPT executor - runs PentestGPT in Docker with timeout."""

import asyncio
import subprocess
from datetime import datetime
from pathlib import Path


class PentestGPTExecutor:
    """Executes PentestGPT in Docker container with timeout handling."""

    def __init__(
        self,
        container_name: str = "pentestgpt",
        model: str | None = None
    ):
        """
        Initialize executor.

        Args:
            container_name: Name of PentestGPT Docker container
            model: Claude model to use (optional)
        """
        self.container_name = container_name
        self.model = model
        self._ensure_container_running()

    def _ensure_container_running(self):
        """Check if container is running, start if needed."""
        # Check if container exists and is running
        result = subprocess.run(
            ["docker", "ps", "--filter", f"name={self.container_name}", "--format", "{{.Names}}"],
            capture_output=True,
            text=True
        )

        if self.container_name not in result.stdout:
            # Container not running, check if it exists
            result = subprocess.run(
                ["docker", "ps", "-a", "--filter", f"name={self.container_name}", "--format", "{{.Names}}"],
                capture_output=True,
                text=True
            )

            if self.container_name in result.stdout:
                # Container exists but not running, start it
                print(f"Starting {self.container_name} container...")
                subprocess.run(
                    ["docker", "start", self.container_name],
                    capture_output=True,
                    text=True,
                    check=True
                )
                print(f"  {self.container_name} container started")
            else:
                raise RuntimeError(
                    f"Container '{self.container_name}' not found. "
                    f"Please ensure PentestGPT Docker container is created."
                )

    async def execute(
        self,
        target_url: str,
        benchmark_id: str,
        timeout_seconds: int,
        output_file: Path
    ) -> dict:
        """
        Execute PentestGPT with timeout.

        Args:
            target_url: Target URL to test
            benchmark_id: Benchmark identifier for logging
            timeout_seconds: Timeout in seconds
            output_file: Path to write output

        Returns:
            dict with 'output_lines', 'returncode', 'timed_out'
        """
        # Build command
        command = self._build_command(target_url)

        # Ensure output directory exists
        output_file.parent.mkdir(parents=True, exist_ok=True)

        print(f"  Running PentestGPT (timeout: {timeout_seconds}s)...")
        print(f"  Target: {target_url}")

        try:
            # Create subprocess
            process = await asyncio.create_subprocess_exec(
                *command,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.STDOUT
            )

            # Stream output with timeout
            try:
                output_lines = await asyncio.wait_for(
                    self._stream_output(process, output_file),
                    timeout=timeout_seconds
                )

                # Wait for process to complete
                returncode = await process.wait()

                return {
                    "output_lines": output_lines,
                    "returncode": returncode,
                    "timed_out": False
                }

            except TimeoutError:
                print(f"  ⏱ Timeout after {timeout_seconds}s")
                # Kill the process gracefully
                await self._kill_gracefully(process)

                # Read whatever output we have
                output_lines = []
                if output_file.exists():
                    with open(output_file) as f:
                        for line in f:
                            # Extract just the content (after timestamp)
                            if ' ' in line:
                                output_lines.append(line.split(' ', 1)[1].rstrip())

                return {
                    "output_lines": output_lines,
                    "returncode": -1,
                    "timed_out": True
                }

        except Exception as e:
            print(f"  ✗ Execution error: {e}")
            return {
                "output_lines": [],
                "returncode": -1,
                "timed_out": False,
                "error": str(e)
            }

    def _build_command(self, target_url: str) -> list[str]:
        """
        Build command to execute PentestGPT in Docker.

        Args:
            target_url: Target URL

        Returns:
            Command as list of strings
        """
        cmd = [
            "docker", "exec",
            "-w", "/home/pentester",  # Set working directory inside container
            self.container_name,
            "pentestgpt",
            "--target", target_url,
        ]

        # Add model if specified
        if self.model:
            cmd.extend(["--model", self.model])

        cmd.extend([
            "--raw",
            "--no-telemetry"
        ])

        return cmd

    async def _stream_output(
        self,
        process: asyncio.subprocess.Process,
        output_file: Path
    ) -> list[str]:
        """
        Stream stdout line-by-line to file and collect lines.
        Terminates early when [DONE] is detected.

        Args:
            process: Async subprocess
            output_file: File to write output

        Returns:
            List of output lines
        """
        lines = []

        with open(output_file, 'w') as f:
            while True:
                line_bytes = await process.stdout.readline()
                if not line_bytes:
                    break

                line = line_bytes.decode('utf-8', errors='replace').rstrip()

                # Write to file with timestamp
                timestamp = datetime.now().isoformat()
                f.write(f"{timestamp} {line}\n")
                f.flush()

                # Print to console (only important lines)
                if any(marker in line for marker in ["[FLAG]", "[DONE]", "[ERROR]", "[WARN]"]):
                    print(f"  {line}")

                # Collect for parsing
                lines.append(line)

                # Check for [DONE] marker - PentestGPT has completed
                if "[DONE]" in line:
                    print("  PentestGPT completed, terminating early...")
                    # Kill the process gracefully
                    await self._kill_gracefully(process)
                    break

        return lines

    async def _kill_gracefully(self, process: asyncio.subprocess.Process):
        """
        Kill process gracefully: SIGTERM -> wait -> SIGKILL.

        Args:
            process: Process to kill
        """
        try:
            # Send SIGTERM
            process.terminate()

            # Wait up to 5 seconds
            try:
                await asyncio.wait_for(process.wait(), timeout=5)
            except TimeoutError:
                # Still running, force kill
                process.kill()
                await process.wait()
        except Exception as e:
            print(f"  Warning: Error killing process: {e}")