"""execute_python — run Python in the sandbox-runner, collect outputs and figures."""

from __future__ import annotations

import base64
import re
from typing import Any

from pydantic import BaseModel, Field

from app.core.ratelimit import MSG_SANDBOX, limiter
from app.llm.schemas import Artifact, ToolResult
from app.sandbox.client import SandboxClient
from app.services import files as file_service
from app.tools.registry import ToolContext, registry


class ExecArgs(BaseModel):
    code: str = Field(..., description="Python source code")
    description: str = Field("", max_length=300)
    heavy: bool = False


IMAGE_EXT = {".png", ".jpg", ".jpeg", ".svg"}


async def run(args: dict[str, Any], ctx: ToolContext) -> ToolResult:
    if ctx.role == "student":
        limiter.check(f"sandbox:{ctx.user_id}", ctx.settings.RATE_SANDBOX_PER_HOUR, 3600,
                      MSG_SANDBOX)
    await ctx.report("running", "Exécution du code Python…")
    code = args["code"].strip()
    m = re.match(r"^```(?:python|py)?\s*\n(.*?)\n```\s*$", code, re.S)
    if m:
        code = m.group(1)
    args["code"] = code
    client = SandboxClient(ctx.settings)
    files_in = await file_service.sandbox_inputs(ctx.user_id, ctx.file_ids)
    timeout = ctx.settings.SANDBOX_HEAVY_TIMEOUT_S if args.get("heavy") else None
    res = await client.run(args["code"], files_in, timeout)
    if res.error and res.exit_code == 127:
        return ToolResult(content=f"Le sandbox n'a pas pu exécuter le code : {res.error}. "
                          "Explique le calcul à la main ou réessaie plus tard.", error=True,
                          payload={"code": args["code"], "stderr": res.error})

    artifacts: list[Artifact] = []
    file_lines: list[str] = []
    for f in res.files_out:
        name = f.get("name", "sortie")
        data = base64.b64decode(f.get("content_b64", "")) if f.get("content_b64") else b""
        if not data:
            continue
        ext = "." + name.rsplit(".", 1)[-1].lower() if "." in name else ""
        ftype = "image" if ext in IMAGE_EXT else ("xlsx" if ext == ".xlsx" else
                                                  "csv" if ext == ".csv" else "file")
        stored = await file_service.store_artifact(
            user_id=ctx.user_id, conversation_id=ctx.conversation_id, filename=name,
            data=data, ftype=ftype)
        artifacts.append(Artifact(type=ftype, file_id=stored.id, filename=name,
                                  url=f"/api/v1/files/{stored.id}"))
        file_lines.append(f"- {name} ({len(data)} octets)")

    status = "ok" if res.exit_code == 0 else f"exit {res.exit_code}"
    content = (f"Exécution terminée ({status}, {res.duration_ms} ms).\n"
               f"STDOUT:\n{res.stdout or '(vide)'}\n")
    if res.stderr:
        content += f"STDERR:\n{res.stderr}\n"
    if file_lines:
        content += "Fichiers produits (déjà affichés à l'étudiant) :\n" + "\n".join(file_lines)
    if res.truncated:
        content += "\n(sortie tronquée)"
    return ToolResult(
        content=content, artifacts=artifacts, error=res.exit_code != 0,
        payload={"code": args["code"], "description": args.get("description", ""),
                 "stdout": res.stdout, "stderr": res.stderr, "exit_code": res.exit_code,
                 "duration_ms": res.duration_ms, "artifacts": [a.to_dict() for a in artifacts]},
        meta={"summary": args.get("description") or "Exécution Python"},
    )


registry.register("execute_python", run, ExecArgs, heavy=True)