"""Dataset, judging, and sandbox helpers for the PKPO coding run.

All generated code is executed inside Modal containers. This module still applies
per-process time and memory limits and uses a guarded Python runner that blocks
network-oriented imports.
"""
from __future__ import annotations

import base64
import json
import os
import pickle
import random
import re
import resource
import subprocess
import sys
import tempfile
import textwrap
import time
import zlib
from dataclasses import dataclass, asdict
from pathlib import Path
from typing import Any

from agent_core import coding_task_message, extract_code


@dataclass
class Problem:
    source: str
    problem_id: str
    title: str
    statement: str
    public_tests: list[dict[str, str]]
    hidden_tests: list[dict[str, str]]
    starter_code: str = ""
    difficulty: str = ""
    contest_date: str = ""

    def first_user_message(self, max_turns: int = 1) -> str:
        return coding_task_message(
            f"{self.title}\n\n{self.statement}",
            starter_code=self.starter_code,
            max_turns=max_turns,
        )

    def to_public_dict(self) -> dict[str, Any]:
        d = asdict(self)
        d["hidden_tests"] = [{"input": t.get("input", "")[:80], "output_len": len(t.get("output", ""))}
                             for t in self.hidden_tests[:3]]
        d["statement"] = self.statement[:500]
        return d


def _json_or_encoded_tests(value: Any) -> list[dict[str, str]]:
    if value is None or value == "":
        return []
    if isinstance(value, list):
        return value
    if not isinstance(value, str):
        return []
    s = value.strip()
    if not s:
        return []
    try:
        obj = json.loads(s)
        return obj if isinstance(obj, list) else []
    except Exception:
        pass
    raw = zlib.decompress(base64.b64decode(s))
    try:
        obj = pickle.loads(raw)
    except Exception:
        obj = raw.decode("utf-8")
    if isinstance(obj, str):
        obj = json.loads(obj)
    return obj if isinstance(obj, list) else []


def _stdin_tests(tests: list[dict[str, Any]], limit: int | None = None) -> list[dict[str, str]]:
    out = []
    for t in tests:
        if t.get("testtype", "stdin") != "stdin":
            continue
        if "input" not in t or "output" not in t:
            continue
        out.append({"input": str(t["input"]), "output": str(t["output"]), "testtype": "stdin"})
        if limit and len(out) >= limit:
            break
    return out


def load_lcb_v6_subset(limit: int = 12, seed: int = 7341) -> list[Problem]:
    """Latest code_generation_lite release file is v6/test6.jsonl."""
    from datasets import load_dataset

    ds = load_dataset("livecodebench/code_generation_lite", "v6", split="test", trust_remote_code=True)
    candidates: list[Problem] = []
    for row in ds:
        if str(row.get("starter_code") or "").strip():
            continue
        public = _stdin_tests(_json_or_encoded_tests(row.get("public_test_cases")))
        hidden = _stdin_tests(_json_or_encoded_tests(row.get("private_test_cases")))
        if not public or not hidden:
            continue
        candidates.append(Problem(
            source="livecodebench/code_generation_lite:v6",
            problem_id=str(row["question_id"]),
            title=str(row.get("question_title") or row["question_id"]),
            statement=str(row.get("question_content") or ""),
            public_tests=public,
            hidden_tests=hidden,
            starter_code=str(row.get("starter_code") or ""),
            difficulty=str(row.get("difficulty") or ""),
            contest_date=str(row.get("contest_date") or ""),
        ))
    candidates.sort(key=lambda p: (p.contest_date, p.problem_id), reverse=True)
    rng = random.Random(seed)
    latest_pool = candidates[: max(limit * 4, limit)]
    rng.shuffle(latest_pool)
    return latest_pool[:limit]


def _tests_from_io_dict(value: Any, limit: int | None = None) -> list[dict[str, str]]:
    if not isinstance(value, dict):
        return []
    ins = value.get("input") or []
    outs = value.get("output") or []
    if len(ins) != len(outs):
        return []
    out = []
    for i, o in zip(ins, outs):
        out.append({"input": str(i), "output": str(o), "testtype": "stdin"})
        if limit and len(out) >= limit:
            break
    return out


