# SPDX-License-Identifier: GPL-2.0-only
#
# llamachat - a small native chat client for a local llama.cpp router
# Copyright (C) 2026 Danilo M. <danix@danix.xyz>
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License version 2 as
# published by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
"""SQLite history store with FTS5 search."""

import sqlite3
import time
from pathlib import Path

SCHEMA = """
PRAGMA journal_mode = WAL;
PRAGMA foreign_keys = ON;

CREATE TABLE IF NOT EXISTS sessions (
    id            INTEGER PRIMARY KEY,
    mode          TEXT    NOT NULL,
    title         TEXT,
    model         TEXT,
    prompt_name   TEXT    NOT NULL DEFAULT '',
    prompt_custom TEXT    NOT NULL DEFAULT '',
    created_at    INTEGER NOT NULL,
    updated_at    INTEGER NOT NULL
);

CREATE TABLE IF NOT EXISTS messages (
    id                INTEGER PRIMARY KEY,
    session_id        INTEGER NOT NULL REFERENCES sessions(id) ON DELETE CASCADE,
    role              TEXT    NOT NULL,
    content           TEXT    NOT NULL,
    reasoning         TEXT    NOT NULL DEFAULT '',
    searches          TEXT,
    prompt_tokens     INTEGER,
    completion_tokens INTEGER,
    model             TEXT,
    created_at        INTEGER NOT NULL
);

CREATE INDEX IF NOT EXISTS idx_messages_session ON messages(session_id, id);
CREATE INDEX IF NOT EXISTS idx_sessions_updated
    ON sessions(updated_at DESC, id DESC);

CREATE TABLE IF NOT EXISTS attachments (
    id         INTEGER PRIMARY KEY,
    message_id INTEGER NOT NULL REFERENCES messages(id) ON DELETE CASCADE,
    path       TEXT    NOT NULL,
    kind       TEXT    NOT NULL,
    mime       TEXT,
    size       INTEGER,
    sha256     TEXT,
    thumb      BLOB,
    truncated  INTEGER NOT NULL DEFAULT 0
);

CREATE INDEX IF NOT EXISTS idx_attachments_message ON attachments(message_id);

CREATE VIRTUAL TABLE IF NOT EXISTS messages_fts USING fts5(
    content,
    content='messages',
    content_rowid='id',
    tokenize='unicode61'
);

CREATE TRIGGER IF NOT EXISTS messages_ai AFTER INSERT ON messages BEGIN
    INSERT INTO messages_fts(rowid, content) VALUES (new.id, new.content);
END;

CREATE TRIGGER IF NOT EXISTS messages_ad AFTER DELETE ON messages BEGIN
    INSERT INTO messages_fts(messages_fts, rowid, content)
        VALUES ('delete', old.id, old.content);
END;

CREATE TRIGGER IF NOT EXISTS messages_au AFTER UPDATE ON messages BEGIN
    INSERT INTO messages_fts(messages_fts, rowid, content)
        VALUES ('delete', old.id, old.content);
    INSERT INTO messages_fts(rowid, content) VALUES (new.id, new.content);
END;
"""


def _now() -> int:
    return int(time.time())


