"""Dual-sandbox evaluator orchestration entry point."""

from __future__ import annotations

from collections import deque
import logging
import hashlib
import threading
import time
from concurrent.futures import ThreadPoolExecutor
from typing import Any, Optional, Callable

from .base import BaseEvaluator
from .models import (
    PhaseResult,
    CaseResult,
    UserSubmission,
    CaseStatus,
    PhaseStatus,
)
from .sandbox_pool import create_sandbox, destroy_sandbox
from .runtime.artifact import (
    download_artifact as runtime_download_artifact,
    extract_artifact as runtime_extract_artifact,
    extract_requirement_pkg_name as runtime_extract_requirement_pkg_name,
    filter_requirements as runtime_filter_requirements,
    resolve_entrypoint as runtime_resolve_entrypoint,
)
from .runtime.gateway import (
    attach_llm_usage_delta as runtime_attach_llm_usage_delta,
    create_gateway_token_for_user as runtime_create_gateway_token_for_user,
    revoke_gateway_token as runtime_revoke_gateway_token,
)
from .runtime.history import (
    attach_case_history as runtime_attach_case_history,
    record_action_history as runtime_record_action_history,
    record_observation_history as runtime_record_observation_history,
)
from .runtime.process import SandboxProcessManager
from .runtime.protocol import MessageType
from .runtime.results import parse_case_result as runtime_parse_case_result
from .runtime.pair_session import PairSessionDeps, run_sandbox_pair_session
from .runtime.sandbox import (
    build_judge_envs as runtime_build_judge_envs,
    build_user_envs as runtime_build_user_envs,
    create_grpc_transport as runtime_create_grpc_transport,
    ensure_runtime_env as runtime_ensure_runtime_env,
    install_problem_dependencies as runtime_install_problem_dependencies,
    install_transport_dependencies as runtime_install_transport_dependencies,
    is_likely_mle_exit as runtime_is_likely_mle_exit,
    load_grpc_bridge_support_files as runtime_load_grpc_bridge_support_files,
    resolve_sandbox_resources as runtime_resolve_sandbox_resources,
    write_files_chunked as runtime_write_files_chunked,
)

logger: logging.Logger = logging.getLogger(__name__)
__all__ = ["DualSandboxEvaluator", "MessageType"]


