"""Filesystem and validation helpers for SWE-Bench Lite runs."""

from __future__ import annotations

import os
import re
import tempfile
from collections.abc import Iterable
from datetime import datetime, timezone
from pathlib import Path

from .. import _json as json
from .types import VALID_CONDITIONS, PredictionRecord

_RUN_ID_RE = re.compile(r"^[A-Za-z0-9._-]+$")


def repo_root() -> Path:
    return Path(__file__).resolve().parents[2]


def swebench_runs_root(*, root: Path | None = None) -> Path:
    base = root or repo_root()
    return base / "evals" / "swebench_runs"


def validate_run_id(run_id: str) -> str:
    value = run_id.strip()
    if not value:
        raise ValueError("run_id must be non-empty")
    if _RUN_ID_RE.fullmatch(value) is None:
        raise ValueError(
            "run_id contains unsupported characters; allowed: letters, digits, '.', '_', '-'"
        )
    return value


def run_dir_for_id(run_id: str, *, root: Path | None = None) -> Path:
    safe = validate_run_id(run_id)
    return swebench_runs_root(root=root) / safe


def parse_conditions(raw: str | Iterable[str] | None) -> list[str]:
    if raw is None:
        return list(VALID_CONDITIONS)
    if isinstance(raw, str):
        tokens = [chunk.strip() for chunk in raw.split(",") if chunk.strip()]
    else:
        tokens = [str(chunk).strip() for chunk in raw if str(chunk).strip()]
    if not tokens:
        return list(VALID_CONDITIONS)

    out: list[str] = []
    seen: set[str] = set()
    for token in tokens:
        if token not in VALID_CONDITIONS:
            raise ValueError(
                f"unsupported condition {token!r}; expected one of {', '.join(VALID_CONDITIONS)}"
            )
        if token in seen:
            continue
        seen.add(token)
        out.append(token)
    return out


def parse_paired_conditions(raw: str | Iterable[str] | None) -> list[str]:
    parsed = parse_conditions(raw)
    expected = set(VALID_CONDITIONS)
    if set(parsed) != expected or len(parsed) != len(expected):
        raise ValueError(
            "conditions must include exactly baseline,llmdebug "
            "(both conditions are mandatory for SWE-Bench Lite runs)"
        )
    return list(VALID_CONDITIONS)


def prediction_path(run_dir: Path, condition: str) -> Path:
    _assert_condition(condition)
    return run_dir / condition / "predictions.jsonl"


def metadata_path(run_dir: Path) -> Path:
    return run_dir / "run_metadata.json"


def instance_ids_path(run_dir: Path) -> Path:
    return run_dir / "instance_ids.txt"


def copy_instructions_path(run_dir: Path) -> Path:
    return run_dir / "COPY_INSTRUCTIONS.txt"


def run_index_path(run_dir: Path) -> Path:
    return run_dir / "run_index.json"


def ensure_run_layout(run_dir: Path, *, conditions: Iterable[str]) -> None:
    run_dir.mkdir(parents=True, exist_ok=True)
    for condition in conditions:
        _assert_condition(condition)
        (run_dir / condition).mkdir(parents=True, exist_ok=True)


