core-python/maris_core/training/human_training.py

"""Human-in-the-loop Space helpers for staging training artifacts."""

from __future__ import annotations

import json
import os
from datetime import UTC, datetime
from pathlib import Path
from typing import Any
from uuid import uuid4

from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator

from maris_core.data.quality import (
    DatasetQualityGateConfig,
    apply_quality_gate_to_records,
    build_dataset_quality_report,
)
from maris_core.training.preferences import (
    PreferenceExample,
    build_blind_side_by_side_artifact,
    build_human_eval_summary,
    summarize_preference_dataset,
)
from maris_core.training.space_ui import SAFE_OUTPUT_SEGMENT_RE
from maris_core.utils.env import validate_hf_model, validate_hf_repo_id

HUMAN_TRAINING_STAGE_DIRNAME = "human-training-staging"
HUMAN_TRAINING_REPO_PREFIX = "human-training"
HUMAN_TRAINING_MIN_TEXT_CHARS = 24
HUMAN_TRAINING_BLIND_REVIEW_SEED = 7


def _timestamp() -> str:
    return datetime.now(UTC).replace(microsecond=0).isoformat()


def _new_run_id() -> str:
    timestamp = datetime.now(UTC).strftime("%Y%m%dT%H%M%SZ")
    return f"{timestamp}-{uuid4().hex[:8]}"


def _validate_output_subdir(value: str) -> str:
    normalized = value.strip().strip("/")
    if not normalized:
        raise ValueError("Output apakšdirektorija nedrīkst būt tukša.")
    parts = Path(normalized).parts
    if ".." in parts or not SAFE_OUTPUT_SEGMENT_RE.fullmatch(normalized):
        raise ValueError("Output apakšdirektorijā drīkst būt tikai droši ceļa segmenti.")
    return normalized


def _normalize_lines(values: list[str]) -> list[str]:
    normalized: list[str] = []
    seen: set[str] = set()
    for value in values:
        text = str(value or "").strip()
        if not text:
            continue
        signature = text.casefold()
        if signature in seen:
            continue
        seen.add(signature)
        normalized.append(text)
    return normalized


class HumanTrainingConversationExample(BaseModel):
    model_config = ConfigDict(str_strip_whitespace=True)

    user: str = Field(min_length=1)
    assistant: str = Field(min_length=1)
    context: str = ""


class HumanTrainingPreferencePair(BaseModel):
    model_config = ConfigDict(str_strip_whitespace=True)

    prompt: str = Field(min_length=1)
    chosen: str = Field(min_length=1)
    rejected: str = Field(min_length=1)
    context: str = ""
    confidence: float | None = Field(default=None, ge=0.0, le=1.0)


class HumanTrainingEvalExample(BaseModel):
    model_config = ConfigDict(str_strip_whitespace=True)

    prompt: str = Field(min_length=1)
    completion: str = Field(min_length=1)
    context: str = ""


