from __future__ import annotations

import json
import os
import shlex
from datetime import datetime
from pathlib import Path
from typing import Any

import ray
from ray.job_submission import JobSubmissionClient

from .builders import BuiltCommand, build_advanced_argv, build_training_argv
from .models import AdvancedTaskSpec, JobSpec, RayConfig
from .yaml_io import dump_yaml


def _ts() -> str:
    return datetime.now().strftime("%Y%m%d_%H%M%S")


def _mkdir(p: Path) -> None:
    p.mkdir(parents=True, exist_ok=True)


def _write_text(p: Path, content: str) -> None:
    _mkdir(p.parent)
    p.write_text(content, encoding="utf-8")


def _write_json(p: Path, obj: Any) -> None:
    _write_text(p, json.dumps(obj, indent=2, ensure_ascii=False) + "\n")


def _safe_basename(path: str) -> str:
    return path.rstrip("/").split("/")[-1]


class RayJobTool:
    def __init__(self, cfg: RayConfig):
        self.cfg = cfg
        self.client = JobSubmissionClient(cfg.address)

    def _job_dir(self, submission_id: str) -> str:
        return f"{self.cfg.shared_root}/jobs/{submission_id}"

    def _runtime_env(self, spec: JobSpec, *, extra_env_vars: dict[str, str] | None = None) -> dict[str, Any]:
        return self._runtime_env_for(code_path=spec.code_path, workload=spec.workload, extra_env_vars=extra_env_vars)

    def _runtime_env_for(
        self,
        *,
        code_path: str,
        workload: str,
        extra_pythonpath: list[str] | None = None,
        extra_env_vars: dict[str, str] | None = None,
    ) -> dict[str, Any]:
        env_vars = dict(self.cfg.runtime_env_env_vars)

        # Default HF cache
        env_vars.setdefault("HF_HOME", f"{self.cfg.shared_root}/hf")
        env_vars.setdefault("HUGGINGFACE_HUB_CACHE", f"{self.cfg.shared_root}/hf/hub")
        env_vars.setdefault("TRANSFORMERS_CACHE", f"{self.cfg.shared_root}/hf/transformers")
        env_vars.setdefault("PYTHONUNBUFFERED", "1")

        # Tool code path must be importable on workers (compose mounts `/workspace/mvp` into all containers).
        # Place it before verl code to avoid interfering with verl import priority.
        tool_code_path = os.environ.get("MVP_TOOL_CODE_PATH", "/workspace/mvp/py")

        existing = env_vars.get("PYTHONPATH", "")
        base_parts = [tool_code_path, code_path, self.cfg.user_code_path]
        if extra_pythonpath:
            base_parts.extend(extra_pythonpath)
        prefix = ":".join(base_parts)
        env_vars["PYTHONPATH"] = f"{prefix}:{existing}" if existing else prefix

        # For debugging / log visibility
        env_vars["MVP_CODE_PATH"] = code_path

        # SFT: ensure ray.init() connects to the cluster
        if workload == "sft":
            env_vars.setdefault("RAY_ADDRESS", "auto")

        # Per-job extra env vars (e.g. W&B). Do not override existing keys.
        if extra_env_vars:
            for k, v in dict(extra_env_vars).items():
                if k not in env_vars:
                    env_vars[k] = str(v)

        return {"env_vars": env_vars}

    def submit(
        self,
        spec: JobSpec,
        no_wait: bool,
        job_dir: str | None = None,
        *,
        extra_env_vars: dict[str, str] | None = None,
        extra_overrides: list[str] | None = None,
    ) -> str:
        submission_id = spec.submission_id or f"mvp11_{spec.workload}_{_ts()}_{os.getpid()}"
        job_dir = job_dir or self._job_dir(submission_id)

        built = build_training_argv(spec, submission_id=submission_id, job_dir=job_dir)
        if extra_overrides:
            built = BuiltCommand(argv=[*built.argv, *list(extra_overrides)])
        entrypoint_argv = [
            "python3",
            "-m",
            "argus.ray.driver_entrypoint",
            "--job-dir",
            job_dir,
            *built.argv,
        ]
        entrypoint = " ".join(shlex.quote(x) for x in entrypoint_argv)

        runtime_env = self._runtime_env(spec, extra_env_vars=extra_env_vars)

        # Prepare job artifacts directory
        job_root = Path(job_dir)
        _mkdir(job_root / "config")
        _mkdir(job_root / "logs")
        _mkdir(job_root / "debug")
        _mkdir(job_root / "checkpoints")

        _write_text(job_root / "config" / "ray_config.yaml", dump_yaml(self.cfg.to_public_dict()))
        _write_text(job_root / "config" / "jobspec.yaml", dump_yaml(spec.to_public_dict()))
        _write_json(job_root / "config" / "submit_payload.json", {
            "submission_id": submission_id,
            "address": self.cfg.address,
            "entrypoint": entrypoint,
            "entrypoint_num_cpus": self.cfg.entrypoint_num_cpus,
            "entrypoint_resources": self.cfg.entrypoint_resources,
            "runtime_env": runtime_env,
        })

        # Pre-submit debug snapshot (ray cluster resources via ray.init)
        try:
            ray.init(address="auto", ignore_reinit_error=True, log_to_driver=False)
            _write_json(job_root / "debug" / "ray_cluster_resources_pre.json", ray.cluster_resources())
            _write_json(job_root / "debug" / "ray_available_resources_pre.json", ray.available_resources())
        except Exception as e:
            _write_text(job_root / "debug" / "ray_resources_pre.error.txt", repr(e) + "\n")

        try:
            submitted = self.client.submit_job(
                entrypoint=entrypoint,
                submission_id=submission_id,
                runtime_env=runtime_env,
                entrypoint_num_cpus=self.cfg.entrypoint_num_cpus,
                entrypoint_resources=self.cfg.entrypoint_resources,
            )
        except Exception as e:
            _write_text(job_root / "logs" / "submit.error.txt", repr(e) + "\n")
            raise

        _write_text(job_root / "config" / "ray_submission_id.txt", submitted + "\n")

        # Post-submit debug snapshot via SDK
        try:
            jobs = self.client.list_jobs()
            _write_text(
                job_root / "debug" / "ray_job_list_post.json",
                json.dumps([_job_details_to_dict(j) for j in jobs], indent=2) + "\n",
            )
        except Exception as e:
            _write_text(job_root / "debug" / "ray_job_list_post.error.txt", repr(e) + "\n")

        if not no_wait:
            # caller can separately wait; keep submit non-blocking by default in scripts
            pass

        return submitted

    def submit_advanced(
        self,
        spec: AdvancedTaskSpec,
        no_wait: bool,
        job_dir: str | None = None,
        user_id: str | None = None,
        *,
        extra_env_vars: dict[str, str] | None = None,
    ) -> str:
        submission_id = spec.submission_id or f"mvp11_{spec.workload}_{_ts()}_{os.getpid()}"
        job_dir = job_dir or self._job_dir(submission_id)

        built = build_advanced_argv(spec.command)
        entrypoint_argv = [
            "python3",
            "-m",
            "argus.ray.driver_entrypoint",
            "--job-dir",
            job_dir,
            *built.argv,
        ]
        entrypoint = " ".join(shlex.quote(x) for x in entrypoint_argv)

        extra_pythonpath: list[str] = []
        if user_id:
            extra_pythonpath.append(f"{self.cfg.shared_root}/users/{user_id}/code")

        runtime_env = self._runtime_env_for(
            code_path=self.cfg.verl_code_path,
            workload=spec.workload,
            extra_pythonpath=extra_pythonpath,
            extra_env_vars=extra_env_vars,
        )

        # Prepare job artifacts directory
        job_root = Path(job_dir)
        _mkdir(job_root / "config")
        _mkdir(job_root / "logs")
        _mkdir(job_root / "debug")
        _mkdir(job_root / "checkpoints")

        _write_text(job_root / "config" / "ray_config.yaml", dump_yaml(self.cfg.to_public_dict()))
        _write_text(job_root / "config" / "taskspec.yaml", dump_yaml(spec.to_public_dict()))
        _write_json(
            job_root / "config" / "submit_payload.json",
            {
                "submission_id": submission_id,
                "address": self.cfg.address,
                "entrypoint": entrypoint,
                "entrypoint_num_cpus": self.cfg.entrypoint_num_cpus,
                "entrypoint_resources": self.cfg.entrypoint_resources,
                "runtime_env": runtime_env,
            },
        )

        # Pre-submit debug snapshot (ray cluster resources via ray.init)
        try:
            ray.init(address="auto", ignore_reinit_error=True, log_to_driver=False)
            _write_json(job_root / "debug" / "ray_cluster_resources_pre.json", ray.cluster_resources())
            _write_json(job_root / "debug" / "ray_available_resources_pre.json", ray.available_resources())
        except Exception as e:
            _write_text(job_root / "debug" / "ray_resources_pre.error.txt", repr(e) + "\n")

        try:
            submitted = self.client.submit_job(
                entrypoint=entrypoint,
                submission_id=submission_id,
                runtime_env=runtime_env,
                entrypoint_num_cpus=self.cfg.entrypoint_num_cpus,
                entrypoint_resources=self.cfg.entrypoint_resources,
            )
        except Exception as e:
            _write_text(job_root / "logs" / "submit.error.txt", repr(e) + "\n")
            raise

        _write_text(job_root / "config" / "ray_submission_id.txt", submitted + "\n")

        # Post-submit debug snapshot via SDK
        try:
            jobs = self.client.list_jobs()
            _write_text(
                job_root / "debug" / "ray_job_list_post.json",
                json.dumps([_job_details_to_dict(j) for j in jobs], indent=2) + "\n",
            )
        except Exception as e:
            _write_text(job_root / "debug" / "ray_job_list_post.error.txt", repr(e) + "\n")

        if not no_wait:
            pass

        return submitted

    def status(self, submission_id: str) -> str:
        return str(self.client.get_job_status(submission_id))

    def stop(self, submission_id: str) -> bool:
        return bool(self.client.stop_job(submission_id))

    def logs(self, submission_id: str) -> str:
        return self.client.get_job_logs(submission_id)

    def list(self) -> list[dict[str, Any]]:
        return [_job_details_to_dict(j) for j in self.client.list_jobs()]


def _job_details_to_dict(obj: Any) -> dict[str, Any]:
    # Ray uses pydantic models internally, but depending on bundled pydantic version
    # we might get `.model_dump()` (v2) or `.dict()` (v1).
    if hasattr(obj, "model_dump"):
        return obj.model_dump()  # type: ignore[no-any-return]
    if hasattr(obj, "dict"):
        return obj.dict()  # type: ignore[no-any-return]
    if hasattr(obj, "__dict__"):
        return dict(obj.__dict__)
    return {"repr": repr(obj)}