class History:
    """Session/message/attachment storage. One connection, one thread."""

    def __init__(self, path: Path):
        path.parent.mkdir(parents=True, exist_ok=True)
        self.conn = sqlite3.connect(path)
        self.conn.row_factory = sqlite3.Row
        self.conn.executescript(SCHEMA)
        self._migrate()
        self.conn.commit()

    def _migrate(self) -> None:
        """Add columns introduced after a database was first created."""
        added = {
            "messages": {
                "reasoning": "TEXT NOT NULL DEFAULT ''",
                # Nullable rather than defaulted: NULL means "no searches",
                # which is exactly what every pre-migration row wants.
                "searches": "TEXT",
                # Nullable for the same reason: an old row has no counts,
                # and zero would read as a reply that cost nothing.
                "prompt_tokens": "INTEGER",
                "completion_tokens": "INTEGER",
                # The model on `sessions` is the current one, which prices a
                # switched conversation wrongly. Record what actually replied.
                "model": "TEXT",
            },
            "sessions": {
                "prompt_name": "TEXT NOT NULL DEFAULT ''",
                "prompt_custom": "TEXT NOT NULL DEFAULT ''",
            },
        }
        for table, wanted in added.items():
            existing = {
                row["name"]
                for row in self.conn.execute(f"PRAGMA table_info({table})")
            }
            for column, spec in wanted.items():
                if column not in existing:
                    self.conn.execute(
                        f"ALTER TABLE {table} ADD COLUMN {column} {spec}"
                    )

    def close(self) -> None:
        self.conn.close()

    # -- sessions ---------------------------------------------------------

    def create_session(
        self,
        mode: str,
        model: str,
        title: str = "",
        prompt_name: str = "",
        prompt_custom: str = "",
    ) -> int:
        now = _now()
        cur = self.conn.execute(
            "INSERT INTO sessions"
            " (mode, title, model, prompt_name, prompt_custom,"
            "  created_at, updated_at)"
            " VALUES (?, ?, ?, ?, ?, ?, ?)",
            (mode, title, model, prompt_name, prompt_custom, now, now),
        )
        self.conn.commit()
        return int(cur.lastrowid)

    def set_prompt(self, session_id: int, name: str, custom: str = "") -> None:
        """Record which system prompt a conversation was built with."""
        self.conn.execute(
            "UPDATE sessions SET prompt_name = ?, prompt_custom = ?"
            " WHERE id = ?",
            (name, custom, session_id),
        )
        self.conn.commit()

    def touch_session(self, session_id: int, model: str | None = None) -> None:
        if model:
            self.conn.execute(
                "UPDATE sessions SET updated_at = ?, model = ? WHERE id = ?",
                (_now(), model, session_id),
            )
        else:
            self.conn.execute(
                "UPDATE sessions SET updated_at = ? WHERE id = ?",
                (_now(), session_id),
            )
        self.conn.commit()

    def set_title(self, session_id: int, title: str) -> None:
        self.conn.execute(
            "UPDATE sessions SET title = ? WHERE id = ?", (title, session_id)
        )
        self.conn.commit()

    def get_session(self, session_id: int) -> sqlite3.Row | None:
        return self.conn.execute(
            "SELECT * FROM sessions WHERE id = ?", (session_id,)
        ).fetchone()

    def recent_sessions(self, limit: int = 200) -> list[sqlite3.Row]:
        # id breaks ties: several sessions can share a one-second timestamp.
        return self.conn.execute(
            "SELECT * FROM sessions ORDER BY updated_at DESC, id DESC LIMIT ?",
            (limit,),
        ).fetchall()

    def delete_session(self, session_id: int) -> None:
        self.conn.execute("DELETE FROM sessions WHERE id = ?", (session_id,))
        self.conn.commit()

    # -- messages ---------------------------------------------------------

    def add_message(self, session_id: int, role: str, content: str) -> int:
        cur = self.conn.execute(
            "INSERT INTO messages (session_id, role, content, created_at)"
            " VALUES (?, ?, ?, ?)",
            (session_id, role, content, _now()),
        )
        self.conn.execute(
            "UPDATE sessions SET updated_at = ? WHERE id = ?",
            (_now(), session_id),
        )
        self.conn.commit()
        return int(cur.lastrowid)

    def update_message(
        self,
        message_id: int,
        content: str,
        reasoning: str | None = None,
        searches: str | None = None,
        prompt_tokens: int | None = None,
        completion_tokens: int | None = None,
        model: str | None = None,
    ) -> None:
        """Fill in a streamed reply. Omitted fields keep their stored value.

        ponytail: no caller can write NULL, since None means "keep". No
        caller needs to: every field here is filled once when a reply
        finishes and never cleared. The upgrade path if one ever does is a
        sentinel object distinct from None, not a second method.
        """
        columns = ["content = ?"]
        values: list = [content]
        for column, value in (
            ("reasoning", reasoning),
            ("searches", searches),
            ("prompt_tokens", prompt_tokens),
            ("completion_tokens", completion_tokens),
            ("model", model),
        ):
            if value is not None:
                columns.append(f"{column} = ?")
                values.append(value)
        values.append(message_id)
        self.conn.execute(
            f"UPDATE messages SET {', '.join(columns)} WHERE id = ?", values
        )
        self.conn.commit()

    def messages(self, session_id: int) -> list[sqlite3.Row]:
        return self.conn.execute(
            "SELECT * FROM messages WHERE session_id = ? ORDER BY id", (session_id,)
        ).fetchall()

    # -- attachments ------------------------------------------------------

    def add_attachment(
        self,
        message_id: int,
        path: str,
        kind: str,
        mime: str = "",
        size: int = 0,
        sha256: str = "",
        thumb: bytes | None = None,
        truncated: bool = False,
    ) -> int:
        cur = self.conn.execute(
            "INSERT INTO attachments"
            " (message_id, path, kind, mime, size, sha256, thumb, truncated)"
            " VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
            (message_id, path, kind, mime, size, sha256, thumb, int(truncated)),
        )
        self.conn.commit()
        return int(cur.lastrowid)

    def attachments(self, message_id: int) -> list[sqlite3.Row]:
        return self.conn.execute(
            "SELECT * FROM attachments WHERE message_id = ? ORDER BY id",
            (message_id,),
        ).fetchall()

    # -- search -----------------------------------------------------------

    def search(self, query: str, limit: int = 100) -> list[sqlite3.Row]:
        """Full-text search over message content.

        Returns one row per matching message, newest first, carrying the
        owning session so the UI can offer to reopen it.
        """
        if not query.strip():
            return []
        return self.conn.execute(
            "SELECT m.id AS message_id, m.session_id, m.role, m.created_at,"
            "       s.title, s.mode, s.model,"
            "       snippet(messages_fts, 0, '<b>', '</b>', '…', 16) AS snip"
            " FROM messages_fts"
            " JOIN messages m ON m.id = messages_fts.rowid"
            " JOIN sessions s ON s.id = m.session_id"
            " WHERE messages_fts MATCH ?"
            " ORDER BY m.created_at DESC, m.id DESC LIMIT ?",
            (_fts_query(query), limit),
        ).fetchall()


def _fts_query(text: str) -> str:
    """Turn user input into a safe FTS5 query.

    Every token is quoted, so punctuation that FTS5 would read as syntax
    (a bare '-', 'AND', a stray quote) cannot raise OperationalError.
    """
    tokens = [t.replace('"', '""') for t in text.split() if t]
    return " ".join(f'"{t}"' for t in tokens)
