"""Research job submission and management."""
from __future__ import annotations

import base64
import os
import time
from pathlib import Path
from typing import TYPE_CHECKING, Generator

if TYPE_CHECKING:
    from .client import Client

# File extensions included when bundling a directory
_BUNDLE_EXTENSIONS = frozenset({
    ".py", ".yaml", ".yml", ".json", ".txt", ".toml", ".cfg", ".sh",
})

# Directory names to skip when walking a directory tree
_SKIP_DIRS = frozenset({
    "__pycache__", ".git", "node_modules", ".venv",
})

# Well-known entrypoint filenames, checked in order
_AUTO_ENTRYPOINTS = ("main.py", "train.py", "run.py")


def _bundle_directory(directory: str | Path) -> dict[str, str]:
    """Walk a directory and return {relative_path: base64_content}.

    Only includes files matching ``_BUNDLE_EXTENSIONS``. Skips
    directories in ``_SKIP_DIRS`` and ``*.pyc`` files.
    """
    root = Path(directory).resolve()
    if not root.is_dir():
        raise ValueError(f"Not a directory: {root}")

    files: dict[str, str] = {}
    for dirpath, dirnames, filenames in os.walk(root):
        # Prune skipped directories in-place
        dirnames[:] = [d for d in dirnames if d not in _SKIP_DIRS]

        for fname in filenames:
            fpath = Path(dirpath) / fname
            if fpath.suffix == ".pyc":
                continue
            if fpath.suffix not in _BUNDLE_EXTENSIONS:
                continue

            rel = fpath.relative_to(root).as_posix()
            content = fpath.read_bytes()
            files[rel] = base64.b64encode(content).decode()

    if not files:
        raise ValueError(f"No bundleable files found in {root}")

    return files


def _detect_entrypoint(files: dict[str, str]) -> str | None:
    """Return an entrypoint command if a well-known root script exists."""
    for candidate in _AUTO_ENTRYPOINTS:
        if candidate in files:
            return f"python /home/ray/code/{candidate}"
    return None


class Jobs:
    def __init__(self, client: Client):
        self._client = client

    def submit(
        self,
        name: str,
        *,
        script: str | None = None,
        code: str | None = None,
        directory: str | None = None,
        gpu: str = "rtx-4090",
        gpu_count: int = 1,
        cpu: str = "2",
        memory: str = "4Gi",
        image: str = "",
        entrypoint: str = "python /home/ray/code/main.py",
        pip_packages: list[str] | None = None,
        env_vars: dict[str, str] | None = None,
        secret_names: list[str] | None = None,
        description: str = "",
        max_runtime_hours: float = 0,
        compute_target: str | None = None,
        provider: str | None = None,
        region: str | None = None,
        sovereignty: bool = False,
        min_vram: int | None = None,
    ) -> dict:
        """Submit a training job.

        Provide exactly one of:
        - ``script``: path to a local Python file
        - ``code``: inline Python source string
        - ``directory``: path to a local directory (bundles all supported files)

        The SDK base64-encodes content before sending it to the API.
        """
        sources = sum(1 for s in (script, code, directory) if s)
        if sources == 0:
            raise ValueError("Exactly one of script, code, or directory is required")
        if sources > 1:
            raise ValueError("Provide only one of script, code, or directory")

        # Determine whether the caller left entrypoint at its default
        entrypoint_is_default = entrypoint == "python /home/ray/code/main.py"

        body: dict = {
            "name": name,
            "gpu_type": gpu,
            "gpu_count": gpu_count,
            "cpu": cpu,
            "memory": memory,
            "description": description,
            "max_runtime_hours": max_runtime_hours,
        }

        if directory:
            files_dict = _bundle_directory(directory)
            body["files"] = files_dict

            # Auto-detect entrypoint when the caller did not override it
            if entrypoint_is_default:
                detected = _detect_entrypoint(files_dict)
                if detected:
                    entrypoint = detected
        else:
            if script:
                raw = Path(script).read_text()
            else:
                raw = code  # type: ignore[assignment]
            body["code"] = base64.b64encode(raw.encode()).decode()  # type: ignore[union-attr]

        body["entrypoint"] = entrypoint

        if image:
            body["image"] = image
        if pip_packages:
            body["pip_packages"] = pip_packages
        if env_vars:
            body["env_vars"] = env_vars
        if secret_names:
            body["secret_names"] = secret_names
        if compute_target:
            body["compute_target"] = compute_target
        if provider:
            body["provider"] = provider
        if region:
            body["region"] = region
        if sovereignty:
            body["sovereignty"] = True
        if min_vram:
            body["min_vram"] = min_vram

        return self._client.post("/v1/research/jobs", json=body)

    def get(self, job_id: str) -> dict:
        """Get job details."""
        return self._client.get(f"/v1/research/jobs/{job_id}")

    def status(self, job_id: str) -> dict:
        """Get job status (alias for get)."""
        return self.get(job_id)

    def logs(self, job_id: str, tail: int = 500) -> dict:
        """Fetch job logs."""
        return self._client.get(
            f"/v1/research/jobs/{job_id}/logs", params={"tail": tail}
        )

    def cancel(self, job_id: str) -> None:
        """Cancel a running job."""
        self._client.delete(f"/v1/research/jobs/{job_id}")

    def list(self) -> list[dict]:
        """List all jobs."""
        return self._client.get("/v1/research/jobs")

    def estimate(
        self,
        gpu: str = "rtx-4090",
        gpu_count: int = 1,
        max_runtime_hours: float = 4,
    ) -> dict:
        """Get a cost estimate before submitting."""
        return self._client.post(
            "/v1/research/jobs/estimate",
            json={
                "gpu_type": gpu,
                "gpu_count": gpu_count,
                "max_runtime_hours": max_runtime_hours,
            },
        )

    def wait(
        self,
        job_id: str,
        poll_interval: float = 5.0,
        timeout: float = 3600,
    ) -> dict:
        """Block until a job reaches a terminal state.

        Uses exponential backoff starting at ``poll_interval``, capping at 30s.
        """
        deadline = time.time() + timeout
        interval = poll_interval
        while time.time() < deadline:
            result = self.get(job_id)
            status = result.get("status", "")
            if status in ("SUCCEEDED", "FAILED", "Completed", "Failed", "STOPPED"):
                return result
            time.sleep(min(interval, deadline - time.time()))
            interval = min(interval * 1.5, 30.0)
        raise TimeoutError(f"Job {job_id} did not complete within {timeout}s")

    def stream_logs(
        self,
        job_id: str,
        tail: int = 100,
        poll_interval: float = 2.0,
    ) -> Generator[str, None, None]:
        """Yield new log lines until the job reaches a terminal state."""
        seen = 0
        while True:
            result = self.logs(job_id, tail=tail)
            lines = result.get("logs", "").splitlines()
            for line in lines[seen:]:
                yield line
            seen = len(lines)

            job = self.get(job_id)
            if job.get("status") in (
                "SUCCEEDED",
                "FAILED",
                "Completed",
                "Failed",
                "STOPPED",
            ):
                # Fetch final logs
                result = self.logs(job_id, tail=tail)
                for line in result.get("logs", "").splitlines()[seen:]:
                    yield line
                return
            time.sleep(poll_interval)
