backend/db.py

"""Backend db service."""
from __future__ import annotations

import sqlite3
from collections import defaultdict
from datetime import UTC, datetime
from backend.constants import VALID_RATINGS
from sqlmodel import Session, select
from models import PlayPreference, Kink



def _sqlite(self) -> sqlite3.Connection:
    conn = sqlite3.connect(self.path, timeout=30)
    conn.row_factory = sqlite3.Row
    conn.execute("PRAGMA journal_mode=WAL")
    conn.execute("PRAGMA synchronous=NORMAL")
    conn.execute("PRAGMA busy_timeout=30000")
    conn.execute("PRAGMA foreign_keys=ON")
    return conn

def _now_iso(self) -> str:
    return datetime.now(UTC).isoformat()

def _ensure_fts_schema(self) -> None:
    with self._sqlite() as conn:
        row = conn.execute(
            "SELECT sql FROM sqlite_master WHERE type = 'table' AND name = 'play_fts'"
        ).fetchone()
        if row and row["sql"] and (
            "porter" not in row["sql"].lower() or "content_kind unindexed" not in row["sql"].lower()
        ):
            conn.execute("DROP TABLE IF EXISTS play_fts")
        conn.execute(
            """
            CREATE VIRTUAL TABLE IF NOT EXISTS play_fts USING fts5(
                kink_id UNINDEXED,
                content_kind UNINDEXED,
                name,
                aliases,
                definition,
                examples,
                tokenize = 'porter unicode61 remove_diacritics 2'
            )
            """
        )
        conn.commit()

def _ensure_runtime_indexes(self) -> None:
    with self._sqlite() as conn:
        conn.execute(
            "CREATE INDEX IF NOT EXISTS idx_fetlifecrawljob_state_updated ON fetlifecrawljob(state, updated_at)"
        )
        conn.execute(
            "CREATE INDEX IF NOT EXISTS idx_fetlifeuserfetish_bucket_nickname_kink "
            "ON fetlifeuserfetish(bucket, nickname, kink_id)"
        )
        conn.execute(
            "CREATE INDEX IF NOT EXISTS idx_fetlifeuserfetish_nickname_kink "
            "ON fetlifeuserfetish(nickname, kink_id)"
        )
        conn.execute(
            "CREATE INDEX IF NOT EXISTS idx_similarityedge_type_left "
            "ON similarityedge(similarity_type, left_kink_id)"
        )
        conn.execute(
            "CREATE INDEX IF NOT EXISTS idx_similarityedge_type_right "
            "ON similarityedge(similarity_type, right_kink_id)"
        )
        conn.execute(
            "CREATE INDEX IF NOT EXISTS idx_similarityedge_left_score_right "
            "ON similarityedge(left_kink_id, score DESC, right_kink_id)"
        )
        conn.execute(
            "CREATE INDEX IF NOT EXISTS idx_fetlifesupervisorrun_state_heartbeat "
            "ON fetlifesupervisorrun(state, heartbeat_at)"
        )
        conn.execute(
            "CREATE INDEX IF NOT EXISTS idx_kinkscenarioparent_parent_score "
            "ON kinkscenarioparent(parent_kink_id, score DESC)"
        )
        conn.execute(
            "CREATE INDEX IF NOT EXISTS idx_kinkscenarioparent_scenario_score "
            "ON kinkscenarioparent(scenario_kink_id, score DESC)"
        )
        conn.execute(
            "CREATE INDEX IF NOT EXISTS idx_scenariopreference_user_parent "
            "ON scenariopreference(user_id, parent_kink_id)"
        )
        conn.commit()

def _normalize_interest_state(self, interest_state: str) -> str:
    if interest_state == "no":
        return "hard_no"
    return interest_state

def _migrate_user_table(self) -> None:
    with self._sqlite() as conn:
        row = conn.execute(
            "SELECT sql FROM sqlite_master WHERE type='table' AND name='user'"
        ).fetchone()
        if not row or not row["sql"] or "label" not in row["sql"].lower():
            return
        conn.execute("PRAGMA foreign_keys=OFF")
        conn.execute("BEGIN")
        conn.execute(
            """
            CREATE TABLE user_new (
                id VARCHAR NOT NULL PRIMARY KEY,
                private_token VARCHAR NOT NULL
            )
            """
        )
        conn.execute(
            "INSERT INTO user_new (id, private_token) SELECT id, private_token FROM user"
        )
        conn.execute("DROP TABLE user")
        conn.execute("ALTER TABLE user_new RENAME TO user")
        conn.execute("COMMIT")
        conn.execute("PRAGMA foreign_keys=ON")

