from __future__ import annotations
import httpx, json, re, uuid, asyncio
from typing import Any


class SandboxClient:
    def __init__(self, url: str, token: str):
        self.url = url.rstrip("/")
        self.token = token
        self._client = httpx.AsyncClient(
            timeout=httpx.Timeout(120.0, connect=10.0),
            headers={"Authorization": f"Bearer {token}"} if token else {},
        )

    async def health(self) -> bool:
        try:
            resp = await self._client.get(f"{self.url}/health", timeout=5.0)
            return resp.is_success
        except Exception:
            return False

    async def execute(self, code: str, language: str = "python") -> dict:
        resp = await self._client.post(
            f"{self.url}/execute",
            json={"code": code, "language": language},
        )
        resp.raise_for_status()
        return resp.json()

    async def install(self, packages: list[str]) -> dict:
        resp = await self._client.post(
            f"{self.url}/install",
            json={"packages": packages},
            timeout=httpx.Timeout(60.0),
        )
        resp.raise_for_status()
        return resp.json()

    async def upload_file(self, filename: str, content: str | bytes) -> dict:
        if isinstance(content, str):
            content = content.encode()
        resp = await self._client.post(
            f"{self.url}/upload",
            files={"file": (filename, content)},
        )
        resp.raise_for_status()
        return resp.json()

    async def read_file(self, path: str) -> dict:
        resp = await self._client.get(f"{self.url}/files/{path.lstrip('/')}")
        resp.raise_for_status()
        return resp.json()

    async def restart(self) -> bool:
        try:
            resp = await self._client.post(f"{self.url}/restart", timeout=30.0)
            return resp.is_success
        except Exception:
            return False

    async def close(self):
        await self._client.aclose()


_OUTPUT_PATTERNS: dict[str, re.Pattern] = {
    "image": re.compile(r"(?:!\[.*?\]\(.*?\)|data:image/\w+;base64,)", re.IGNORECASE),
    "audio": re.compile(r"data:audio/\w+;base64,", re.IGNORECASE),
    "dataframe": re.compile(r"<table[^>]*>.*?</table>", re.DOTALL | re.IGNORECASE),
    "plot": re.compile(r"(plt\.show|matplotlib|plotly|bokeh|seaborn)", re.IGNORECASE),
    "error": re.compile(r"(Traceback \(most recent call last\)|Error|Exception)", re.IGNORECASE),
}


def parse_output(output: str) -> list[dict]:
    results = []
    for kind, pattern in _OUTPUT_PATTERNS.items():
        matches = pattern.findall(output)
        if matches:
            results.append({"type": kind, "matches": len(matches)})
    if not results:
        results.append({"type": "text", "matches": 1})
    return results


async def run_code_with_fix_loop(
    client: SandboxClient,
    code: str,
    language: str = "python",
    max_iterations: int = 3,
    system_prompt: str = "",
) -> dict:
    attempt = 0
    last_error = ""
    current_code = code

    while attempt < max_iterations:
        try:
            result = await client.execute(current_code, language=language)
        except Exception as e:
            return {"success": False, "error": str(e), "attempts": attempt + 1}

        stdout = result.get("stdout", "")
        stderr = result.get("stderr", "")
        exit_code = result.get("exit_code", result.get("code", 0))

        if exit_code == 0:
            output_types = parse_output(stdout + stderr)
            return {
                "success": True,
                "stdout": stdout,
                "stderr": stderr,
                "exit_code": exit_code,
                "output_types": output_types,
                "attempts": attempt + 1,
            }

        last_error = stderr or stdout
        attempt += 1

        if attempt >= max_iterations:
            break

        current_code = _generate_fix_prompt(current_code, last_error, system_prompt)

    return {
        "success": False,
        "stdout": stdout if attempt > 0 else "",
        "stderr": last_error,
        "exit_code": exit_code if attempt > 0 else -1,
        "attempts": attempt,
    }


def _generate_fix_prompt(code: str, error: str, system_prompt: str) -> str:
    lines = code.split("\n")
    if system_prompt:
        comment_char = "#" if not lines or not lines[0].strip().startswith("#") else lines[0].strip()[:1]
        lines.insert(0, f"{comment_char} Fix attempt for error: {error[:200]}")
        lines.insert(1, "")
    return "\n".join(lines)