class HumanTrainingRequest(BaseModel):
    """Request payload for building staged human-training artifacts."""

    model_config = ConfigDict(str_strip_whitespace=True)

    dataset_repo: str
    model_repo: str = ""
    hub_model_id: str = ""
    model_preset: str = "balanced"
    model_name: str = ""
    num_epochs: int = Field(default=3, ge=1, le=100)
    all_branches: bool = False
    push_to_hub: bool = True
    output_subdir: str = "maris-human-training"
    continue_from_latest_artifact: bool = True
    continue_model_path: str = ""
    profile_facts: list[str] = Field(default_factory=list)
    profile_preferences: list[str] = Field(default_factory=list)
    response_instructions: list[str] = Field(default_factory=list)
    conversation_examples: list[HumanTrainingConversationExample] = Field(default_factory=list)
    preference_pairs: list[HumanTrainingPreferencePair] = Field(default_factory=list)
    eval_examples: list[HumanTrainingEvalExample] = Field(default_factory=list)

    @field_validator("dataset_repo")
    @classmethod
    def validate_dataset_repo(cls, value: str) -> str:
        try:
            return validate_hf_repo_id(value, "dataset_repo", label="dataset repo")
        except RuntimeError as exc:
            raise ValueError(str(exc)) from exc

    @field_validator("model_repo")
    @classmethod
    def validate_model_repo(cls, value: str) -> str:
        normalized = value.strip()
        if not normalized:
            return ""
        try:
            return validate_hf_model(normalized, "model_repo")
        except RuntimeError as exc:
            raise ValueError(str(exc)) from exc

    @field_validator("hub_model_id")
    @classmethod
    def validate_hub_model_id(cls, value: str) -> str:
        normalized = value.strip()
        if not normalized:
            return ""
        try:
            return validate_hf_model(normalized, "hub_model_id")
        except RuntimeError as exc:
            raise ValueError(str(exc)) from exc

    @field_validator("model_name")
    @classmethod
    def validate_model_name(cls, value: str) -> str:
        normalized = value.strip()
        if not normalized:
            return ""
        try:
            return validate_hf_model(normalized, "model_name")
        except RuntimeError as exc:
            raise ValueError(str(exc)) from exc

    @field_validator("output_subdir")
    @classmethod
    def validate_output_subdir(cls, value: str) -> str:
        return _validate_output_subdir(value)

    @field_validator("continue_model_path")
    @classmethod
    def validate_continue_model_path(cls, value: str) -> str:
        normalized = value.strip()
        if not normalized:
            return ""
        return _validate_output_subdir(normalized)

    @field_validator("profile_facts", "profile_preferences", "response_instructions")
    @classmethod
    def dedupe_lines(cls, value: list[str]) -> list[str]:
        return _normalize_lines(value)

    @model_validator(mode="after")
    def validate_has_training_signal(self) -> HumanTrainingRequest:
        resolved_model_repo = self.hub_model_id or self.model_repo
        if not resolved_model_repo:
            raise ValueError("Jānorāda hub_model_id vai model_repo.")
        self.hub_model_id = resolved_model_repo
        self.model_repo = resolved_model_repo
        if not any(
            (
                self.profile_facts,
                self.profile_preferences,
                self.response_instructions,
                self.conversation_examples,
                self.preference_pairs,
                self.eval_examples,
            )
        ):
            raise ValueError("Human training pieprasījumā jābūt vismaz vienam ievades blokam.")
        if not self.model_name and not self.model_preset:
            self.model_preset = "balanced"
        return self


class HumanTrainingExecuteRequest(BaseModel):
    model_config = ConfigDict(str_strip_whitespace=True)

    run_id: str = Field(min_length=6)
    publish_artifacts: bool = True
    start_training: bool = False

    @model_validator(mode="after")
    def validate_execution_flags(self) -> HumanTrainingExecuteRequest:
        if self.start_training and not self.publish_artifacts:
            raise ValueError(
                "Lai sāktu treniņu, artefakti vispirms jāpublicē dataset repozitorijā."
            )
        return self


class HumanTrainingLaunchSpec(BaseModel):
    model_config = ConfigDict(str_strip_whitespace=True)

    dataset_repo: str
    model_repo: str
    hub_model_id: str = ""
    model_preset: str = ""
    model_name: str = ""
    num_epochs: int
    all_branches: bool
    push_to_hub: bool
    output_subdir: str
    continue_from_latest_artifact: bool = True
    continue_model_path: str = ""

    @model_validator(mode="after")
    def validate_model_target(self) -> HumanTrainingLaunchSpec:
        resolved_model_repo = self.hub_model_id or self.model_repo
        if not resolved_model_repo:
            raise ValueError("Jānorāda hub_model_id vai model_repo.")
        self.hub_model_id = resolved_model_repo
        self.model_repo = resolved_model_repo
        return self


def resolve_human_training_stage_dir(persistent_dir: str, run_id: str) -> Path:
    root = Path(persistent_dir).expanduser().resolve()
    target = (root / HUMAN_TRAINING_STAGE_DIRNAME / run_id).resolve()
    if os.path.commonpath([str(root), str(target)]) != str(root):
        raise ValueError(
            "Human training staging direktorijai jāatrodas persistent storage ietvaros."
        )
    return target


