"""Sandboxed Python execution tool for data processing."""

from __future__ import annotations

import asyncio
import json
import os
import tempfile
from pathlib import Path
from typing import Any

from loguru import logger

from madpod_agent.agent.tools.base import Tool

_MAX_OUTPUT = 30_000  # chars
_DEFAULT_TIMEOUT = 30  # seconds


class PythonExecTool(Tool):
    """Execute Python code in an isolated subprocess for data processing."""

    def __init__(
        self,
        workspace: Path,
        timeout: int = _DEFAULT_TIMEOUT,
        allowed_dir: Path | None = None,
    ):
        self._workspace = workspace
        self._timeout = timeout
        self._allowed_dir = allowed_dir

    @property
    def name(self) -> str:
        return "python_exec"

    @property
    def description(self) -> str:
        return (
            "Execute Python code in an isolated subprocess and return the output. "
            "Use this for data processing, CSV/JSON/Excel manipulation, calculations, "
            "text parsing, regex extraction, date math, or any computation. "
            "The code runs with the agent's workspace as the working directory. "
            "Common libraries available: json, csv, re, math, datetime, pathlib, collections, itertools. "
            "Print results to stdout — only stdout is returned."
        )

    @property
    def parameters(self) -> dict[str, Any]:
        return {
            "type": "object",
            "properties": {
                "code": {
                    "type": "string",
                    "description": "Python code to execute. Use print() to output results.",
                },
                "timeout": {
                    "type": "integer",
                    "description": f"Execution timeout in seconds (default {_DEFAULT_TIMEOUT})",
                    "minimum": 1,
                    "maximum": 120,
                },
            },
            "required": ["code"],
        }

    async def execute(self, code: str, timeout: int | None = None, **kwargs: Any) -> str:
        exec_timeout = min(timeout or self._timeout, 120)

        # Write code to a temp file to avoid shell injection
        with tempfile.NamedTemporaryFile(
            mode="w", suffix=".py", dir=str(self._workspace), delete=False
        ) as f:
            f.write(code)
            script_path = f.name

        try:
            env = os.environ.copy()
            env["PYTHONDONTWRITEBYTECODE"] = "1"
            # Restrict file access to workspace if configured
            if self._allowed_dir:
                env["MADPOD_WORKSPACE"] = str(self._allowed_dir)

            proc = await asyncio.create_subprocess_exec(
                "python3", script_path,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
                cwd=str(self._workspace),
                env=env,
            )

            try:
                stdout, stderr = await asyncio.wait_for(proc.communicate(), timeout=exec_timeout)
            except asyncio.TimeoutError:
                proc.kill()
                try:
                    await asyncio.wait_for(proc.wait(), timeout=5)
                except asyncio.TimeoutError:
                    pass
                return json.dumps({
                    "error": f"Execution timed out after {exec_timeout}s",
                    "exit_code": -1,
                })

            stdout_text = stdout.decode("utf-8", errors="replace")
            stderr_text = stderr.decode("utf-8", errors="replace")

            truncated = False
            if len(stdout_text) > _MAX_OUTPUT:
                stdout_text = stdout_text[:_MAX_OUTPUT] + "\n... (truncated)"
                truncated = True

            result: dict[str, Any] = {
                "ok": proc.returncode == 0,
                "exit_code": proc.returncode,
                "output": stdout_text,
            }
            if truncated:
                result["truncated"] = True
            if stderr_text.strip():
                result["stderr"] = stderr_text[:5000]

            return json.dumps(result)

        except Exception as exc:
            logger.error("Python exec failed: {}", exc)
            return json.dumps({"error": str(exc)})
        finally:
            Path(script_path).unlink(missing_ok=True)