def load_codecontest_train(
    limit: int = 12,
    seed: int = 20260709,
    exclude_problem_ids: set[str] | None = None,
    max_rows: int | None = None,
    min_cf_rating: int = 800,
    max_cf_rating: int = 1300,
) -> list[Problem]:
    """Sample valid old CodeContests train problems without an early-prefix bias.

    ``streaming=True`` yields a stable dataset prefix; shuffling only after
    prematurely stopping at ``limit`` is not a sample.  Reservoir sampling lets a
    bounded scan supply a deterministic, broader curriculum while keeping SFT and
    RL problem ids disjoint.

    A Qwen3-8B base policy earns ~0 reward on unrated/hard Codeforces problems,
    which starves PKPO of nonzero groups, so the RL pool is restricted to
    problems with a KNOWN cf_rating inside [min_cf_rating, max_cf_rating].
    """
    from datasets import load_dataset

    stream = load_dataset("deepmind/code_contests", split="train", streaming=True)
    excluded = exclude_problem_ids or set()
    rng = random.Random(seed)
    rows: list[Problem] = []
    seen = 0
    max_rows = max_rows or max(3000, limit * 45)
    for row_idx, row in enumerate(stream, start=1):
        if row_idx > max_rows:
            break
        rating = int(row.get("cf_rating") or 0)
        if rating < min_cf_rating or rating > max_cf_rating:
            continue
        desc = str(row.get("description") or "")
        if not desc or len(desc) > 4500:
            continue
        low = desc.lower()
        if "interactive" in low or "output-only" in low:
            continue
        if str(row.get("input_file") or "").strip() or str(row.get("output_file") or "").strip():
            continue
        public = _tests_from_io_dict(row.get("public_tests"), limit=3)
        generated = _tests_from_io_dict(row.get("generated_tests"), limit=18)
        private = _tests_from_io_dict(row.get("private_tests"), limit=18)
        hidden = generated or private
        problem_id = str(row.get("name") or f"codecontest-{row_idx}")
        if problem_id in excluded or not public or len(hidden) < 6:
            continue
        problem = Problem(
            source="deepmind/code_contests:train",
            problem_id=problem_id,
            title=str(row.get("name") or "CodeContests problem"),
            statement=desc,
            public_tests=public,
            hidden_tests=hidden,
            difficulty=str(row.get("difficulty") or ""),
        )
        seen += 1
        if len(rows) < limit:
            rows.append(problem)
        else:
            replace_idx = rng.randrange(seen)
            if replace_idx < limit:
                rows[replace_idx] = problem
    rng.shuffle(rows)
    return rows


def python3_verified_solutions(
    limit: int = 8,
    seed: int = 20260710,
    candidate_multiplier: int = 2,
    max_rows: int | None = None,
    min_cf_rating: int = 800,
    max_cf_rating: int = 1500,
) -> list[tuple[Problem, str]]:
    """Return runnable CodeContests Python-3 reference solutions only.

    ``solutions.language`` is a ClassLabel integer in CodeContests: 1 is Python
    (Python 2) while 3 is Python 3.  The earlier loader ignored that paired field,
    so it could SFT on ``raw_input``/bare-``print`` programs that fail our Python 3
    judge.  Each candidate here is also compiled and run on held-out generated
    tests before it is admitted to the warm-up set.
    """
    from datasets import load_dataset

    stream = load_dataset("deepmind/code_contests", split="train", streaming=True)
    pairs: list[tuple[Problem, str]] = []
    wanted = max(limit, limit * max(1, candidate_multiplier))
    max_rows = max_rows or min(6000, max(500, wanted * 40))

    for row_idx, row in enumerate(stream, start=1):
        if row_idx > max_rows:
            break
        rating = int(row.get("cf_rating") or 0)
        if rating < min_cf_rating or rating > max_cf_rating:
            continue
        desc = str(row.get("description") or "")
        if not desc or len(desc) > 4500:
            continue
        low = desc.lower()
        if "interactive" in low or "output-only" in low:
            continue
        if str(row.get("input_file") or "").strip() or str(row.get("output_file") or "").strip():
            continue
        sols = row.get("solutions") or {}
        languages = sols.get("language") or []
        solutions = sols.get("solution") or []
        public = _tests_from_io_dict(row.get("public_tests"), limit=2)
        generated = _tests_from_io_dict(row.get("generated_tests"), limit=3)
        private = _tests_from_io_dict(row.get("private_tests"), limit=3)
        held_out = generated or private
        if not public or not held_out:
            continue

        for language, raw_solution in zip(languages, solutions):
            # CodeContests' ClassLabel id 3 is PYTHON3.  Do not accept PYTHON (id
            # 1), even if it happens to compile under the local interpreter.
            if language != 3:
                continue
            solution = textwrap.dedent(str(raw_solution)).strip()
            if not 40 <= len(solution) <= 7000:
                continue
            try:
                compile(solution, "<codecontests-python3>", "exec")
            except (SyntaxError, ValueError, TypeError):
                continue
            verdict = judge_code(solution, held_out, timeout_s=3, memory_mb=768, max_tests=3)
            if not verdict.get("passed"):
                continue
            prob = Problem(
                source="deepmind/code_contests:train",
                problem_id=str(row.get("name") or f"sft-{len(pairs)}"),
                title=str(row.get("name") or "CodeContests problem"),
                statement=desc,
                public_tests=public,
                hidden_tests=held_out,
                difficulty=str(row.get("difficulty") or ""),
            )
            pairs.append((prob, solution))
            break
        if len(pairs) >= wanted:
            break

    rng = random.Random(seed)
    rng.shuffle(pairs)
    return pairs[:limit]