def stage_human_training_artifacts(
    request: HumanTrainingRequest,
    *,
    persistent_dir: str,
) -> dict[str, Any]:
    run_id = _new_run_id()
    stage_dir = resolve_human_training_stage_dir(persistent_dir, run_id)
    stage_dir.mkdir(parents=True, exist_ok=True)

    train_records = _build_train_records(request)
    eval_records = _build_eval_records(request)
    preference_dataset = _build_preference_dataset(request)

    quality_config = DatasetQualityGateConfig(min_text_chars=HUMAN_TRAINING_MIN_TEXT_CHARS)
    filtered_train, train_report = apply_quality_gate_to_records(
        train_records,
        split_name="train",
        config=quality_config,
    )
    filtered_eval, eval_report = (
        apply_quality_gate_to_records(
            eval_records,
            split_name="eval",
            config=quality_config,
        )
        if eval_records
        else ([], None)
    )
    quality_report = build_dataset_quality_report(
        config=quality_config,
        train_report=train_report,
        eval_report=eval_report,
    ).to_dict()

    preference_examples = [PreferenceExample(**item) for item in preference_dataset["preferences"]]
    preference_summary = (
        summarize_preference_dataset(preference_examples) if preference_examples else None
    )
    blind_review = (
        build_blind_side_by_side_artifact(
            preference_examples,
            seed=HUMAN_TRAINING_BLIND_REVIEW_SEED,
        )
        if preference_examples
        else None
    )
    human_eval_summary = (
        build_human_eval_summary(preference_examples) if preference_examples else None
    )

    artifact_specs: dict[str, dict[str, Any]] = {
        "train_dataset": {
            "repo_path": f"data/{HUMAN_TRAINING_REPO_PREFIX}/{run_id}/train.json",
            "payload": filtered_train,
            "record_count": len(filtered_train),
        },
        "dataset_quality_report": {
            "repo_path": f"artifacts/{HUMAN_TRAINING_REPO_PREFIX}/{run_id}/dataset-quality-report.json",
            "payload": quality_report,
            "record_count": quality_report["kept_records"],
        },
    }
    if filtered_eval:
        artifact_specs["eval_dataset"] = {
            "repo_path": f"artifacts/{HUMAN_TRAINING_REPO_PREFIX}/{run_id}/eval.json",
            "payload": filtered_eval,
            "record_count": len(filtered_eval),
        }
    if preference_dataset["preferences"]:
        artifact_specs["preference_dataset"] = {
            "repo_path": f"artifacts/{HUMAN_TRAINING_REPO_PREFIX}/{run_id}/preferences.json",
            "payload": preference_dataset,
            "record_count": len(preference_dataset["preferences"]),
        }
    if preference_summary is not None:
        artifact_specs["preference_summary"] = {
            "repo_path": f"artifacts/{HUMAN_TRAINING_REPO_PREFIX}/{run_id}/preference-summary.json",
            "payload": preference_summary,
            "record_count": preference_summary["total_examples"],
        }
    if blind_review is not None:
        artifact_specs["blind_review"] = {
            "repo_path": f"artifacts/{HUMAN_TRAINING_REPO_PREFIX}/{run_id}/blind-review.json",
            "payload": blind_review,
            "record_count": blind_review["total_pairs"],
        }
    if human_eval_summary is not None:
        artifact_specs["human_eval_summary"] = {
            "repo_path": f"artifacts/{HUMAN_TRAINING_REPO_PREFIX}/{run_id}/human-eval-summary.json",
            "payload": human_eval_summary,
            "record_count": human_eval_summary["total_examples"],
        }

    artifacts: dict[str, dict[str, Any]] = {}
    for name, spec in artifact_specs.items():
        local_path = stage_dir / Path(spec["repo_path"]).name
        _write_json(local_path, spec["payload"])
        artifacts[name] = {
            "local_path": str(local_path),
            "repo_path": spec["repo_path"],
            "record_count": spec["record_count"],
        }

    launch_spec = HumanTrainingLaunchSpec(
        dataset_repo=request.dataset_repo,
        model_repo=request.model_repo,
        hub_model_id=request.hub_model_id,
        model_preset="" if request.model_name else request.model_preset,
        model_name=request.model_name,
        num_epochs=request.num_epochs,
        all_branches=request.all_branches,
        push_to_hub=request.push_to_hub,
        output_subdir=request.output_subdir,
        continue_from_latest_artifact=request.continue_from_latest_artifact,
        continue_model_path=request.continue_model_path,
    )
    manifest = {
        "artifact_type": "human-training-manifest",
        "run_id": run_id,
        "staged_at": _timestamp(),
        "stage_dir": str(stage_dir),
        "dataset_repo": request.dataset_repo,
        "model_repo": request.model_repo,
        "input_summary": {
            "profile_facts": len(request.profile_facts),
            "profile_preferences": len(request.profile_preferences),
            "response_instructions": len(request.response_instructions),
            "conversation_examples": len(request.conversation_examples),
            "preference_pairs": len(request.preference_pairs),
            "eval_examples": len(request.eval_examples),
        },
        "quality_report": quality_report,
        "artifacts": artifacts,
        "training_request": launch_spec.model_dump(),
        "ready_for_review": True,
        "ready_for_training": bool(filtered_train),
    }
    _write_json(stage_dir / "manifest.json", manifest)
    return manifest


