"""Dual-sandbox pair session lifecycle orchestration."""

from __future__ import annotations

import json
import logging
import time
from dataclasses import dataclass
from typing import Any, Callable, Optional

from ..models import CaseResult, PhaseConfig, UserSubmission
from ..transport import SandboxTransport

from .router import run_pair_protocol_router

logger: logging.Logger = logging.getLogger(__name__)


@dataclass
class PairSessionDeps:
    create_sandbox: Callable[..., Any]
    destroy_sandbox: Callable[[Any], None]
    resolve_sandbox_resources: Callable[[], tuple[Optional[int], Optional[int]]]
    load_bridge_support_files: Callable[[], dict[str, bytes]]
    write_files_chunked: Callable[[Any, dict[str, bytes], str], None]
    ensure_runtime_env: Callable[[Any], None]
    install_transport_dependencies: Callable[[Any], None]
    install_problem_dependencies: Callable[[Any], None]
    build_judge_envs: Callable[[UserSubmission, Optional[str]], dict[str, str]]
    build_user_envs: Callable[[UserSubmission, Optional[str]], dict[str, str]]
    resolve_entrypoint: Callable[[], str]
    start_background_python: Callable[..., Any]
    create_transport: Callable[[Any, int], SandboxTransport]
    stop_process: Callable[[Optional[Any]], None]
    is_process_alive: Callable[[Optional[Any]], bool]
    describe_process: Callable[[Optional[Any]], str]
    is_likely_mle_exit: Callable[[Any, str], bool]
    parse_case_result: Callable[[dict[str, Any], int], CaseResult]
    attach_case_history: Callable[[CaseResult, dict[int, list[dict[str, Any]]], int], None]
    record_observation_history: Callable[[dict[int, list[dict[str, Any]]], int, dict[str, Any]], None]
    record_action_history: Callable[[dict[int, list[dict[str, Any]]], int, dict[str, Any]], None]


