"""Task runner for executing agent-generated code and verifying results."""

import sys
import traceback
import linecache
from dataclasses import dataclass
from typing import Optional
from io import StringIO
import signal
from contextlib import contextmanager


@dataclass
class TaskResult:
    """Result of running a single task."""

    task_id: str
    success: bool
    error_message: Optional[str] = None
    execution_output: Optional[str] = None
    verification_passed: bool = False
    code_executed: Optional[str] = None

    def to_dict(self) -> dict:
        return {
            "task_id": self.task_id,
            "success": self.success,
            "error_message": self.error_message,
            "execution_output": self.execution_output,
            "verification_passed": self.verification_passed,
            "code_executed": self.code_executed,
        }


class TimeoutError(Exception):
    """Raised when code execution times out."""

    pass


@contextmanager
def timeout(seconds: int):
    """Context manager for timing out code execution."""

    def signal_handler(signum, frame):
        raise TimeoutError(f"Execution timed out after {seconds} seconds")

    # Set up the signal handler
    old_handler = signal.signal(signal.SIGALRM, signal_handler)
    signal.alarm(seconds)

    try:
        yield
    finally:
        signal.alarm(0)
        signal.signal(signal.SIGALRM, old_handler)


def register_source_with_linecache(filename: str, source: str) -> None:
    """
    Register source code with linecache so inspect.getsourcelines() works.

    This is necessary because nnsight uses inspect.getsourcelines() to capture
    source code for tracing, but exec'd code has no source file. By registering
    the source with linecache, we make it available to the inspect module.

    Args:
        filename: The pseudo-filename to register the source under
        source: The source code to register
    """
    lines = source.splitlines(keepends=True)
    # Ensure each line ends with a newline
    if lines and not lines[-1].endswith("\n"):
        lines[-1] += "\n"
    # Register with linecache - the format is (size, mtime, lines, fullname)
    linecache.cache[filename] = (
        len(source),  # size
        None,  # mtime (None means don't check for updates)
        lines,  # lines
        filename,  # fullname
    )


def execute_code(
    setup_code: str, agent_code: str, timeout_seconds: int = 60
) -> tuple[dict, str]:
    """
    Execute setup code and agent code, returning the namespace and stdout.

    Args:
        setup_code: Code to set up the environment (model loading, imports)
        agent_code: Code generated by the agent
        timeout_seconds: Maximum execution time

    Returns:
        Tuple of (namespace dict, stdout output)
    """
    # Capture stdout
    old_stdout = sys.stdout
    sys.stdout = mystdout = StringIO()

    namespace = {}

    # Create pseudo-filenames for the code
    # NOTE: These filenames must NOT start with "<nnsight" because nnsight's
    # capture() method skips those when looking for source files. Using names
    # that look like regular Python files allows inspect.getsourcelines() to
    # find the source we registered in linecache.
    setup_filename = "__nnsight_eval_setup__.py"
    agent_filename = "__nnsight_eval_agent__.py"

    try:
        with timeout(timeout_seconds):
            # Register setup code with linecache and compile with filename
            register_source_with_linecache(setup_filename, setup_code)
            setup_compiled = compile(setup_code, setup_filename, "exec")
            exec(setup_compiled, namespace)

            # Register agent code with linecache and compile with filename
            register_source_with_linecache(agent_filename, agent_code)
            agent_compiled = compile(agent_code, agent_filename, "exec")
            exec(agent_compiled, namespace)
    finally:
        sys.stdout = old_stdout
        # Clean up linecache entries
        linecache.cache.pop(setup_filename, None)
        linecache.cache.pop(agent_filename, None)

    output = mystdout.getvalue()
    return namespace, output


def run_task(task, agent_code: str) -> TaskResult:
    """
    Run a single task with agent-generated code.

    Args:
        task: Task object with setup_code, verify function, etc.
        agent_code: Code generated by the agent

    Returns:
        TaskResult with success/failure and details
    """
    try:
        # Execute the code
        namespace, output = execute_code(
            task.setup_code, agent_code, timeout_seconds=task.timeout_seconds
        )

        # Run verification
        try:
            verification_passed = task.verify(namespace)
        except Exception as e:
            return TaskResult(
                task_id=task.id,
                success=False,
                error_message=f"Verification error: {str(e)}",
                execution_output=output,
                verification_passed=False,
                code_executed=agent_code,
            )

        return TaskResult(
            task_id=task.id,
            success=verification_passed,
            error_message=None if verification_passed else "Verification failed",
            execution_output=output,
            verification_passed=verification_passed,
            code_executed=agent_code,
        )

    except TimeoutError as e:
        return TaskResult(
            task_id=task.id,
            success=False,
            error_message=str(e),
            execution_output=None,
            verification_passed=False,
            code_executed=agent_code,
        )
    except SyntaxError as e:
        return TaskResult(
            task_id=task.id,
            success=False,
            error_message=f"Syntax error: {str(e)}",
            execution_output=None,
            verification_passed=False,
            code_executed=agent_code,
        )
    except Exception as e:
        return TaskResult(
            task_id=task.id,
            success=False,
            error_message=f"Execution error: {str(e)}\n{traceback.format_exc()}",
            execution_output=None,
            verification_passed=False,
            code_executed=agent_code,
        )


def extract_code_from_response(response: str) -> str:
    """
    Extract Python code from an LLM response.

    Handles:
    - Code blocks with ```python ... ```
    - Code blocks with ``` ... ```
    - Plain code without markers
    """
    # Try to find python code block
    if "```python" in response:
        start = response.find("```python") + len("```python")
        end = response.find("```", start)
        if end != -1:
            return response[start:end].strip()

    # Try to find generic code block
    if "```" in response:
        start = response.find("```") + 3
        # Skip language identifier if present
        newline = response.find("\n", start)
        if newline != -1:
            start = newline + 1
        end = response.find("```", start)
        if end != -1:
            return response[start:end].strip()

    # Return as-is (assume it's plain code)
    return response.strip()