def write_json(path: Path, payload: dict[str, object]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(payload, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")


def load_json(path: Path) -> dict[str, object]:
    payload = json.loads(path.read_text(encoding="utf-8"))
    if not isinstance(payload, dict):
        raise ValueError(f"expected JSON object in {path}")
    return payload


def load_instance_ids(file_path: Path) -> list[str]:
    ids: list[str] = []
    seen: set[str] = set()
    for raw in file_path.read_text(encoding="utf-8", errors="replace").splitlines():
        value = raw.strip()
        if not value:
            continue
        if value in seen:
            continue
        seen.add(value)
        ids.append(value)
    return ids


def write_instance_ids(path: Path, instance_ids: Iterable[str]) -> None:
    cleaned: list[str] = []
    seen: set[str] = set()
    for raw in instance_ids:
        value = str(raw).strip()
        if not value or value in seen:
            continue
        seen.add(value)
        cleaned.append(value)
    path.write_text("\n".join(cleaned) + ("\n" if cleaned else ""), encoding="utf-8")


def read_predictions(path: Path) -> list[PredictionRecord]:
    out: list[PredictionRecord] = []
    if not path.exists():
        return out
    for line_no, line in enumerate(
        path.read_text(encoding="utf-8", errors="replace").splitlines(), 1
    ):
        if not line.strip():
            continue
        try:
            payload = json.loads(line)
        except json.JSONDecodeError as exc:
            raise ValueError(f"{path}:{line_no}: invalid JSON ({exc})") from exc
        if not isinstance(payload, dict):
            raise ValueError(f"{path}:{line_no}: prediction row must be object")

        instance_id = payload.get("instance_id")
        model_patch = payload.get("model_patch")
        model_name_or_path = payload.get("model_name_or_path")
        if not isinstance(instance_id, str) or not instance_id:
            raise ValueError(f"{path}:{line_no}: missing instance_id")
        if not isinstance(model_patch, str):
            raise ValueError(f"{path}:{line_no}: model_patch must be string")
        if not isinstance(model_name_or_path, str) or not model_name_or_path:
            raise ValueError(f"{path}:{line_no}: missing model_name_or_path")

        out.append(
            PredictionRecord(
                instance_id=instance_id,
                model_patch=model_patch,
                model_name_or_path=model_name_or_path,
            )
        )
    return out


def write_predictions(path: Path, records: Iterable[PredictionRecord]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temp_path: Path | None = None
    try:
        with tempfile.NamedTemporaryFile(
            mode="w",
            encoding="utf-8",
            dir=path.parent,
            prefix=f".{path.name}.",
            suffix=".tmp",
            delete=False,
        ) as handle:
            temp_path = Path(handle.name)
            for record in records:
                handle.write(json.dumps(record.as_json(), ensure_ascii=False) + "\n")
            handle.flush()
            os.fsync(handle.fileno())
        os.replace(temp_path, path)
    finally:
        if temp_path is not None and temp_path.exists():
            temp_path.unlink(missing_ok=True)


def collect_prediction_ids(path: Path) -> list[str]:
    records = read_predictions(path)
    return [record.instance_id for record in records]


def validate_prediction_file(path: Path) -> tuple[int, int]:
    records = read_predictions(path)
    seen: set[str] = set()
    duplicates = 0
    for record in records:
        if record.instance_id in seen:
            duplicates += 1
        seen.add(record.instance_id)
    return len(records), duplicates


def validate_paired_predictions(
    run_dir: Path,
    *,
    conditions: Iterable[str],
    expected_ids: Iterable[str] | None = None,
    require_files: bool = True,
    require_nonempty: bool = False,
) -> None:
    condition_list = list(parse_conditions(conditions))
    id_sets: dict[str, set[str]] = {}
    for condition in condition_list:
        condition_path = prediction_path(run_dir, condition)
        if require_files and not condition_path.exists():
            raise ValueError(f"{condition}: missing predictions file ({condition_path})")
        records = read_predictions(condition_path)
        if require_nonempty and not records:
            raise ValueError(f"{condition}: predictions file is empty ({condition_path})")
        ids = [record.instance_id for record in records]
        id_sets[condition] = set(ids)
        if len(ids) != len(id_sets[condition]):
            raise ValueError(f"{condition}: prediction file contains duplicate instance_id rows")

    first = condition_list[0]
    first_ids = id_sets[first]
    for condition in condition_list[1:]:
        if id_sets[condition] != first_ids:
            missing = sorted(first_ids - id_sets[condition])
            extra = sorted(id_sets[condition] - first_ids)
            raise ValueError(
                f"paired prediction mismatch for {condition}: "
                f"missing={missing[:5]} extra={extra[:5]}"
            )

    if expected_ids is not None:
        expected = {item.strip() for item in expected_ids if item.strip()}
        if first_ids != expected:
            missing = sorted(expected - first_ids)
            extra = sorted(first_ids - expected)
            raise ValueError(
                "prediction instance set differs from instance_ids.txt: "
                f"missing={missing[:5]} extra={extra[:5]}"
            )


def copy_instructions_text(run_dir: Path, *, conditions: Iterable[str]) -> str:
    condition_list = list(parse_conditions(conditions))
    lines: list[str] = [
        "SWE-Bench Lite Handoff",
        "======================",
        "",
        f"Run dir: {run_dir}",
        "",
        "Required files:",
        f"- {instance_ids_path(run_dir)}",
        f"- {metadata_path(run_dir)}",
    ]
    lines.extend(f"- {prediction_path(run_dir, condition)}" for condition in condition_list)

    lines.extend(
        [
            "",
            "On the Docker scoring machine:",
            "1. Copy this entire run directory.",
            "2. Run:",
            "   ./scripts/score_swebench_lite_official.sh --run-dir <copied_run_dir>",
            "",
            "Example:",
            "   ./scripts/score_swebench_lite_official.sh \\",
            "     --run-dir /path/to/evals/swebench_runs/<run_id>",
            "",
            "Outputs will be written under:",
            f"- {run_dir / 'official_eval'}",
            "",
            f"Generated at: {datetime.now(timezone.utc).isoformat()}",
        ]
    )
    return "\n".join(lines) + "\n"


def write_copy_instructions(run_dir: Path, *, conditions: Iterable[str]) -> Path:
    text = copy_instructions_text(run_dir, conditions=conditions)
    path = copy_instructions_path(run_dir)
    path.write_text(text, encoding="utf-8")
    return path


def write_run_index(
    run_dir: Path,
    *,
    run_id: str,
    conditions: Iterable[str],
) -> Path:
    condition_list = list(parse_conditions(conditions))
    payload: dict[str, object] = {
        "schema_version": "swebench_handoff_index_v1",
        "generated_at_utc": datetime.now(timezone.utc).replace(microsecond=0).isoformat(),
        "run_id": run_id,
        "run_dir": str(run_dir),
        "instance_ids_file": str(instance_ids_path(run_dir)),
        "metadata_file": str(metadata_path(run_dir)),
        "predictions": {
            condition: str(prediction_path(run_dir, condition)) for condition in condition_list
        },
        "official_eval_root": str(run_dir / "official_eval"),
        "required_files": [
            str(instance_ids_path(run_dir)),
            str(metadata_path(run_dir)),
            *(str(prediction_path(run_dir, condition)) for condition in condition_list),
        ],
    }
    path = run_index_path(run_dir)
    write_json(path, payload)
    return path


def _assert_condition(condition: str) -> None:
    if condition not in VALID_CONDITIONS:
        raise ValueError(
            f"unsupported condition {condition!r}; expected one of {', '.join(VALID_CONDITIONS)}"
        )