def _migrate_partner_group_member(self) -> None:
    with self._sqlite() as conn:
        cols = conn.execute("PRAGMA table_info(partnergroupmember)").fetchall()
        col_names = {c[1] for c in cols}
        if "share_full_list" not in col_names:
            conn.execute("ALTER TABLE partnergroupmember ADD COLUMN share_full_list INTEGER DEFAULT 0")
            conn.commit()


def _migrate_duplicate_pair_groups(self) -> None:
    """Collapse old mirrored automatic Together groups for the same linked pair."""
    with self._sqlite() as conn:
        links = {
            tuple(sorted((row["left_user_id"], row["right_user_id"])))
            for row in conn.execute("SELECT left_user_id, right_user_id FROM partnerlink").fetchall()
        }
        if not links:
            return
        groups = conn.execute(
            "SELECT id, owner_user_id, name FROM partnergroup WHERE lower(name) = 'together'"
        ).fetchall()
        if not groups:
            return
        member_rows = conn.execute(
            "SELECT group_id, member_user_id, COALESCE(share_full_list, 0) AS share_full_list FROM partnergroupmember"
        ).fetchall()
        members_by_group: dict[str, list[sqlite3.Row]] = defaultdict(list)
        for row in member_rows:
            members_by_group[row["group_id"]].append(row)

        groups_by_pair: dict[tuple[str, str], list[sqlite3.Row]] = defaultdict(list)
        for group in groups:
            participants = {group["owner_user_id"]}
            participants.update(row["member_user_id"] for row in members_by_group.get(group["id"], []))
            if len(participants) != 2:
                continue
            pair = tuple(sorted(participants))
            if pair not in links:
                continue
            groups_by_pair[pair].append(group)

        changed = False
        for pair, pair_groups in groups_by_pair.items():
            if len(pair_groups) < 2:
                continue
            keep = sorted(pair_groups, key=lambda g: (0 if g["owner_user_id"] == pair[0] else 1, g["id"]))[0]
            keep_id = keep["id"]
            keep_owner = keep["owner_user_id"]
            share_by_user = {
                user_id: any(
                    bool(row["share_full_list"])
                    for group in pair_groups
                    for row in members_by_group.get(group["id"], [])
                    if row["member_user_id"] == user_id
                )
                for user_id in pair
            }

            for user_id in pair:
                should_have_row = user_id != keep_owner or share_by_user[user_id]
                if not should_have_row:
                    continue
                existing = conn.execute(
                    "SELECT share_full_list FROM partnergroupmember WHERE group_id = ? AND member_user_id = ?",
                    (keep_id, user_id),
                ).fetchone()
                share = 1 if share_by_user[user_id] else 0
                if existing:
                    if share and not bool(existing["share_full_list"]):
                        conn.execute(
                            "UPDATE partnergroupmember SET share_full_list = 1 WHERE group_id = ? AND member_user_id = ?",
                            (keep_id, user_id),
                        )
                        changed = True
                else:
                    conn.execute(
                        "INSERT INTO partnergroupmember (group_id, member_user_id, share_full_list) VALUES (?, ?, ?)",
                        (keep_id, user_id, share),
                    )
                    changed = True

            for group in pair_groups:
                if group["id"] == keep_id:
                    continue
                conn.execute("DELETE FROM partnergroupmember WHERE group_id = ?", (group["id"],))
                conn.execute("DELETE FROM partnergroup WHERE id = ?", (group["id"],))
                changed = True

        if changed:
            conn.commit()

def _migrate_legacy_ratings(self) -> None:
    with self._sqlite() as conn:
        rating_exists = conn.execute(
            "SELECT 1 FROM sqlite_master WHERE type='table' AND name='rating'"
        ).fetchone()
    if not rating_exists:
        return
    with self._sqlite() as conn:
        legacy_rows = conn.execute("SELECT user_id, kink_id, rating FROM rating").fetchall()
    if not legacy_rows:
        with self._sqlite() as conn:
            conn.execute("DROP TABLE rating")
            conn.commit()
        return
    with Session(self.engine) as session:
        for row in legacy_rows:
            normalized = self._normalize_interest_state(row["rating"])
            if normalized not in VALID_RATINGS:
                continue
            existing = session.get(PlayPreference, (row["user_id"], row["kink_id"]))
            if existing:
                continue
            kink = session.get(Kink, row["kink_id"])
            kink_payload = {"name": kink.name, "cluster": kink.cluster} if kink else None
            directions = self._default_directions_for_kink(kink_payload, normalized) if kink_payload else []
            session.add(
                PlayPreference(
                    user_id=row["user_id"],
                    kink_id=row["kink_id"],
                    interest_state=normalized,
                    directions_csv=self._serialize_directions(directions),
                )
            )
        session.commit()
    with self._sqlite() as conn:
        conn.execute("DROP TABLE rating")
        conn.commit()