GUARD = r'''
import builtins
import sys

blocked = {
    "_socket", "socket", "ssl", "urllib", "http", "ftplib", "requests",
    "subprocess", "multiprocessing", "ctypes",
}
real_import = builtins.__import__

def guarded_import(name, globals=None, locals=None, fromlist=(), level=0):
    root = name.split(".", 1)[0]
    if root in blocked:
        raise ImportError(f"blocked import: {name}")
    return real_import(name, globals, locals, fromlist, level)

builtins.__import__ = guarded_import
path = sys.argv[1]
with open(path, "r", encoding="utf-8", errors="replace") as f:
    src = f.read()
ns = {"__name__": "__main__", "__file__": path}
exec(compile(src, path, "exec"), ns, ns)
'''


def _limit_child(timeout_s: int, memory_mb: int, cwd: str):
    def inner():
        os.chdir(cwd)
        cpu = max(1, int(timeout_s) + 1)
        for res, lim in [
            (resource.RLIMIT_CPU, (cpu, cpu + 1)),
            (resource.RLIMIT_AS, (memory_mb * 1024 * 1024, memory_mb * 1024 * 1024)),
            (resource.RLIMIT_FSIZE, (32 * 1024 * 1024, 32 * 1024 * 1024)),
            (resource.RLIMIT_NOFILE, (64, 64)),
        ]:
            try:
                resource.setrlimit(res, lim)
            except Exception:
                pass
    return inner


def normalize_output(text: str) -> str:
    text = text.replace("\r\n", "\n").replace("\r", "\n")
    return "\n".join(line.rstrip() for line in text.strip().split("\n")).strip()


def judge_code(code: str, tests: list[dict[str, str]], timeout_s: int = 3,
               memory_mb: int = 768, max_tests: int | None = None) -> dict[str, Any]:
    tests = _stdin_tests(tests, limit=max_tests)
    started = time.time()
    if not code.strip():
        return {
            "passed": False,
            "passed_tests": 0,
            "total_tests": len(tests),
            "error": "empty code",
            "seconds": time.time() - started,
        }
    with tempfile.TemporaryDirectory(prefix="pkpo_exec_") as td:
        sol = Path(td) / "solution.py"
        runner = Path(td) / "guarded_runner.py"
        sol.write_text(code, encoding="utf-8")
        runner.write_text(GUARD, encoding="utf-8")
        for idx, test in enumerate(tests):
            try:
                cp = subprocess.run(
                    [sys.executable, "-I", str(runner), str(sol)],
                    input=test["input"],
                    text=True,
                    stdout=subprocess.PIPE,
                    stderr=subprocess.PIPE,
                    timeout=timeout_s + 1,
                    preexec_fn=_limit_child(timeout_s, memory_mb, td),
                )
            except subprocess.TimeoutExpired:
                return {
                    "passed": False,
                    "passed_tests": idx,
                    "total_tests": len(tests),
                    "error": "timeout",
                    "seconds": time.time() - started,
                }
            if cp.returncode != 0:
                return {
                    "passed": False,
                    "passed_tests": idx,
                    "total_tests": len(tests),
                    "error": "runtime_error",
                    "exit_code": cp.returncode,
                    "stderr": cp.stderr[-500:],
                    "seconds": time.time() - started,
                }
            got = normalize_output(cp.stdout)
            want = normalize_output(test["output"])
            if got != want:
                return {
                    "passed": False,
                    "passed_tests": idx,
                    "total_tests": len(tests),
                    "error": "wrong_answer",
                    "got": got[:300],
                    "want": want[:300],
                    "seconds": time.time() - started,
                }
    return {"passed": True, "passed_tests": len(tests), "total_tests": len(tests), "seconds": time.time() - started}


def judge_final_answer(final_answer: str, tests: list[dict[str, str]], **kwargs) -> dict[str, Any]:
    return judge_code(extract_code(final_answer), tests, **kwargs)


def compact_json_dump(path: str | Path, obj: Any):
    Path(path).parent.mkdir(parents=True, exist_ok=True)
    Path(path).write_text(json.dumps(obj, indent=2, sort_keys=True) + "\n", encoding="utf-8")


def short_completion_for_sft(solution: str) -> str:
    solution = textwrap.dedent(solution).strip()
    return (
        "I will provide a direct Python solution.\n</think>\n"
        "<answer>Tool type: final\nTool query: ```python\n"
        + solution
        + "\n```</answer>"
    )