class DualSandboxEvaluator(BaseEvaluator):

    _judge_sandbox: Any = None
    _user_sandbox: Any = None
    _gateway_token_info: Optional[dict[str, Any]] = None
    _prev_usage_chars: int = 0
    _prev_usage_requests: int = 0

    def cleanup(self) -> None:
        try:
            if self._judge_sandbox:
                destroy_sandbox(self._judge_sandbox)
        except Exception:
            pass
        try:
            if self._user_sandbox:
                destroy_sandbox(self._user_sandbox)
        except Exception:
            pass
        self._judge_sandbox = None
        self._user_sandbox = None

    def _run_sandbox_pair(
        self,
        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,
    ) -> list[CaseResult]:
        deps = PairSessionDeps(
            create_sandbox=create_sandbox,
            destroy_sandbox=destroy_sandbox,
            resolve_sandbox_resources=lambda: runtime_resolve_sandbox_resources(self.config),
            load_bridge_support_files=lambda: runtime_load_grpc_bridge_support_files(__file__),
            write_files_chunked=runtime_write_files_chunked,
            ensure_runtime_env=lambda sandbox: runtime_ensure_runtime_env(
                sandbox,
                int(self.config.user_deps_timeout),
            ),
            install_transport_dependencies=lambda sandbox: runtime_install_transport_dependencies(
                sandbox,
                int(self.config.user_deps_timeout),
            ),
            install_problem_dependencies=lambda sandbox: runtime_install_problem_dependencies(
                config=self.config,
                sandbox=sandbox,
                deps_timeout_seconds=int(self.config.user_deps_timeout),
                extract_requirement_pkg_name=runtime_extract_requirement_pkg_name,
            ),
            build_judge_envs=lambda sub, token: runtime_build_judge_envs(
                config=self.config,
                submission=sub,
                get_client=self._get_client,
                gateway_token=token,
            ),
            build_user_envs=lambda sub, token: runtime_build_user_envs(
                config=self.config,
                submission=sub,
                get_client=self._get_client,
                gateway_token=token,
            ),
            resolve_entrypoint=lambda: runtime_resolve_entrypoint(self.config),
            start_background_python=SandboxProcessManager.start_background_python,
            create_transport=runtime_create_grpc_transport,
            stop_process=SandboxProcessManager.stop_process,
            is_process_alive=SandboxProcessManager.is_process_alive,
            describe_process=SandboxProcessManager.describe_process,
            is_likely_mle_exit=lambda sandbox, stderr_path: runtime_is_likely_mle_exit(
                self.config,
                sandbox,
                stderr_path,
            ),
            parse_case_result=runtime_parse_case_result,
            attach_case_history=runtime_attach_case_history,
            record_observation_history=runtime_record_observation_history,
            record_action_history=runtime_record_action_history,
        )
        return run_sandbox_pair_session(
            deps=deps,
            config=self.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=self._with_step_deadline,
            attach_llm_usage_delta=(
                (lambda c: runtime_attach_llm_usage_delta(self, c, submission))
                if track_per_case_usage
                else None
            ),
        )

    def _run_parallel_pairs(
        self,
        submission: UserSubmission,
        gateway_token: Optional[str],
        artifact_files: dict[str, bytes],
        user_files_bytes: dict[str, bytes],
        user_req_path: str,
        num_cases: int,
        parallel: int,
        on_case_start: Optional[Callable[[int], None]],
        on_case_end: Optional[Callable[[int, CaseResult], None]],
        deadline: float,
        track_per_case_usage: bool = False,
    ) -> list[CaseResult]:
        active_pairs = max(1, min(parallel, num_cases))
        case_queue: deque[int] = deque(range(num_cases))
        queue_lock = threading.Lock()

        logger.info(
            f"[{submission.submit_id}] dynamic eval: pairs={active_pairs}, num_cases={num_cases}"
        )

        lock = threading.Lock()

        def safe_on_case_start(idx: int) -> None:
            if on_case_start:
                with lock:
                    on_case_start(idx)

        def safe_on_case_end(idx: int, result: CaseResult) -> None:
            if on_case_end:
                with lock:
                    on_case_end(idx, result)

        def pop_next_case() -> Optional[int]:
            with queue_lock:
                if not case_queue:
                    return None
                return case_queue.popleft()

        def run_group() -> list[CaseResult]:
            try:
                return self._run_sandbox_pair(
                    submission=submission,
                    gateway_token=gateway_token,
                    artifact_files=artifact_files,
                    user_files_bytes=user_files_bytes,
                    user_req_path=user_req_path,
                    case_provider=pop_next_case,
                    track_per_case_usage=track_per_case_usage,
                    on_case_start=safe_on_case_start,
                    on_case_end=safe_on_case_end,
                    deadline=deadline,
                )
            except Exception as exc:
                logger.exception(
                    f"[{submission.submit_id}] dynamic group failed: {exc}"
                )
                return []

        all_cases: list[CaseResult] = []

        with ThreadPoolExecutor(max_workers=active_pairs) as pool:
            futures = [pool.submit(run_group) for _ in range(active_pairs)]
            for fut in futures:
                group_cases = fut.result()
                all_cases.extend(group_cases)

        all_cases.sort(key=lambda c: c.case_index)
        return all_cases

    def evaluate(
        self,
        submission: UserSubmission,
        parallel_cases: int = 1,
        on_case_start: Optional[Callable[[int], None]] = None,
        on_case_end: Optional[Callable[[int, CaseResult], None]] = None,
    ) -> PhaseResult:
        start_time = time.time()

        try:
            num_cases = self.config.num_cases
            is_parallel = parallel_cases > 1 and num_cases > 1
            effective_parallel = (
                min(parallel_cases, num_cases) if is_parallel else 1
            )

            try:
                gateway_token = runtime_create_gateway_token_for_user(
                    self,
                    submission,
                    aggregate_limit=True,
                )
            except ValueError as e:
                return PhaseResult(
                    status=PhaseStatus.ERROR,
                    error=str(e),
                    total_time=int((time.time() - start_time) * 1000),
                )
            self._prev_usage_chars = 0
            self._prev_usage_requests = 0

            artifact_bytes = runtime_download_artifact(
                self.config.artifact_url, self.config.artifact_checksum,
            )
            artifact_sha256 = hashlib.sha256(artifact_bytes).hexdigest()
            artifact_files = runtime_extract_artifact(artifact_bytes)
            entrypoint = runtime_resolve_entrypoint(self.config)
            entrypoint_bytes = artifact_files.get(entrypoint)
            entrypoint_sha256 = (
                hashlib.sha256(entrypoint_bytes).hexdigest()
                if entrypoint_bytes is not None
                else None
            )
            bridge_path = "_user_bridge.py (generated)"
            bridge_sha256 = None
            bridge_has_action_request = False
            logger.info(
                f"[{submission.submit_id}] artifact audit: "
                f"url={self.config.artifact_url!r}, "
                f"config_checksum={self.config.artifact_checksum!r}, "
                f"downloaded_sha256={artifact_sha256}, "
                f"entrypoint={entrypoint!r}, "
                f"entrypoint_sha256={entrypoint_sha256!r}, "
                f"entrypoint_exists={entrypoint_bytes is not None}, "
                f"user_bridge={bridge_path!r}, "
                f"user_bridge_sha256={bridge_sha256!r}, "
                f"user_bridge_exists={True}, "
                f"user_bridge_has_action_request={bridge_has_action_request}, "
                f"artifact_file_count={len(artifact_files)}"
            )
            if entrypoint_bytes is None:
                sample = sorted(list(artifact_files.keys()))[:20]
                logger.warning(
                    f"[{submission.submit_id}] artifact entrypoint not found: "
                    f"{entrypoint!r}, sample_files={sample}"
                )

            user_files_text = submission.code_files or {}
            if not user_files_text:
                user_files_text = self._get_client().download_code(
                    submission.code_url,
                    expected_checksum=submission.code_checksum,
                )
            if "requirements.txt" not in user_files_text:
                raise ValueError("Missing requirements.txt (submission zip root must include this file)")
            if (
                "requirements.txt" in user_files_text
                and self.config.allowed_packages
            ):
                filtered = runtime_filter_requirements(
                    user_files_text["requirements.txt"],
                    self.config.allowed_packages,
                )
                user_files_text["requirements.txt"] = filtered

            user_files_bytes: dict[str, bytes] = {
                p: s.encode("utf-8") for p, s in user_files_text.items()
            }
            user_files_bytes["_agent_wrapper.py"] = (
                self._generate_user_wrapper().encode("utf-8")
            )
            problem_adapter_path = "sandbox/user_adapter.py"
            if problem_adapter_path in artifact_files:
                user_files_bytes["eval_runtime/problem_adapter.py"] = artifact_files[
                    problem_adapter_path
                ]

            solve_attr_name = str(getattr(self.config, "solve_attr_name", "") or "").strip()
            adapter_preset = str(getattr(self.config, "adapter_preset", "") or "").strip()
            if not solve_attr_name:
                raise ValueError("phase_config.solve_attr_name must not be empty")
            if not adapter_preset:
                raise ValueError("phase_config.adapter_preset must not be empty")
            if "sandbox/user_adapter.py" not in artifact_files:
                raise ValueError("Missing sandbox/user_adapter.py")
            user_files_bytes["_user_bridge.py"] = self._generate_managed_user_bridge(
                solve_attr_name=solve_attr_name,
                adapter_preset=adapter_preset,
            ).encode("utf-8")

            user_req_path = (
                "requirements.txt"
                if "requirements.txt" in user_files_text
                else ""
            )

            eval_start_time = time.time()
            deadline = eval_start_time + float(self.config.sandbox_timeout)

            cases = self._run_parallel_pairs(
                submission=submission,
                gateway_token=gateway_token,
                artifact_files=artifact_files,
                user_files_bytes=user_files_bytes,
                user_req_path=user_req_path,
                num_cases=num_cases,
                parallel=effective_parallel,
                on_case_start=on_case_start,
                on_case_end=on_case_end,
                deadline=deadline,
                track_per_case_usage=(effective_parallel <= 1),
            )

            observed = {
                int(c.case_index)
                for c in cases
                if isinstance(c.case_index, int) and 0 <= int(c.case_index) < num_cases
            }
            missing = [idx for idx in range(num_cases) if idx not in observed]
            phase_error: Optional[str] = None
            if missing:
                logger.error(
                    f"[{submission.submit_id}] evaluation interrupted: "
                    f"missing {len(missing)}/{num_cases} cases, mark as error"
                )
                phase_error = (
                    f"evaluation interrupted: missing {len(missing)}/{num_cases} cases"
                )
                for idx in missing:
                    cases.append(
                        CaseResult(
                            case_index=idx,
                            status=CaseStatus.ERROR,
                            score=0,
                            error="evaluation interrupted before case completion",
                        )
                    )
                cases.sort(key=lambda c: c.case_index)

            wall_elapsed_ms = int((time.time() - eval_start_time) * 1000)
            total_cases = len(cases)
            passed_cases = sum(1 for c in cases if c.status == CaseStatus.PASSED)
            score = sum(int(c.score or 0) for c in cases)
            cases_elapsed_ms = sum(max(0, int(c.time_used or 0)) for c in cases)
            result_total_time_ms = cases_elapsed_ms if cases else wall_elapsed_ms

            result_total_chars = sum(c.chars_used for c in cases)
            result_total_requests = sum(c.requests_used for c in cases)

            if is_parallel and self._gateway_token_info:
                try:
                    usage = self._get_client().get_gateway_token_usage(
                        submission.submit_id,
                    )
                    if usage:
                        result_total_chars = int(
                            usage.get("used_chars", 0)
                        )
                        result_total_requests = int(
                            usage.get("used_requests", 0)
                        )
                except Exception:
                    pass

            return PhaseResult(
                status=(
                    PhaseStatus.ERROR
                    if phase_error
                    else (
                        PhaseStatus.SUCCESS
                        if passed_cases == total_cases
                        else PhaseStatus.FAILED
                    )
                ),
                score=score,
                passed_cases=passed_cases,
                total_cases=total_cases,
                cases=cases,
                total_time=result_total_time_ms,
                total_chars=result_total_chars,
                total_requests=result_total_requests,
                error=phase_error,
            )

        except Exception as e:
            logger.exception(
                f"[{submission.submit_id}] dual sandbox error: {e}"
            )
            return PhaseResult(
                status=PhaseStatus.ERROR,
                error=f"Dual-sandbox evaluation system error: {type(e).__name__}: {str(e)}",
                total_time=int((time.time() - start_time) * 1000),
            )
        finally:
            runtime_revoke_gateway_token(self, submission)

    def _with_step_deadline(self, deadline: float) -> float:
        raw_idle = getattr(self.config, "case_idle_timeout", 0)
        try:
            idle_sec = int(raw_idle)
        except Exception:
            idle_sec = 0
        if idle_sec <= 0:
            return deadline
        return min(deadline, time.time() + float(idle_sec))

    def _generate_user_wrapper(self) -> str:
        return r'''#!/usr/bin/env python3
from __future__ import annotations

import sys
import traceback


def main() -> None:
    try:
        from _user_bridge import serve  # type: ignore
    except ImportError as e:
        print(f"[wrapper] import _user_bridge failed: {e}", file=sys.stderr, flush=True)
        raise RuntimeError("bridge-only mode requires _user_bridge.py") from e
    try:
        serve()
    except Exception as e:
        print(f"[wrapper] bridge exception: {type(e).__name__}: {e}", file=sys.stderr, flush=True)
        print(traceback.format_exc(), file=sys.stderr, flush=True)
        raise


if __name__ == "__main__":
    main()
'''

    def _generate_managed_user_bridge(
        self,
        *,
        solve_attr_name: str,
        adapter_preset: str,
    ) -> str:
        return f'''#!/usr/bin/env python3
from __future__ import annotations

from importlib import import_module


def serve() -> None:
    try:
        runtime_mod = import_module("eval_runtime.user_runtime")
    except ImportError:
        runtime_mod = import_module("agent_genesis.runtime.user_runtime")
    serve_user_runtime = getattr(runtime_mod, "serve_user_runtime")
    serve_user_runtime(
        solve_attr_name={repr(solve_attr_name)},
        adapter_preset={repr(adapter_preset)},
    )


if __name__ == "__main__":
    serve()
'''


