from __future__ import annotations

import subprocess
import resource
import sys
import threading
import shlex
from pathlib import Path
from datetime import datetime
from typing import BinaryIO

from one_run.backends.base import BackendQueryResult, BackendStartResult, RunContext
from one_run.errors import BackendError
from one_run.runtime import local_runtime_snapshot
from one_run.utils.shell import command_exists, run_capture
from one_run.utils.templating import render_command


class LocalBackend:
    @staticmethod
    def _tee_stream(pipe: BinaryIO, log_file: BinaryIO, console: BinaryIO | None) -> None:
        while True:
            chunk = pipe.readline()
            if not chunk:
                break
            log_file.write(chunk)
            log_file.flush()
            if console is not None:
                console.write(chunk)
                console.flush()

    def _render(self, run_ctx: RunContext) -> tuple[str, str | None]:
        if run_ctx.manifest.environment.kind == "docker":
            return self._render_docker(run_ctx)

        return (
            render_command(
                run_ctx.manifest.entrypoint.command,
                {
                    "config": str(run_ctx.paths.config_snapshot),
                    "run_dir": str(run_ctx.paths.run_dir),
                    "seed": run_ctx.manifest.experiment.seed,
                    "project_root": str(run_ctx.project_root),
                },
            ),
            None,
        )

    @staticmethod
    def _docker_image_exists(image: str) -> bool:
        proc = run_capture(["docker", "image", "inspect", image])
        return proc.returncode == 0

    @staticmethod
    def _docker_pull(image: str) -> None:
        proc = run_capture(["docker", "pull", image])
        if proc.returncode != 0:
            msg = (proc.stderr or proc.stdout).strip() or "docker pull failed"
            raise BackendError(f"docker pull failed for image '{image}': {msg}")

    @staticmethod
    def _docker_image_digest(image: str) -> str | None:
        proc = run_capture(["docker", "image", "inspect", "--format", "{{index .RepoDigests 0}}", image])
        if proc.returncode != 0:
            return None
        digest = proc.stdout.strip()
        return digest or None

    @staticmethod
    def _resolve_env_file(path_str: str, project_root: Path) -> str:
        p = Path(path_str)
        if not p.is_absolute():
            p = project_root / p
        return str(p)

    def _ensure_docker_image(self, run_ctx: RunContext, image: str) -> str | None:
        env_cfg = run_ctx.manifest.environment
        pull_policy = env_cfg.pull_policy or "if_missing"
        exists = self._docker_image_exists(image)

        if pull_policy == "always":
            self._docker_pull(image)
        elif pull_policy == "if_missing":
            if not exists:
                self._docker_pull(image)
        elif pull_policy == "never":
            if not exists:
                raise BackendError(
                    f"docker image '{image}' not found locally and pull_policy=never"
                )

        return self._docker_image_digest(image)

    def _render_docker(self, run_ctx: RunContext) -> tuple[str, str | None]:
        env_cfg = run_ctx.manifest.environment
        image = env_cfg.image
        if not image:
            raise BackendError("environment.image is required when environment.kind=docker")

        digest = self._ensure_docker_image(run_ctx, image)

        container_project_root = "/workspace"
        container_run_dir = "/run"
        container_config = f"{container_run_dir}/config.snapshot.yaml"
        workdir = env_cfg.workdir or container_project_root

        inner_command = render_command(
            run_ctx.manifest.entrypoint.command,
            {
                "config": container_config,
                "run_dir": container_run_dir,
                "seed": run_ctx.manifest.experiment.seed,
                "project_root": container_project_root,
            },
        )

        cmd_parts = [
            "docker",
            "run",
            "--rm",
            "-v",
            f"{run_ctx.project_root}:{container_project_root}",
            "-v",
            f"{run_ctx.paths.run_dir}:{container_run_dir}",
            "-w",
            workdir,
        ]

        if env_cfg.mounts:
            for mount in env_cfg.mounts:
                cmd_parts.extend(["-v", mount])
        if env_cfg.env_file:
            cmd_parts.extend(["--env-file", self._resolve_env_file(env_cfg.env_file, run_ctx.project_root)])
        if env_cfg.env:
            for key, value in env_cfg.env.items():
                cmd_parts.extend(["-e", f"{key}={value}"])

        labels = {
            "one_run_id": run_ctx.run_id,
            "one_run_experiment": run_ctx.manifest.experiment.name,
            "one_run_backend": run_ctx.manifest.backend.kind,
        }
        if env_cfg.labels:
            labels.update(env_cfg.labels)
        for key, value in labels.items():
            cmd_parts.extend(["--label", f"{key}={value}"])

        resources = run_ctx.manifest.resources
        if resources and resources.gpus:
            cmd_parts.extend(["--gpus", str(resources.gpus)])

        cmd_parts.extend([image, "/bin/sh", "-lc", inner_command])
        return " ".join(shlex.quote(part) for part in cmd_parts), digest

    def prepare(self, run_ctx: RunContext) -> None:
        if run_ctx.manifest.environment.kind == "docker" and not command_exists("docker"):
            raise BackendError("docker not found. local docker execution is unavailable")
        return

    def start(self, run_ctx: RunContext, *, emit_console: bool = True) -> BackendStartResult:
        command, image_digest = self._render(run_ctx)

        started_at = datetime.now()
        with run_ctx.paths.stdout_log.open("ab") as stdout_f, run_ctx.paths.stderr_log.open("ab") as stderr_f:
            process = subprocess.Popen(
                command,
                shell=True,
                cwd=str(run_ctx.project_root),
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
            )
            run_ctx.paths.local_pid.write_text(f"{process.pid}\n", encoding="utf-8")
            stdout_pipe = process.stdout
            stderr_pipe = process.stderr
            if stdout_pipe is None or stderr_pipe is None:
                raise RuntimeError("failed to capture process streams")

            out_thread = threading.Thread(
                target=self._tee_stream,
                args=(stdout_pipe, stdout_f, sys.stdout.buffer if emit_console else None),
                daemon=True,
            )
            err_thread = threading.Thread(
                target=self._tee_stream,
                args=(stderr_pipe, stderr_f, sys.stderr.buffer if emit_console else None),
                daemon=True,
            )
            out_thread.start()
            err_thread.start()
            try:
                exit_code = process.wait()
            except KeyboardInterrupt:
                process.terminate()
                try:
                    process.wait(timeout=10)
                except subprocess.TimeoutExpired:
                    process.kill()
                    process.wait()
                out_thread.join(timeout=2)
                err_thread.join(timeout=2)
                run_ctx.paths.local_exit_code.write_text("130\n", encoding="utf-8")
                return BackendStartResult(
                    state="failed",
                    started_at=started_at,
                    ended_at=datetime.now(),
                    exit_code=130,
                    message="interrupted",
                    docker_image_digest=image_digest,
                    runtime=local_runtime_snapshot(
                        duration_sec=(datetime.now() - started_at).total_seconds(),
                        peak_rss_raw=resource.getrusage(resource.RUSAGE_CHILDREN).ru_maxrss,
                    ).model_dump(mode="json"),
                )
            out_thread.join()
            err_thread.join()
            run_ctx.paths.local_exit_code.write_text(f"{exit_code}\n", encoding="utf-8")

        ended_at = datetime.now()
        rss_after = resource.getrusage(resource.RUSAGE_CHILDREN).ru_maxrss
        state = "success" if exit_code == 0 else "failed"
        return BackendStartResult(
            state=state,
            started_at=started_at,
            ended_at=ended_at,
            exit_code=exit_code,
            message=None if state == "success" else f"command exited with code {exit_code}",
            docker_image_digest=image_digest,
            runtime=local_runtime_snapshot(
                duration_sec=(ended_at - started_at).total_seconds(),
                peak_rss_raw=rss_after,
            ).model_dump(mode="json"),
        )

    def start_detached(self, run_ctx: RunContext) -> BackendStartResult:
        command, image_digest = self._render(run_ctx)
        started_at = datetime.now()
        wrapper = (
            f"{command}; rc=$?; "
            f"echo $rc > '{run_ctx.paths.local_exit_code}'"
        )
        with run_ctx.paths.stdout_log.open("ab") as stdout_f, run_ctx.paths.stderr_log.open("ab") as stderr_f:
            process = subprocess.Popen(
                ["/bin/sh", "-c", wrapper],
                cwd=str(run_ctx.project_root),
                stdout=stdout_f,
                stderr=stderr_f,
                start_new_session=True,
            )
        run_ctx.paths.local_pid.write_text(f"{process.pid}\n", encoding="utf-8")
        return BackendStartResult(
            state="running",
            started_at=started_at,
            ended_at=None,
            exit_code=None,
            message=None,
            docker_image_digest=image_digest,
        )

    def query(self, run_ctx: RunContext, status):
        return BackendQueryResult(state=status.state, exit_code=status.exit_code, message=status.message)

    def cancel(self, run_ctx: RunContext) -> None:
        return
