import os
import time
from dataclasses import dataclass

from .errors import HardsimJobCanceledError, HardsimJobFailedError, HardsimJobTimeoutError

TERMINAL_STATES = {"succeeded", "failed", "canceled"}


@dataclass
class Job:
    client: "HardsimClient"
    job_id: str

    def status(self) -> dict:
        return self.client.status(self.job_id)

    def cancel(self) -> dict:
        return self.client.cancel(self.job_id)

    def wait(
        self,
        poll_interval_s: float = 2.0,
        timeout_s: float = 1800.0,
        raise_on_error: bool = True,
    ) -> dict:
        start = time.time()
        while True:
            data = self.status()
            status = data["status"]
            if status in TERMINAL_STATES:
                if raise_on_error and status == "failed":
                    raise HardsimJobFailedError(job_id=self.job_id, status=status, payload=data)
                if raise_on_error and status == "canceled":
                    raise HardsimJobCanceledError(job_id=self.job_id, status=status, payload=data)
                return data
            if time.time() - start > timeout_s:
                raise HardsimJobTimeoutError(f"job {self.job_id} timed out after {timeout_s} seconds")
            time.sleep(poll_interval_s)

    def download(self, output_dir: str) -> list[str]:
        data = self.status()
        status = data.get("status")
        if status == "failed":
            raise HardsimJobFailedError(job_id=self.job_id, status=status, payload=data)
        if status == "canceled":
            raise HardsimJobCanceledError(job_id=self.job_id, status=status, payload=data)

        artifacts = data.get("artifacts", [])
        os.makedirs(output_dir, exist_ok=True)

        downloaded_paths: list[str] = []
        for artifact in artifacts:
            artifact_name = artifact["name"]
            content = self.client.download_artifact(self.job_id, artifact_name)

            path = os.path.join(output_dir, artifact_name)
            with open(path, "wb") as f:
                f.write(content)
            downloaded_paths.append(path)
        return downloaded_paths
