from __future__ import annotations

import importlib.util
import shutil
import tempfile
from collections.abc import Callable
from importlib import resources
from pathlib import Path

import chess


WorkspaceEvalFn = Callable[[chess.Board], int]


class WorkspaceManager:
    allowed_files = {"eval.py"}

    def __init__(self) -> None:
        self.root = Path(tempfile.mkdtemp(prefix="zero960_"))
        self._write_template_files()

    def _write_template_files(self) -> None:
        template = resources.files("zero960.workspace_template").joinpath("eval.py").read_text()
        self.write_file("eval.py", template)

    def cleanup(self) -> None:
        shutil.rmtree(self.root, ignore_errors=True)

    def resolve_path(self, relative_path: str) -> Path:
        if relative_path not in self.allowed_files:
            raise ValueError(f"writes are restricted to {sorted(self.allowed_files)}")
        return self.root / relative_path

    def read_file(self, relative_path: str) -> str:
        path = self.resolve_path(relative_path)
        return path.read_text()

    def write_file(self, relative_path: str, content: str) -> None:
        path = self.resolve_path(relative_path)
        path.parent.mkdir(parents=True, exist_ok=True)
        path.write_text(content)

    def load_eval_function(self) -> WorkspaceEvalFn:
        path = self.resolve_path("eval.py")
        module_name = f"zero960_eval_{id(path)}"
        spec = importlib.util.spec_from_file_location(module_name, path)
        if spec is None or spec.loader is None:
            raise RuntimeError("failed to create module spec for eval.py")

        module = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(module)

        eval_fn = getattr(module, "evaluate", None)
        if eval_fn is None or not callable(eval_fn):
            raise RuntimeError("eval.py must define evaluate(board)")
        return eval_fn