database.py

# Optimized SQLite DB layer for Character Card CLI v3.5 — with JSON field parsing + complete

import sqlite3
import os
import json
from typing import List, Dict, Optional, Iterable
from contextlib import contextmanager

class CardDatabase:
    """
    SQLite wrapper with WAL/PRAGMA tuning, UPSERT by filename, transactional bulk ops,
    and normalized tag storage (cards, tags, card_tags).
    """
    
    def __init__(self, db_folder: str):
        if not os.path.exists(db_folder):
             # Если папка не существует, мы не можем создать там БД.
             # Логика создания папки должна быть снаружи.
             raise FileNotFoundError(f"Card folder not found: {db_folder}")

        self.db_folder = db_folder
        self.db_path = os.path.join(db_folder, "cards.db")
        self.conn = sqlite3.connect(self.db_path, check_same_thread=False)
        self.conn.row_factory = sqlite3.Row
        self._tune()
        self._create_tables()
    
    def _tune(self):
        cur = self.conn.cursor()
        cur.execute("PRAGMA journal_mode=WAL;")
        cur.execute("PRAGMA synchronous=NORMAL;")
        cur.execute("PRAGMA foreign_keys=ON;")
        cur.execute("PRAGMA temp_store=MEMORY;")
        cur.execute("PRAGMA cache_size=-20000;")  # ~20MB
        self.conn.commit()
    
    def _create_tables(self):
        cur = self.conn.cursor()
        
        cur.execute("""
        CREATE TABLE IF NOT EXISTS cards (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            filename TEXT UNIQUE NOT NULL,
            filepath TEXT NOT NULL,
            creator TEXT DEFAULT 'Unknown',
            name TEXT DEFAULT '',
            description TEXT DEFAULT '',
            first_mes TEXT DEFAULT '',
            scenario TEXT DEFAULT '',
            personality TEXT DEFAULT '',
            mes_example TEXT DEFAULT '',
            avatar TEXT DEFAULT '',
            system_prompt TEXT DEFAULT '',
            post_history_instructions TEXT DEFAULT '',
            alternate_greetings TEXT DEFAULT '', -- JSON array
            character_book TEXT DEFAULT '', -- JSON object  
            extensions TEXT DEFAULT '', -- JSON object
            spec TEXT DEFAULT 'chara_card_v2',
            spec_version TEXT DEFAULT '2.0',
            -- Поля из метаданных
            create_date TEXT DEFAULT '',
            chat TEXT DEFAULT '',
            -- Файловая информация
            file_size INTEGER DEFAULT 0,
            file_mtime REAL DEFAULT 0,
            file_ctime REAL DEFAULT 0,
            signature TEXT DEFAULT '', -- SHA256 checksum
            -- Ссылки (предварительно обработанные)
            hub_links TEXT DEFAULT '', -- JSON array
            -- Chub.ai Sync
            chub_id TEXT DEFAULT '', 
            chub_url TEXT DEFAULT '',
            chub_updated_at TEXT DEFAULT '',
            content_hash TEXT DEFAULT '', -- Semantic hash of core fields
            -- Кэшированные теги для поиска
            tags_text TEXT DEFAULT '',
            last_scanned REAL DEFAULT 0,
            created_at REAL DEFAULT (julianday('now')),
            updated_at REAL DEFAULT (julianday('now'))
        );
        """)
        
        cur.execute("""
        CREATE TABLE IF NOT EXISTS tags (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            tag_name TEXT UNIQUE NOT NULL
        );
        """)
        
        cur.execute("""
        CREATE TABLE IF NOT EXISTS card_tags (
            card_id INTEGER NOT NULL,
            tag_id INTEGER NOT NULL,
            PRIMARY KEY (card_id, tag_id),
            FOREIGN KEY (card_id) REFERENCES cards(id) ON DELETE CASCADE,
            FOREIGN KEY (tag_id) REFERENCES tags(id) ON DELETE CASCADE
        );
        """)
        
        # FTS5 Virtual Table for Full Text Search
        # Мы используем content='cards' (External Content Table) для экономии места.
        cur.execute("""
        CREATE VIRTUAL TABLE IF NOT EXISTS cards_fts USING fts5(
            name, 
            creator, 
            description, 
            first_mes, 
            scenario, 
            personality, 
            mes_example, 
            system_prompt, 
            tags_text, 
            content='cards', 
            content_rowid='id'
        );
        """)

        # Triggers to keep FTS index in sync with the main table
        cur.execute("""
        CREATE TRIGGER IF NOT EXISTS cards_ai AFTER INSERT ON cards BEGIN
            INSERT INTO cards_fts(rowid, name, creator, description, first_mes, scenario, personality, mes_example, system_prompt, tags_text)
            VALUES (new.id, new.name, new.creator, new.description, new.first_mes, new.scenario, new.personality, new.mes_example, new.system_prompt, new.tags_text);
        END;
        """)

        cur.execute("""
        CREATE TRIGGER IF NOT EXISTS cards_ad AFTER DELETE ON cards BEGIN
            INSERT INTO cards_fts(cards_fts, rowid, name, creator, description, first_mes, scenario, personality, mes_example, system_prompt, tags_text)
            VALUES('delete', old.id, old.name, old.creator, old.description, old.first_mes, old.scenario, old.personality, old.mes_example, old.system_prompt, old.tags_text);
        END;
        """)

        cur.execute("""
        CREATE TRIGGER IF NOT EXISTS cards_au AFTER UPDATE ON cards BEGIN
            INSERT INTO cards_fts(cards_fts, rowid, name, creator, description, first_mes, scenario, personality, mes_example, system_prompt, tags_text)
            VALUES('delete', old.id, old.name, old.creator, old.description, old.first_mes, old.scenario, old.personality, old.mes_example, old.system_prompt, old.tags_text);
            INSERT INTO cards_fts(rowid, name, creator, description, first_mes, scenario, personality, mes_example, system_prompt, tags_text)
            VALUES (new.id, new.name, new.creator, new.description, new.first_mes, new.scenario, new.personality, new.mes_example, new.system_prompt, new.tags_text);
        END;
        """)
        
        cur.execute("CREATE INDEX IF NOT EXISTS idx_cards_filename ON cards(filename);")
        cur.execute("CREATE INDEX IF NOT EXISTS idx_cards_creator ON cards(creator);")
        cur.execute("CREATE INDEX IF NOT EXISTS idx_cards_mtime ON cards(file_mtime);")
        cur.execute("CREATE INDEX IF NOT EXISTS idx_cards_signature ON cards(signature);")
        cur.execute("CREATE INDEX IF NOT EXISTS idx_tags_name ON tags(tag_name);")
        
        # Добавляем новые столбцы в существующие таблицы (миграция)
        self._migrate_schema()
        self.conn.commit()
    
    def _migrate_schema(self):
        """Добавляем новые столбцы в существующие таблицы"""
        cur = self.conn.cursor()
        
        # Проверяем какие столбцы уже есть
        cur.execute("PRAGMA table_info(cards);")
        existing_columns = {row[1] for row in cur.fetchall()}
        
        # Добавляем недостающие столбцы  
        new_columns = {
            'create_date': 'TEXT DEFAULT ""',
            'chat': 'TEXT DEFAULT ""', 
            'file_ctime': 'REAL DEFAULT 0',
            'signature': 'TEXT DEFAULT ""',
            'hub_links': 'TEXT DEFAULT ""',
            'tags_text': 'TEXT DEFAULT ""',
            'chub_id': 'TEXT DEFAULT ""',
            'chub_url': 'TEXT DEFAULT ""',
            'chub_updated_at': 'TEXT DEFAULT ""',
            'content_hash': 'TEXT DEFAULT ""'
        }
        
        for col_name, col_def in new_columns.items():
            if col_name not in existing_columns:
                try:
                    cur.execute(f"ALTER TABLE cards ADD COLUMN {col_name} {col_def};")
                except sqlite3.OperationalError:
                    pass  # Column already exists
    
    def close(self):
        if self.conn:
            self.conn.close()
    
    def __enter__(self):
        return self
    
    def __exit__(self, exc_type, exc, tb):
        self.close()
    
    @contextmanager
    def transaction(self):
        """
        BEGIN IMMEDIATE for better writer throughput; commit/rollback automatically.
        """
        cur = self.conn.cursor()
        cur.execute("BEGIN IMMEDIATE;")
        try:
            yield
            self.conn.commit()
        except Exception:
            self.conn.rollback()
            raise
    
    def upsert_card(self, filename: str, filepath: str, creator: str = "Unknown",
                   tags: List[str] = None, metadata: Dict = None,
                   file_size: int = 0, file_mtime: float = 0, file_ctime: float = 0,
                   signature: str = "", hub_links: List[str] = None,
                   create_date: str = "", chat: str = "",
                   chub_id: str = "", chub_url: str = "", 
                   chub_updated_at: str = "", content_hash: str = "") -> int:
        """
        Insert or update a card by filename (UPSERT).
        Returns card id.
        """
        tags = tags or []
        metadata = metadata or {}
        hub_links = hub_links or []
        
        cur = self.conn.cursor()
        
        alt_greetings = json.dumps(metadata.get("alternate_greetings", []), ensure_ascii=False)
        character_book = json.dumps(metadata.get("character_book", {}), ensure_ascii=False)
        extensions = json.dumps(metadata.get("extensions", {}), ensure_ascii=False)
        hub_links_json = json.dumps(hub_links, ensure_ascii=False)
        
        cur.execute("""
        INSERT INTO cards (
            filename, filepath, creator, name, description, first_mes, scenario,
            personality, mes_example, avatar, system_prompt, post_history_instructions,
            alternate_greetings, character_book, extensions, spec, spec_version,
            create_date, chat, file_size, file_mtime, file_ctime, signature, hub_links,
            chub_id, chub_url, chub_updated_at, content_hash,
            last_scanned, updated_at
        ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, julianday('now'), julianday('now'))
        ON CONFLICT(filename) DO UPDATE SET
            filepath=excluded.filepath,
            creator=excluded.creator,
            name=excluded.name,
            description=excluded.description,
            first_mes=excluded.first_mes,
            scenario=excluded.scenario,
            personality=excluded.personality,
            mes_example=excluded.mes_example,
            avatar=excluded.avatar,
            system_prompt=excluded.system_prompt,
            post_history_instructions=excluded.post_history_instructions,
            alternate_greetings=excluded.alternate_greetings,
            character_book=excluded.character_book,
            extensions=excluded.extensions,
            spec=excluded.spec,
            spec_version=excluded.spec_version,
            create_date=excluded.create_date,
            chat=excluded.chat,
            file_size=excluded.file_size,
            file_mtime=excluded.file_mtime,
            file_ctime=excluded.file_ctime,
            signature=excluded.signature,
            hub_links=excluded.hub_links,
            chub_id=excluded.chub_id,
            chub_url=excluded.chub_url,
            chub_updated_at=excluded.chub_updated_at,
            content_hash=excluded.content_hash,
            last_scanned=julianday('now'),
            updated_at=julianday('now');
        """, (
            filename, filepath, creator,
            metadata.get('name', ''),
            metadata.get('description', ''),
            metadata.get('first_mes', ''),
            metadata.get('scenario', ''),
            metadata.get('personality', ''),
            metadata.get('mes_example', ''),
            metadata.get('avatar', ''),
            metadata.get('system_prompt', ''),
            metadata.get('post_history_instructions', ''),
            alt_greetings, character_book, extensions,
            metadata.get('spec', 'chara_card_v2'),
            metadata.get('spec_version', '2.0'),
            create_date, chat,
            int(file_size), float(file_mtime), float(file_ctime), signature, hub_links_json,
            chub_id, chub_url, chub_updated_at, content_hash
        ))
        
        # Retrieve id reliably
        cur.execute("SELECT id FROM cards WHERE filename = ?;", (filename,))
        row = cur.fetchone()
        card_id = row["id"] if row else None
        
        # Update tags
        self._set_tags(card_id, tags)
        
        return card_id
    
    def _set_tags(self, card_id: Optional[int], tags: List[str]):
        if not card_id:
            return
            
        cur = self.conn.cursor()
        tags = [t.strip() for t in (tags or []) if t and t.strip()]
        
        # Insert any new tags
        cur.executemany("INSERT OR IGNORE INTO tags (tag_name) VALUES (?);", [(t,) for t in tags])
        
        # Map tag_name -> id
        if tags:
            placeholders = ",".join(["?"] * len(tags))
            cur.execute(f"SELECT id, tag_name FROM tags WHERE tag_name IN ({placeholders});", tags)
            tag_map = {r["tag_name"]: r["id"] for r in cur.fetchall()}
        else:
            tag_map = {}
        
        # Replace links
        cur.execute("DELETE FROM card_tags WHERE card_id = ?;", (card_id,))
        cur.executemany(
            "INSERT OR IGNORE INTO card_tags (card_id, tag_id) VALUES (?, ?);",
            [(card_id, tag_map[t]) for t in tags if t in tag_map]
        )
        
        # UPDATE FTS TAGS INDEX (через обновление основной таблицы, триггер подхватит)
        tags_str = " ".join(tags)
        cur.execute("UPDATE cards SET tags_text = ? WHERE id = ?;", (tags_str, card_id))
    
    def _parse_json_field(self, field_value):
        """Парсим JSON-поле, если это строка"""
        if isinstance(field_value, str) and field_value.strip():
            try:
                return json.loads(field_value)
            except Exception:
                return field_value
        return field_value
    
    def search_fts(self, query_text: str) -> List[Dict]:
        """
        Полнотекстовый поиск по всем текстовым полям с использованием FTS5.
        Поддерживает синтаксис FTS (AND, OR, NOT, "phrase", pre*).
        """
        if not query_text or not query_text.strip():
            return []

        cur = self.conn.cursor()
        
        # Очистка запроса от опасных символов (простая защита)
        # В FTS5 лучше экранировать двойные кавычки, если они есть
        safe_query = query_text.replace('"', '""')
        
        # Формируем запрос. Используем подзапрос для получения ID из FTS
        sql = """
        SELECT c.*, cards_fts.rank 
        FROM cards c
        JOIN cards_fts ON c.id = cards_fts.rowid
        WHERE cards_fts MATCH ?
        ORDER BY cards_fts.rank;
        """
        
        try:
            cur.execute(sql, (safe_query,))
        except sqlite3.OperationalError:
            # Fallback если синтаксис запроса некорректен для FTS (например, незакрытые кавычки)
            # Можно попробовать поискать как обычную фразу
            try:
                cur.execute(sql, (f'"{safe_query}"',))
            except Exception:
                 return []

        rows = cur.fetchall()
        cards: List[Dict] = []
        
        if not rows:
            return cards

        # Preload tags logic (similar to get_all_cards)
        ids = [r["id"] for r in rows]
        placeholders = ",".join(["?"] * len(ids))
        
        cur.execute(f"""
        SELECT ct.card_id, t.tag_name
        FROM card_tags ct JOIN tags t ON t.id = ct.tag_id
        WHERE ct.card_id IN ({placeholders});
        """, ids)
        
        tag_rows = cur.fetchall()
        tags_by_card: Dict[int, List[str]] = {}
        for tr in tag_rows:
            tags_by_card.setdefault(tr["card_id"], []).append(tr["tag_name"])
        
        for row in rows:
            c = dict(row)
            c["tags"] = tags_by_card.get(c["id"], [])
            c["alternate_greetings"] = self._parse_json_field(c["alternate_greetings"])
            c["character_book"] = self._parse_json_field(c["character_book"])
            c["extensions"] = self._parse_json_field(c["extensions"])
            c["hub_links"] = self._parse_json_field(c["hub_links"])
            cards.append(c)
            
        return cards
    
    def get_card_by_filename(self, filename: str) -> Optional[Dict]:
        cur = self.conn.cursor()
        cur.execute("SELECT * FROM cards WHERE filename = ?;", (filename,))
        row = cur.fetchone()
        if not row:
            return None
        
        card = dict(row)
        card["tags"] = self._get_tags_for_card(card["id"])
        
        # Парсим JSON-поля
        card["alternate_greetings"] = self._parse_json_field(card["alternate_greetings"])
        card["character_book"] = self._parse_json_field(card["character_book"])
        card["extensions"] = self._parse_json_field(card["extensions"])
        card["hub_links"] = self._parse_json_field(card["hub_links"])
        
        return card
    
    def _get_tags_for_card(self, card_id: int) -> List[str]:
        cur = self.conn.cursor()  
        cur.execute("""
        SELECT t.tag_name FROM tags t
        JOIN card_tags ct ON t.id = ct.tag_id
        WHERE ct.card_id = ?;
        """, (card_id,))
        return [r["tag_name"] for r in cur.fetchall()]
    
    def get_all_cards(self, order_by: str = "filename") -> List[Dict]:
        cur = self.conn.cursor()
        if order_by not in ("filename","creator","file_mtime"):
            order_by = "filename"
        
        cur.execute(f"SELECT * FROM cards ORDER BY {order_by} COLLATE NOCASE;")
        cards: List[Dict] = []
        rows = cur.fetchall()
        if not rows:
            return cards
        
        # Preload all tags to cut N+1
        cur.execute("""
        SELECT ct.card_id, t.tag_name
        FROM card_tags ct JOIN tags t ON t.id = ct.tag_id;
        """)
        tag_rows = cur.fetchall()
        tags_by_card: Dict[int, List[str]] = {}
        for tr in tag_rows:
            tags_by_card.setdefault(tr["card_id"], []).append(tr["tag_name"])
        
        for row in rows:
            c = dict(row)
            c["tags"] = tags_by_card.get(c["id"], [])
            
            # Парсим JSON-поля
            c["alternate_greetings"] = self._parse_json_field(c["alternate_greetings"])
            c["character_book"] = self._parse_json_field(c["character_book"])
            c["extensions"] = self._parse_json_field(c["extensions"])
            c["hub_links"] = self._parse_json_field(c["hub_links"])
            
            cards.append(c)
        
        return cards
    
    def remove_missing_files(self, existing_files: List[str]) -> int:
        """
        Remove cards whose filename not present in 'existing_files'.
        """
        cur = self.conn.cursor()
        if existing_files:
            placeholders = ",".join(["?"] * len(existing_files))
            cur.execute(f"DELETE FROM cards WHERE filename NOT IN ({placeholders});", existing_files)
        else:
            cur.execute("DELETE FROM cards;")
        
        removed = cur.rowcount
        self.conn.commit()
        return removed
    
    def get_all_creators(self) -> List[str]:
        cur = self.conn.cursor()
        cur.execute("SELECT DISTINCT creator FROM cards WHERE creator != 'Unknown' ORDER BY creator COLLATE NOCASE;")
        return [r["creator"] for r in cur.fetchall()]
    
    def get_all_tags(self) -> List[str]:
        cur = self.conn.cursor()
        cur.execute("SELECT tag_name FROM tags ORDER BY tag_name COLLATE NOCASE;")
        return [r["tag_name"] for r in cur.fetchall()]
    
    def search_cards(self, search_term: str = "", search_type: str = "filename") -> List[Dict]:
        cur = self.conn.cursor()
        if not search_term:
            return self.get_all_cards()
        
        if search_type == "filename":
            cur.execute("SELECT * FROM cards WHERE filename LIKE ? ORDER BY filename COLLATE NOCASE;",
                       (f"%{search_term}%",))
        elif search_type == "creator":
            cur.execute("SELECT * FROM cards WHERE creator LIKE ? ORDER BY filename COLLATE NOCASE;",
                       (f"%{search_term}%",))
        elif search_type == "tag":
            cur.execute("""
            SELECT DISTINCT c.* FROM cards c
            JOIN card_tags ct ON c.id = ct.card_id
            JOIN tags t ON ct.tag_id = t.id
            WHERE t.tag_name = ?
            ORDER BY c.filename COLLATE NOCASE;
            """, (search_term,))
        else:
            return []
        
        rows = cur.fetchall()
        cards: List[Dict] = []
        
        for row in rows:
            c = dict(row)
            c["tags"] = self._get_tags_for_card(c["id"])
            
            # Парсим JSON-поля
            c["alternate_greetings"] = self._parse_json_field(c["alternate_greetings"])
            c["character_book"] = self._parse_json_field(c["character_book"])
            c["extensions"] = self._parse_json_field(c["extensions"])
            c["hub_links"] = self._parse_json_field(c["hub_links"])
            
            cards.append(c)
        
        return cards
    
    def get_stats(self) -> Dict[str, int]:
        cur = self.conn.cursor()
        
        cur.execute("SELECT COUNT(*) AS total_cards FROM cards;")
        total_cards = cur.fetchone()["total_cards"]
        
        cur.execute("SELECT COUNT(DISTINCT creator) AS unique_creators FROM cards WHERE creator != 'Unknown';")
        unique_creators = cur.fetchone()["unique_creators"]
        
        cur.execute("SELECT COUNT(*) AS total_tags FROM tags;")
        total_tags = cur.fetchone()["total_tags"]
        
        return {
            "total_cards": total_cards,
            "unique_creators": unique_creators,
            "total_tags": total_tags
        }