src/session_history.py

import time
import uuid
import json
import os


_HISTORY_FILE = os.path.join(os.path.dirname(os.path.dirname(__file__)), "data", "session_history.json")


class HistoryEntry:
    __slots__ = ("id", "timestamp", "prompt", "model", "rating", "num_variations",
                 "creativity", "weight_mode", "results", "neg_results",
                 "categories", "seed", "liked_indices")

    def __init__(self, prompt: str, results: list[str], model: str = "anima",
                 rating: str = "pg", num_variations: int = 5,
                 creativity: str = "medium", weight_mode: str = "off",
                 neg_results: list[str] | None = None,
                 categories: list[str] | None = None,
                 seed: int | None = None):
        self.id = uuid.uuid4().hex[:8]
        self.timestamp = time.time()
        self.prompt = prompt
        self.model = model
        self.rating = rating
        self.num_variations = num_variations
        self.creativity = creativity
        self.weight_mode = weight_mode
        self.results = list(results)
        self.neg_results = list(neg_results) if neg_results else []
        self.categories = list(categories) if categories else []
        self.seed = seed
        self.liked_indices: list[int] = []

    @classmethod
    def from_dict(cls, d: dict) -> "HistoryEntry":
        entry = cls.__new__(cls)
        entry.id = d.get("id", uuid.uuid4().hex[:8])
        entry.timestamp = d.get("timestamp", time.time())
        entry.prompt = d.get("prompt", "")
        entry.model = d.get("model", "anima")
        entry.rating = d.get("rating", "pg")
        entry.num_variations = d.get("num_variations", 5)
        entry.creativity = d.get("creativity", "medium")
        entry.weight_mode = d.get("weight_mode", "off")
        entry.results = list(d.get("results", []))
        entry.neg_results = list(d.get("neg_results", []))
        entry.categories = list(d.get("categories", []))
        entry.seed = d.get("seed")
        entry.liked_indices = list(d.get("liked_indices", []))
        return entry

    def toggle_like(self, index: int) -> bool:
        if index in self.liked_indices:
            self.liked_indices.remove(index)
            return False
        self.liked_indices.append(index)
        return True

    def is_liked(self, index: int) -> bool:
        return index in self.liked_indices

    def get_liked_results(self) -> list[tuple[int, str]]:
        return [(i, self.results[i]) for i in self.liked_indices if i < len(self.results)]

    def to_dict(self) -> dict:
        return {
            "id": self.id,
            "timestamp": self.timestamp,
            "prompt": self.prompt,
            "model": self.model,
            "rating": self.rating,
            "num_variations": self.num_variations,
            "creativity": self.creativity,
            "weight_mode": self.weight_mode,
            "results": self.results,
            "neg_results": self.neg_results,
            "categories": self.categories,
            "seed": self.seed,
            "liked_indices": self.liked_indices,
        }


class SessionHistory:
    def __init__(self, max_entries: int = 50, file_path: str | None = None):
        self._entries: list[HistoryEntry] = []
        self._max_entries = max_entries
        self._file_path = file_path or _HISTORY_FILE
        self._load()

    def persist(self):
        self._save()

    def toggle_like(self, entry_index: int, result_index: int) -> bool:
        if entry_index < 0 or entry_index >= len(self._entries):
            return False
        liked = self._entries[entry_index].toggle_like(result_index)
        self._save()
        return liked

    def add(self, entry: HistoryEntry) -> str:
        self._entries.insert(0, entry)
        if len(self._entries) > self._max_entries:
            self._entries.pop()
        self._save()
        return entry.id

    def get(self, entry_id: str) -> HistoryEntry | None:
        for e in self._entries:
            if e.id == entry_id:
                return e
        return None

    def get_all(self) -> list[HistoryEntry]:
        return list(self._entries)

    def get_favorites(self) -> list[tuple[HistoryEntry, int, str]]:
        items = []
        for entry in self._entries:
            for idx, text in entry.get_liked_results():
                items.append((entry, idx, text))
        return items

    def clear(self):
        self._entries.clear()
        self._save()

    def _save(self):
        data = [e.to_dict() for e in self._entries]
        tmp_path = self._file_path + ".tmp"
        try:
            with open(tmp_path, "w", encoding="utf-8") as f:
                json.dump(data, f, ensure_ascii=False, indent=2)
            if os.path.exists(self._file_path):
                os.replace(self._file_path, self._file_path + ".bak")
            os.replace(tmp_path, self._file_path)
        except (OSError, TypeError):
            try:
                if os.path.exists(tmp_path):
                    os.remove(tmp_path)
            except OSError:
                pass

    def _load(self):
        # Try primary, then .bak (recovers from an interrupted write), then reset.
        for path in (self._file_path, self._file_path + ".bak"):
            if not os.path.exists(path):
                continue
            try:
                with open(path, "r", encoding="utf-8") as f:
                    data = json.load(f)
                self._entries = [HistoryEntry.from_dict(d) for d in data[:self._max_entries]]
                return
            except (json.JSONDecodeError, UnicodeDecodeError, ValueError, OSError):
                continue
        self._entries = []

    def __len__(self) -> int:
        return len(self._entries)


_history = SessionHistory()


def get_history() -> SessionHistory:
    return _history


def add_to_history(prompt: str, results: list[str], **kwargs) -> str:
    entry = HistoryEntry(prompt, results, **kwargs)
    return _history.add(entry)