from __future__ import annotations

import subprocess
import resource
import sys
import threading
from datetime import datetime
from typing import BinaryIO

from one_run.backends.base import BackendQueryResult, BackendStartResult, RunContext
from one_run.runtime import local_runtime_snapshot
from one_run.templating import render_command


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

    def _render(self, run_ctx: RunContext) -> str:
        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),
            },
        )

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

    def start(self, run_ctx: RunContext) -> BackendStartResult:
        command = 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,
            )
            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),
                daemon=True,
            )
            err_thread = threading.Thread(
                target=self._tee_stream,
                args=(stderr_pipe, stderr_f, sys.stderr.buffer),
                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)
                return BackendStartResult(
                    state="failed",
                    started_at=started_at,
                    ended_at=datetime.now(),
                    exit_code=130,
                    message="interrupted",
                    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()

        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}",
            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 = 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,
        )

    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
