src/app_kit/storage.py

from __future__ import annotations

import json
import sqlite3
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any
import shutil
import uuid

from .embedding import SimpleEmbeddingIndex


DEFAULT_DB_PATH = Path('app_data.sqlite3')


def init_db(db_path: str | Path = DEFAULT_DB_PATH, artifact_dir: str | Path | None = None) -> None:
    db_path = Path(db_path)
    artifact_dir = Path(artifact_dir) if artifact_dir is not None else db_path.parent / 'artifacts'
    store = SQLiteStore(db_path, artifact_dir)
    store.close()


def reset_db(db_path: str | Path = DEFAULT_DB_PATH, artifact_dir: str | Path | None = None) -> None:
    db_path = Path(db_path)
    if db_path.exists():
        db_path.unlink()
    init_db(db_path, artifact_dir)

SCHEMA = """
CREATE TABLE IF NOT EXISTS artifacts (
    id TEXT PRIMARY KEY,
    project TEXT NOT NULL,
    pack_id TEXT NOT NULL,
    type TEXT NOT NULL,
    path TEXT NOT NULL,
    created_at TEXT NOT NULL,
    metadata_json TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS records (
    id TEXT PRIMARY KEY,
    project TEXT NOT NULL,
    pack_id TEXT NOT NULL,
    title TEXT NOT NULL,
    primary_text TEXT NOT NULL,
    json_blob TEXT NOT NULL,
    status TEXT NOT NULL,
    created_at TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS embeddings (
    record_id TEXT PRIMARY KEY,
    project TEXT NOT NULL,
    vector_json TEXT NOT NULL,
    metadata_json TEXT NOT NULL,
    created_at TEXT NOT NULL
);
"""


@dataclass
class StoredPackResult:
    record_id: str
    title: str
    primary_text: str
    json_blob: dict[str, Any]
    status: str


def utc_now() -> str:
    return datetime.now(timezone.utc).isoformat()


class SQLiteStore:
    def __init__(self, db_path: str | Path, artifact_dir: str | Path):
        self.db_path = Path(db_path)
        self.artifact_dir = Path(artifact_dir)
        self.db_path.parent.mkdir(parents=True, exist_ok=True)
        self.artifact_dir.mkdir(parents=True, exist_ok=True)
        self._conn = sqlite3.connect(self.db_path, check_same_thread=False)
        self._conn.row_factory = sqlite3.Row
        self._conn.executescript(SCHEMA)
        self._conn.commit()

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

    def store_artifact(self, project: str, pack_id: str, source_path: Path, kind: str, metadata: dict[str, Any] | None = None) -> str:
        artifact_id = str(uuid.uuid4())
        dest = self.artifact_dir / f'{artifact_id}{source_path.suffix or ".bin"}'
        shutil.copy2(source_path, dest)
        self._conn.execute(
            'INSERT INTO artifacts VALUES (?, ?, ?, ?, ?, ?, ?)',
            (artifact_id, project, pack_id, kind, str(dest), utc_now(), json.dumps(metadata or {}, ensure_ascii=False)),
        )
        self._conn.commit()
        return artifact_id

    def store_record(self, project: str, pack_id: str, title: str, primary_text: str, payload: dict[str, Any], status: str = 'stored', record_id: str | None = None) -> str:
        record_id = record_id or str(uuid.uuid4())
        self._conn.execute(
            'INSERT OR REPLACE INTO records VALUES (?, ?, ?, ?, ?, ?, ?, ?)',
            (record_id, project, pack_id, title, primary_text, json.dumps(payload, ensure_ascii=False), status, utc_now()),
        )
        self._conn.commit()
        return record_id

    def store_embedding(self, record_id: str, project: str, text: str, metadata: dict[str, Any] | None = None) -> None:
        vec = SimpleEmbeddingIndex()
        vec.add(record_id, text)
        vector_json = json.dumps({token: count for token, count in vec.entries[record_id].items()}, ensure_ascii=False)
        self._conn.execute(
            'INSERT OR REPLACE INTO embeddings VALUES (?, ?, ?, ?, ?)',
            (record_id, project, vector_json, json.dumps(metadata or {}, ensure_ascii=False), utc_now()),
        )
        self._conn.commit()

    def _embedding_index(self, project: str) -> SimpleEmbeddingIndex:
        index = SimpleEmbeddingIndex()
        rows = self._conn.execute('SELECT record_id, vector_json FROM embeddings WHERE project = ?', (project,)).fetchall()
        for row in rows:
            index.entries[row['record_id']] = __import__('collections').Counter(json.loads(row['vector_json']))
        return index

    def search_records(self, project: str, query: str, limit: int = 5) -> list[dict[str, Any]]:
        index = self._embedding_index(project)
        scored = index.search(query, limit=limit)
        if not scored:
            return []
        ids = [record_id for record_id, score in scored if score > 0]
        if not ids:
            ids = [record_id for record_id, _ in scored]
        out = []
        for record_id in ids:
            row = self._conn.execute('SELECT * FROM records WHERE id = ?', (record_id,)).fetchone()
            if row:
                out.append(dict(row))
        return out[:limit]

    def list_records(self, project: str) -> list[dict[str, Any]]:
        rows = self._conn.execute('SELECT * FROM records WHERE project = ? ORDER BY created_at DESC', (project,)).fetchall()
        return [dict(row) for row in rows]

    def get_record(self, record_id: str) -> dict[str, Any] | None:
        row = self._conn.execute('SELECT * FROM records WHERE id = ?', (record_id,)).fetchone()
        return dict(row) if row else None

    def inbox(self, project: str) -> list[dict[str, Any]]:
        return self.list_records(project)

    def history(self, project: str) -> list[dict[str, Any]]:
        return self.list_records(project)