class _SandboxPairSession:
    def __init__(
        self,
        *,
        deps: PairSessionDeps,
        config: PhaseConfig,
        submission: UserSubmission,
        gateway_token: Optional[str],
        artifact_files: dict[str, bytes],
        user_files_bytes: dict[str, bytes],
        user_req_path: str,
        case_provider: Callable[[], Optional[int]],
        track_per_case_usage: bool,
        on_case_start: Optional[Callable[[int], None]],
        on_case_end: Optional[Callable[[int, CaseResult], None]],
        deadline: float,
        compute_step_deadline: Callable[[float], float],
        attach_llm_usage_delta: Optional[Callable[[CaseResult], CaseResult]],
    ) -> None:
        self.deps = deps
        self.config = config
        self.submission = submission
        self.gateway_token = gateway_token
        self.artifact_files = artifact_files
        self.user_files_bytes = user_files_bytes
        self.user_req_path = user_req_path
        self.case_provider = case_provider
        self.track_per_case_usage = track_per_case_usage
        self.on_case_start = on_case_start
        self.on_case_end = on_case_end
        self.deadline = deadline
        self.compute_step_deadline = compute_step_deadline
        self.attach_llm_usage_delta = attach_llm_usage_delta

        self.judge_sb: Any = None
        self.user_sb: Any = None
        self.judge_process: Optional[Any] = None
        self.user_process: Optional[Any] = None
        self.judge_transport: Optional[SandboxTransport] = None
        self.user_transport: Optional[SandboxTransport] = None
        self.cases: list[CaseResult] = []

        self.logged_user_timeout_snapshot: bool = False
        self.logged_user_exit_snapshot: bool = False
        self.judge_envs: dict[str, str] = {}
        self.user_envs: dict[str, str] = {}
        self.user_port: int = 50052

    def run(self) -> list[CaseResult]:
        try:
            self._setup_sandboxes_and_runtime()
            self._run_router()
            return list(self.cases)
        finally:
            self._cleanup()

    def _setup_sandboxes_and_runtime(self) -> None:
        cpu_count, memory_mb = self.deps.resolve_sandbox_resources()
        self.judge_sb = self.deps.create_sandbox(
            sandbox_timeout=int(self.config.sandbox_timeout),
            cpu_count=cpu_count,
            memory_mb=memory_mb,
        )
        self.user_sb = self.deps.create_sandbox(
            sandbox_timeout=int(self.config.sandbox_timeout),
            cpu_count=cpu_count,
            memory_mb=memory_mb,
        )

        judge_id = getattr(self.judge_sb, "id", "unknown")
        user_id = getattr(self.user_sb, "id", "unknown")
        logger.info(
            "[%s] pair(dynamic scheduler): judge=%s, user=%s",
            self.submission.submit_id,
            judge_id,
            user_id,
        )

        bridge_support_files = self.deps.load_bridge_support_files()
        self._prepare_judge_runtime(bridge_support_files)
        self._prepare_user_runtime(bridge_support_files)

        self.judge_envs = self.deps.build_judge_envs(self.submission, self.gateway_token)
        self.user_envs = self.deps.build_user_envs(self.submission, self.gateway_token)

        judge_port = 50051
        self.user_port = 50052
        self.judge_envs["SANDBOX_GRPC_PORT"] = str(judge_port)
        self.user_envs["SANDBOX_GRPC_PORT"] = str(self.user_port)

        entrypoint = self.deps.resolve_entrypoint()
        self.judge_process = self.deps.start_background_python(
            sandbox=self.judge_sb,
            workdir="/workspace/judge",
            script_rel=entrypoint,
            envs=self.judge_envs,
            stderr_path="/workspace/judge_stderr.log",
        )
        logger.info(
            "[daemon] started judge runtime: %s",
            self.deps.describe_process(self.judge_process),
        )

        self.judge_transport = self.deps.create_transport(self.judge_sb, judge_port)
        judge_ready_timeout = max(5, min(60, max(1, int(self.deadline - time.time()))))
        if not self.judge_transport.wait_for_ready(timeout=judge_ready_timeout):
            raise RuntimeError("judge bridge not ready")

    def _prepare_judge_runtime(self, bridge_support_files: dict[str, bytes]) -> None:
        self.judge_sb.commands.run(
            "sudo mkdir -p /workspace/judge && sudo chmod 777 /workspace",
            timeout=30,
        )
        judge_runtime_files = dict(self.artifact_files)
        judge_runtime_files.update(bridge_support_files)
        self.deps.write_files_chunked(self.judge_sb, judge_runtime_files, "/workspace/judge")
        self.deps.ensure_runtime_env(self.judge_sb)
        self.deps.install_transport_dependencies(self.judge_sb)
        self.deps.install_problem_dependencies(self.judge_sb)

    def _prepare_user_runtime(self, bridge_support_files: dict[str, bytes]) -> None:
        self.user_sb.commands.run(
            "sudo mkdir -p /workspace/user && sudo chmod 777 /workspace",
            timeout=30,
        )
        user_runtime_files = dict(self.user_files_bytes)
        user_runtime_files.update(bridge_support_files)
        self.deps.write_files_chunked(self.user_sb, user_runtime_files, "/workspace/user")
        self.deps.ensure_runtime_env(self.user_sb)
        self.deps.install_transport_dependencies(self.user_sb)
        self.deps.install_problem_dependencies(self.user_sb)
        if self.user_req_path:
            self.user_sb.commands.run(
                f"cd /workspace && source .venv/bin/activate && "
                f"uv pip install -r /workspace/user/{self.user_req_path} -q",
                timeout=int(self.config.user_deps_timeout),
            )

    def _run_router(self) -> None:
        route_state = run_pair_protocol_router(
            submission_id=self.submission.submit_id,
            deadline=self.deadline,
            case_provider=self.case_provider,
            compute_step_deadline=self.compute_step_deadline,
            poll_judge_line=self._poll_judge_line,
            restart_user_runtime=self._restart_user_process_for_case,
            ensure_user_runtime=self._ensure_user_runtime,
            request_user_action=self._request_user_action,
            send_to_judge=self._send_to_judge,
            send_to_user=self._send_to_user,
            parse_case_result=self.deps.parse_case_result,
            attach_case_history=self.deps.attach_case_history,
            record_observation_history=self.deps.record_observation_history,
            record_action_history=self.deps.record_action_history,
            on_case_start=self.on_case_start,
            on_case_end=self.on_case_end,
            track_per_case_usage=self.track_per_case_usage,
            attach_llm_usage_delta=self.attach_llm_usage_delta,
        )
        self.cases.extend(route_state.cases)

    def _poll_judge_line(self, stop_deadline: float) -> tuple[Optional[str], bool]:
        if not self.judge_transport:
            return None, True
        return self._poll_transport_line_with_process_exit_until_deadline(
            transport=self.judge_transport,
            process=self.judge_process,
            stop_deadline=stop_deadline,
            poll_seconds=3,
        )

    def _ensure_user_runtime(self) -> None:
        if self.user_process is None:
            self._restart_user_process_for_case()

    def _restart_user_process_for_case(self) -> None:
        if self.user_transport:
            try:
                self.user_transport.close()
            except Exception:
                pass
            self.user_transport = None

        if self.user_process is not None:
            self.deps.stop_process(self.user_process)

        self.user_process = self.deps.start_background_python(
            sandbox=self.user_sb,
            workdir="/workspace/user",
            script_rel="_agent_wrapper.py",
            envs=self.user_envs,
            stderr_path="/workspace/user_stderr.log",
        )
        logger.info(
            "[daemon] started user runtime: %s",
            self.deps.describe_process(self.user_process),
        )

        self.user_transport = self.deps.create_transport(self.user_sb, self.user_port)
        ready_timeout = max(3, min(30, max(1, int(self.deadline - time.time()))))
        if not self.user_transport.wait_for_ready(timeout=ready_timeout):
            self._log_user_runtime_snapshot("probe_timeout")
            raise RuntimeError("user bridge not ready")

    def _poll_transport_line_with_process_exit_until_deadline(
        self,
        *,
        transport: SandboxTransport,
        process: Optional[Any],
        stop_deadline: float,
        poll_seconds: int,
    ) -> tuple[Optional[str], bool]:
        while True:
            remaining = stop_deadline - time.time()
            if remaining <= 0:
                return None, False
            slice_timeout = max(1, int(min(float(poll_seconds), remaining)))
            try:
                line = transport.recv_message(timeout=slice_timeout)
            except Exception:
                if not self.deps.is_process_alive(process):
                    return None, True
                return None, False
            if line is not None:
                return line, False
            if not self.deps.is_process_alive(process):
                return None, True

    def _request_user_action(self, trigger_msg: dict[str, Any]) -> dict[str, Any]:
        user_msg = dict(trigger_msg)
        user_msg.pop("history_events", None)
        if not self.user_transport:
            return {
                "type": "action",
                "error": "user transport unavailable",
                "status": "error",
            }

        remaining = max(1, int(self.deadline - time.time()))
        try:
            self.user_transport.send_message(user_msg, timeout=remaining)
        except Exception as exc:
            return {
                "type": "action",
                "error": f"user transport send failed: {type(exc).__name__}: {exc}",
                "status": "error",
            }

        user_wait_deadline = self.compute_step_deadline(self.deadline)
        user_line, user_process_exited = self._poll_transport_line_with_process_exit_until_deadline(
            transport=self.user_transport,
            process=self.user_process,
            stop_deadline=user_wait_deadline,
            poll_seconds=2,
        )
        if user_line is None:
            if user_process_exited:
                self._log_user_runtime_snapshot("process_exited")
                is_mle = self.deps.is_likely_mle_exit(
                    self.user_sb,
                    "/workspace/user_stderr.log",
                )
                action_msg: dict[str, Any] = {
                    "type": "action",
                    "error": "user process exited",
                    "status": "error",
                }
                if is_mle:
                    action_msg["status"] = "mle"
                    action_msg["error"] = "user process exited (mle)"
                return action_msg
            action_msg = {
                "type": "action",
                "error": (
                    "user idle timeout"
                    if user_wait_deadline < self.deadline
                    else "user timeout"
                ),
                "status": "tle",
            }
            self._log_user_runtime_snapshot(action_msg["error"])
            return action_msg

        try:
            action_msg = json.loads(user_line)
        except Exception:
            action_msg = {
                "type": "action",
                "error": f"invalid user output: {user_line[:200]}",
                "status": "error",
            }
        if action_msg.get("error") and "status" not in action_msg:
            action_msg["status"] = "error"
        return action_msg

    def _send_to_judge(self, msg: dict[str, Any], timeout: int) -> None:
        if not self.judge_transport:
            raise RuntimeError("judge transport unavailable")
        self.judge_transport.send_message(msg, timeout=timeout)

    def _send_to_user(self, msg: dict[str, Any], timeout: int) -> None:
        if not self.user_transport:
            raise RuntimeError("user transport unavailable")
        self.user_transport.send_message(msg, timeout=timeout)

    def _log_user_runtime_snapshot(self, reason: str) -> None:
        if reason == "process_exited":
            if self.logged_user_exit_snapshot:
                return
            self.logged_user_exit_snapshot = True
        elif reason in ("user idle timeout", "user timeout"):
            if self.logged_user_timeout_snapshot:
                return
            self.logged_user_timeout_snapshot = True

        stderr_snippet = ""
        process_desc = self.deps.describe_process(self.user_process)
        try:
            r_stderr = self.user_sb.commands.run(
                "bash -lc \"tail -c 8192 /workspace/user_stderr.log 2>/dev/null || true\"",
                timeout=6,
            )
            stderr_text = (r_stderr.stdout or "")
            stderr_snippet = stderr_text.strip().replace("\n", "\\n")
            if len(stderr_snippet) > 600:
                stderr_snippet = f"{stderr_snippet[:600]}...(truncated)"
        except Exception:
            pass

        logger.warning(
            "[%s] user runtime snapshot (%s): process=%s, stderr=%r",
            self.submission.submit_id,
            reason,
            process_desc,
            stderr_snippet,
        )

    def _cleanup(self) -> None:
        for transport in [self.judge_transport, self.user_transport]:
            try:
                if transport:
                    transport.close()
            except Exception:
                pass

        if not self.cases:
            for label, sb, stderr_path, process in [
                ("judge", self.judge_sb, "/workspace/judge_stderr.log", self.judge_process),
                ("user", self.user_sb, "/workspace/user_stderr.log", self.user_process),
            ]:
                try:
                    if sb:
                        r = sb.commands.run(
                            f"bash -lc \"tail -c 4096 {stderr_path} 2>/dev/null || echo '[file not found]'\"",
                            timeout=10,
                        )
                        snippet = (r.stdout or "").strip()
                        if snippet and snippet != "[file not found]":
                            logger.error(
                                "[%s] %s stderr (%s): %s",
                                self.submission.submit_id,
                                label,
                                self.deps.describe_process(process),
                                snippet[:2000],
                            )
                        else:
                            logger.warning(
                                "[%s] %s stderr empty or missing (%s)",
                                self.submission.submit_id,
                                label,
                                self.deps.describe_process(process),
                            )
                except Exception as diag_err:
                    logger.warning(
                        "[%s] %s diag failed: %s",
                        self.submission.submit_id,
                        label,
                        diag_err,
                    )

        for process in [self.judge_process, self.user_process]:
            try:
                self.deps.stop_process(process)
            except Exception:
                pass
        for sb in [self.judge_sb, self.user_sb]:
            try:
                if sb:
                    self.deps.destroy_sandbox(sb)
            except Exception:
                pass


def run_sandbox_pair_session(
    *,
    deps: PairSessionDeps,
    config: PhaseConfig,
    submission: UserSubmission,
    gateway_token: Optional[str],
    artifact_files: dict[str, bytes],
    user_files_bytes: dict[str, bytes],
    user_req_path: str,
    case_provider: Callable[[], Optional[int]],
    track_per_case_usage: bool,
    on_case_start: Optional[Callable[[int], None]],
    on_case_end: Optional[Callable[[int, CaseResult], None]],
    deadline: float,
    compute_step_deadline: Callable[[float], float],
    attach_llm_usage_delta: Optional[Callable[[CaseResult], CaseResult]],
) -> list[CaseResult]:
    session = _SandboxPairSession(
        deps=deps,
        config=config,
        submission=submission,
        gateway_token=gateway_token,
        artifact_files=artifact_files,
        user_files_bytes=user_files_bytes,
        user_req_path=user_req_path,
        case_provider=case_provider,
        track_per_case_usage=track_per_case_usage,
        on_case_start=on_case_start,
        on_case_end=on_case_end,
        deadline=deadline,
        compute_step_deadline=compute_step_deadline,
        attach_llm_usage_delta=attach_llm_usage_delta,
    )
    return session.run()
