from __future__ import annotations

import hashlib
import json
from collections.abc import Iterable
from pathlib import Path
from typing import TypeVar

from pydantic import BaseModel

from .schemas import RewardRecord, TaskRecord

T = TypeVar("T", bound=BaseModel)


def read_jsonl(path: str | Path, model: type[T]) -> list[T]:
    records: list[T] = []
    with Path(path).open("r", encoding="utf-8") as handle:
        for line_number, line in enumerate(handle, start=1):
            if not line.strip():
                continue
            try:
                records.append(model.model_validate_json(line))
            except Exception as exc:
                raise ValueError(f"Invalid JSONL at {path}:{line_number}: {exc}") from exc
    return records


def write_jsonl(path: str | Path, records: Iterable[BaseModel]) -> None:
    destination = Path(path)
    destination.parent.mkdir(parents=True, exist_ok=True)
    temporary = destination.with_suffix(destination.suffix + ".tmp")
    with temporary.open("w", encoding="utf-8") as handle:
        for record in records:
            handle.write(record.model_dump_json() + "\n")
    temporary.replace(destination)


def load_tasks(path: str | Path) -> list[TaskRecord]:
    return read_jsonl(path, TaskRecord)


def load_rewards(path: str | Path) -> list[RewardRecord]:
    records = read_jsonl(path, RewardRecord)
    if not records:
        raise ValueError(f"No reward records found in {path}")
    expected = records[0].worker_ids
    for record in records[1:]:
        if record.worker_ids != expected:
            raise ValueError(
                f"Worker order changed at task {record.task_id}; expected {expected}, "
                f"got {record.worker_ids}"
            )
    return records


def assigned_split(task_id: str, explicit: str | None) -> str:
    if explicit:
        return explicit
    bucket = int(hashlib.sha256(task_id.encode("utf-8")).hexdigest()[:8], 16) % 10
    if bucket == 0:
        return "test"
    if bucket == 1:
        return "validation"
    return "train"


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