def load_human_training_manifest(persistent_dir: str, run_id: str) -> dict[str, Any]:
    manifest_path = resolve_human_training_stage_dir(persistent_dir, run_id) / "manifest.json"
    if not manifest_path.is_file():
        raise FileNotFoundError(f"Human training staging ieraksts {run_id} nav atrasts.")
    return json.loads(manifest_path.read_text(encoding="utf-8"))


def publish_human_training_artifacts(
    manifest: dict[str, Any],
    *,
    save_file: Any,
) -> list[dict[str, Any]]:
    published: list[dict[str, Any]] = []
    dataset_repo = str(manifest["dataset_repo"])
    run_id = str(manifest["run_id"])
    for artifact_name, artifact in manifest.get("artifacts", {}).items():
        repo_path = artifact.get("repo_path")
        local_path = artifact.get("local_path")
        if not isinstance(repo_path, str) or not isinstance(local_path, str):
            continue
        content = Path(local_path).read_text(encoding="utf-8")
        result = save_file(
            repo_id=dataset_repo,
            repo_type="dataset",
            path_in_repo=repo_path,
            content=content,
            commit_message=f"Add human training artifacts for {run_id}",
        )
        published.append({"artifact": artifact_name, **result})
    return published


def build_human_training_launch_spec(manifest: dict[str, Any]) -> HumanTrainingLaunchSpec:
    return HumanTrainingLaunchSpec.model_validate(manifest["training_request"])


def _build_train_records(request: HumanTrainingRequest) -> list[dict[str, Any]]:
    records: list[dict[str, Any]] = []
    profile_record = _build_profile_record(request)
    if profile_record is not None:
        records.append(profile_record)
    for example in request.conversation_examples:
        record = {
            "user": example.user,
            "assistant": example.assistant,
            "metadata": {"source": "human_training_space"},
        }
        if example.context:
            record["context"] = example.context
        records.append(record)
    return records


def _build_profile_record(request: HumanTrainingRequest) -> dict[str, Any] | None:
    sections: list[str] = []
    if request.profile_facts:
        sections.append(
            "Fakti par lietotāju:\n" + "\n".join(f"- {item}" for item in request.profile_facts)
        )
    if request.profile_preferences:
        sections.append(
            "Lietotāja preferences:\n"
            + "\n".join(f"- {item}" for item in request.profile_preferences)
        )
    if request.response_instructions:
        sections.append(
            "Atbildēšanas instrukcijas:\n"
            + "\n".join(f"- {item}" for item in request.response_instructions)
        )
    if not sections:
        return None
    return {
        "prompt": "Iegaumē šo lietotāja profilu un atbildē saskaņā ar to turpmākajās sarunās.",
        "completion": "\n\n".join(sections),
        "metadata": {
            "source": "human_training_space",
            "artifact": "profile_memory",
        },
    }


def _build_eval_records(request: HumanTrainingRequest) -> list[dict[str, Any]]:
    records: list[dict[str, Any]] = []
    for example in request.eval_examples:
        record = {
            "prompt": example.prompt,
            "completion": example.completion,
            "metadata": {"source": "human_training_space", "artifact": "eval"},
        }
        if example.context:
            record["context"] = example.context
        records.append(record)
    return records


def _build_preference_dataset(request: HumanTrainingRequest) -> dict[str, Any]:
    preferences: list[dict[str, Any]] = []
    for index, pair in enumerate(request.preference_pairs, start=1):
        item = {
            "prompt": pair.prompt,
            "chosen": pair.chosen,
            "rejected": pair.rejected,
            "context": pair.context,
            "source": "human_training_space",
            "source_type": "real_reviewer",
            "reviewer_segment": "self-training",
            "preference_outcome": "chosen",
            "confidence": pair.confidence,
            "pair_id": f"human-training-{index:04d}",
            "blind": True,
            "production_like": True,
            "tags": ["human-training"],
        }
        preferences.append(item)
    return {
        "artifact_type": "human-training-preferences",
        "preferences": preferences,
    }


def _write_json(path: Path, payload: Any) -> None:
    path.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")