"""Modal backend: spawn a deployed worker function (worker/modal_app.py) and poll its result."""

from __future__ import annotations

import importlib
from types import ModuleType

from .sim_backend import BackendError


class ModalClient:
    name = "modal"
    # Functions deployed as their own app (worker/modal_isaac.py), so the big image never slows the main deploy.
    APP_SUFFIX = {"isaac_render": "-isaac"}

    def __init__(self, app_name: str, modal_module: ModuleType | None = None):
        if modal_module is None:
            try:
                modal_module = importlib.import_module("modal")
            except ImportError as e:
                raise BackendError("the modal package is not installed (pip install -r requirements.txt)") from e
        self._modal = modal_module
        self.app_name = app_name

    def supports(self, function: str) -> bool:
        return True

    def submit(self, job_input: dict, execution_timeout_s: int, function: str) -> dict:
        """Start the job and return at once. Time limits are set on the deployed functions (see modal_app.py)."""
        try:
            fn = self._modal.Function.from_name(self.app_name + self.APP_SUFFIX.get(function, ""), function)
            call = fn.spawn(job_input)
        except Exception as e:
            raise BackendError(f"Modal could not start {self.app_name}/{function}: {type(e).__name__}: {e}") from e
        return {"id": call.object_id, "status": "IN_QUEUE"}

    def status(self, job_id: str) -> dict:
        exc = self._modal.exception
        try:
            output = self._modal.FunctionCall.from_id(job_id).get(timeout=0)
        except exc.OutputExpiredError:
            return {"status": "FAILED", "error": "Modal no longer holds this run's result (results expire)"}
        except exc.FunctionTimeoutError as e:
            return {"status": "TIMED_OUT", "error": str(e) or "the run hit the function's time limit"}
        except TimeoutError:  # builtin: not finished yet (Modal does not tell queued from running apart)
            return {"status": "IN_PROGRESS"}
        except exc.InputCancellation:
            return {"status": "CANCELLED", "error": "the run was cancelled"}
        except (exc.RemoteError, exc.ExecutionError, exc.InternalError) as e:
            return {"status": "FAILED", "error": f"{type(e).__name__}: {e}"}
        except Exception as e:  # network or auth trouble on our side: keep the current status and retry later
            raise BackendError(f"Modal status check failed: {type(e).__name__}: {e}") from e
        return {"status": "COMPLETED", "output": output}