src/pino/feedback.py

from __future__ import annotations

import json
import logging
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

from huggingface_hub import hf_hub_download, upload_file

logger = logging.getLogger("pino.feedback")

DEFAULT_FEEDBACK_PATH = Path("data/incremental_feedback.jsonl")
DEFAULT_DATASET_REPO = "mattbitzesty/pino-synthetic-dataset"


@dataclass
class Critique:
    """Structured human or agent critique of a fragrance prediction."""

    formula_id: str
    rating: int | None = None  # 1-5 overall rating
    notes: str = ""
    axis_scores: dict[str, float] = field(default_factory=dict)
    correct_gender: float | None = None  # -1.0 to +1.0
    correct_notes: dict[str, float] = field(default_factory=dict)  # index -> value overrides

    def to_record(self) -> dict[str, Any]:
        return {
            "formula_id": self.formula_id,
            "rating": self.rating,
            "notes": self.notes,
            "axis_scores": self.axis_scores,
            "correct_gender": self.correct_gender,
            "correct_notes": self.correct_notes,
        }

    @classmethod
    def from_record(cls, record: dict[str, Any]) -> "Critique":
        return cls(
            formula_id=record["formula_id"],
            rating=record.get("rating"),
            notes=record.get("notes", ""),
            axis_scores=record.get("axis_scores", {}),
            correct_gender=record.get("correct_gender"),
            correct_notes=record.get("correct_notes", {}),
        )


class FeedbackLoopManager:
    """
    Atomic append-only recorder for real-world corrections.

    Feedback critiques live in a dedicated ``data/incremental_feedback.jsonl``
    file and are NEVER appended to the formulation dataset JSONL. When the
    flywheel triggers, only valid formulation records from the original pool are
    pushed to the Hugging Face Dataset Hub so a cloud training job can consume a
    clean corpus.
    """

    def __init__(
        self,
        feedback_path: str | Path = DEFAULT_FEEDBACK_PATH,
        dataset_repo: str = DEFAULT_DATASET_REPO,
    ) -> None:
        self.feedback_path = Path(feedback_path)
        self.dataset_repo = dataset_repo
        self._ensure_feedback_file()

    def _ensure_feedback_file(self) -> None:
        self.feedback_path.parent.mkdir(parents=True, exist_ok=True)
        if not self.feedback_path.exists():
            self.feedback_path.write_text("", encoding="utf-8")

    def append(self, critique: Critique) -> int:
        """Append a critique atomically and return the new total feedback count."""
        with self.feedback_path.open("a", encoding="utf-8") as f:
            f.write(json.dumps(critique.to_record(), ensure_ascii=False) + "\n")

        count = self.count()
        logger.info("Logged critique for %s; total feedback rows: %d", critique.formula_id, count)
        return count

    def count(self) -> int:
        """Return the number of feedback rows currently stored."""
        if not self.feedback_path.exists():
            return 0
        with self.feedback_path.open("r", encoding="utf-8") as f:
            return sum(1 for line in f if line.strip())

    def read_all(self) -> list[Critique]:
        """Return all critiques from the feedback file."""
        critiques = []
        if not self.feedback_path.exists():
            return critiques
        with self.feedback_path.open("r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if line:
                    critiques.append(Critique.from_record(json.loads(line)))
        return critiques

    @staticmethod
    def _is_valid_formulation_record(record: dict[str, Any]) -> bool:
        """
        Return True only for records that conform to the formulation schema.

        A valid formulation record carries both the ``formula`` payload and the
        ``trajectory`` array produced by the dataset builder, which is exactly
        what :class:`~pino.pimt_model.FragranceTrajectoryDataset` requires on
        every line. Critique/feedback rows (which carry ``feedback`` instead)
        are explicitly excluded so they never leak into the formulation dataset
        JSONL.
        """
        return bool(
            isinstance(record, dict)
            and record.get("formula")
            and record.get("trajectory")
            and not record.get("feedback")
        )

    def _load_original_pool(
        self,
        original_path: str | Path | None = None,
    ) -> list[dict[str, Any]]:
        """Load and return only valid formulation records from the original pool."""
        if original_path is None:
            original_path = hf_hub_download(
                repo_id=self.dataset_repo,
                filename="synthetic_dataset_v2.jsonl",
                repo_type="dataset",
            )

        original_path = Path(original_path)
        records: list[dict[str, Any]] = []
        with original_path.open("r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                record = json.loads(line)
                if self._is_valid_formulation_record(record):
                    records.append(record)

        logger.debug(
            "Loaded %d valid formulation records from %s", len(records), original_path
        )
        return records

    def merge_with_original_pool(
        self,
        original_path: str | Path | None = None,
    ) -> list[dict[str, Any]]:
        """
        Return only valid formulation records from the original pool.

        Feedback critiques are intentionally NOT appended here. They live in the
        separate ``incremental_feedback.jsonl`` file so the formulation dataset
        schema stays clean. Use ``read_all()`` to access raw feedback rows.
        """
        return self._load_original_pool(original_path)

    def push_merged_pool(
        self,
        original_path: str | Path | None = None,
        path_in_repo: str = "synthetic_dataset_v2.jsonl",
    ) -> str:
        """Write valid formulation records to JSONL and push to the Hugging Face Dataset Hub."""
        merged = self.merge_with_original_pool(original_path)
        local_path = Path("/tmp/pino_merged_dataset.jsonl")
        with local_path.open("w", encoding="utf-8") as f:
            for r in merged:
                f.write(json.dumps(r, ensure_ascii=False) + "\n")

        logger.info("Uploading merged dataset (%d rows) to %s", len(merged), self.dataset_repo)
        upload_file(
            path_or_fileobj=str(local_path),
            path_in_repo=path_in_repo,
            repo_id=self.dataset_repo,
            repo_type="dataset",
        )
        return f"{self.dataset_repo}/{path_in_repo}"

    def sync_flywheel(
        self,
        original_path: str | Path | None = None,
        job_config_path: str | Path = "hf_job.yaml",
    ) -> dict[str, Any]:
        """
        If new feedback exists, push the clean formulation pool to the Hub and
        trigger a fresh training job. Feedback rows are preserved in the separate
        feedback file and never mixed into the formulation dataset.
        """
        new_count = self.count()
        if new_count == 0:
            return {"action": "none", "reason": "no feedback samples", "merged_samples": 0}

        hub_path = self.push_merged_pool(original_path)
        self._trigger_training_job(job_config_path)
        return {
            "action": "sync",
            "feedback_samples": new_count,
            "dataset_hub_path": hub_path,
            "training_job_triggered": True,
        }

    def _trigger_training_job(self, job_config_path: str | Path) -> None:
        """Trigger a fresh HF training container via the configured job YAML."""
        # The canonical invocation is `hf jobs create hf_job.yaml` from the terminal.
        # We log it here and leave the actual subprocess call to the CLI layer.
        logger.info("Training retrigger requested via %s", job_config_path)

    def clear(self) -> None:
        """Clear the feedback file after a successful sync."""
        if self.feedback_path.exists():
            self.feedback_path.write_text("", encoding="utf-8")
            logger.info("Cleared feedback file: %s", self.feedback_path)