src/pageparse/store.py

from __future__ import annotations

import hashlib
import json
import sqlite3
from pathlib import Path

from pageparse.config import settings
from pageparse.schema import SourceParseResult


class Store:
    def __init__(self, db_path: str | Path | None = None) -> None:
        self.db_path = Path(db_path or settings.db_path)
        self._conn: sqlite3.Connection | None = None

    @property
    def conn(self) -> sqlite3.Connection:
        if self._conn is None:
            self._conn = sqlite3.connect(str(self.db_path), check_same_thread=False)
            self._conn.row_factory = sqlite3.Row
        return self._conn

    def init_db(self) -> None:
        cursor = self.conn.cursor()
        try:
            cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='pages'")
            has_pages = cursor.fetchone() is not None
        except sqlite3.OperationalError:
            has_pages = False

        try:
            cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='sources'")
            has_sources = cursor.fetchone() is not None
        except sqlite3.OperationalError:
            has_sources = False

        self.conn.executescript(
            """
            CREATE TABLE IF NOT EXISTS sources (
                id            INTEGER PRIMARY KEY AUTOINCREMENT,
                filename      TEXT NOT NULL,
                source_type   TEXT CHECK (source_type IN ('image','audio','video','document','spreadsheet')),
                source_path   TEXT,
                captured_date TEXT,
                title         TEXT,
                summary       TEXT,
                raw_text      TEXT,
                created_at    TEXT DEFAULT (datetime('now')),
                image_path    TEXT,
                cleaned_image_path TEXT,
                content_hash  TEXT,
                barcodes      TEXT,
                tables        TEXT
            );

            CREATE TABLE IF NOT EXISTS records (
                id          INTEGER PRIMARY KEY AUTOINCREMENT,
                source_id   INTEGER REFERENCES sources(id),
                type        TEXT,
                content     TEXT NOT NULL,
                due_date    TEXT,
                priority    TEXT CHECK (priority IN ('high','medium','low') OR priority IS NULL),
                category    TEXT,
                speaker     TEXT,
                timestamp   TEXT,
                status      TEXT DEFAULT 'todo',
                confidence  REAL,
                user_edited INTEGER DEFAULT 0
            );
        """
        )

        try:
            cursor.execute("PRAGMA table_info(sources)")
            existing_columns = {row[1] for row in cursor.fetchall()}
            for col_name, col_type in [("content_hash", "TEXT"), ("barcodes", "TEXT"), ("tables", "TEXT")]:
                if col_name not in existing_columns:
                    self.conn.execute(f"ALTER TABLE sources ADD COLUMN {col_name} {col_type}")
            self.conn.commit()
        except Exception as e:
            print(f"Error migrating sources table columns: {e}")

        if has_pages and not has_sources:
            try:
                self.conn.execute(
                    """
                    INSERT INTO sources (id, filename, source_type, captured_date, raw_text, created_at, image_path, cleaned_image_path, summary)
                    SELECT id, filename, 'image', captured_date, raw_ocr_text, created_at, image_path, cleaned_image_path, summary
                    FROM pages
                """
                )
                self.conn.execute(
                    """
                    INSERT INTO records (source_id, type, content, due_date, priority, category, status, confidence)
                    SELECT page_id, 'task', task, due_date, priority, category, status, ocr_confidence
                    FROM tasks
                """
                )
                self.conn.execute("DROP TABLE tasks")
                self.conn.execute("DROP TABLE pages")
                self.conn.commit()
            except Exception as e:
                print(f"Migration error: {e}")
                self.conn.rollback()

        self.conn.commit()

    def save(
        self,
        result: SourceParseResult,
        raw_text: str,
        image_path: str | None = None,
        cleaned_image_path: str | None = None,
        content_hash: str | None = None,
    ) -> int:
        cursor = self.conn.execute(
            "INSERT INTO sources (filename, source_type, captured_date, title, summary, raw_text, image_path, cleaned_image_path, content_hash, barcodes, tables) "
            "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
            (
                result.source_file,
                result.source_type,
                str(result.captured_date) if result.captured_date else None,
                result.title,
                result.summary,
                raw_text,
                image_path,
                cleaned_image_path,
                content_hash,
                json.dumps(result.barcodes) if result.barcodes else None,
                json.dumps(result.tables) if result.tables else None,
            ),
        )
        source_id = cursor.lastrowid
        assert source_id is not None

        for record in result.records:
            self.conn.execute(
                "INSERT INTO records (source_id, type, content, due_date, priority, category, speaker, timestamp, status, confidence) "
                "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
                (
                    source_id,
                    record.type,
                    record.content,
                    str(record.due_date) if record.due_date else None,
                    record.priority,
                    record.category,
                    record.speaker,
                    record.timestamp,
                    record.status,
                    record.confidence,
                ),
            )
        self.conn.commit()

        self._fire_webhook(source_id, result, raw_text)
        return source_id

    def _fire_webhook(self, source_id: int, result: SourceParseResult, raw_text: str) -> None:
        url = settings.webhook_url
        if not url:
            return
        try:
            import httpx
            payload = {
                "event": "source.created",
                "source_id": source_id,
                "source_file": result.source_file,
                "source_type": result.source_type,
                "records_count": len(result.records),
                "records": [r.model_dump(mode="json") for r in result.records],
                "raw_text": raw_text[:5000],
            }
            headers = {"Content-Type": "application/json"}
            if settings.webhook_secret:
                headers["X-Webhook-Secret"] = settings.webhook_secret
            httpx.post(url, json=payload, headers=headers, timeout=10)
        except Exception as e:
            print(f"Webhook notification failed: {e}")

    def update_source_summary(self, source_id: int, summary: str) -> bool:
        self.conn.execute(
            "UPDATE sources SET summary = ? WHERE id = ?",
            (summary, source_id),
        )
        self.conn.commit()
        return True

    def list_sources(self) -> list[dict]:
        rows = self.conn.execute("SELECT * FROM sources ORDER BY created_at DESC").fetchall()
        return [dict(r) for r in rows]

    def get_records(
        self,
        source_id: int | None = None,
        priority: str | None = None,
        type: str | None = None,
    ) -> list[dict]:
        query = "SELECT r.*, s.filename FROM records r JOIN sources s ON r.source_id = s.id"
        params: list = []
        conditions = []
        if source_id is not None:
            conditions.append("r.source_id = ?")
            params.append(source_id)
        if priority:
            conditions.append("r.priority = ?")
            params.append(priority)
        if type:
            conditions.append("r.type = ?")
            params.append(type)
        if conditions:
            query += " WHERE " + " AND ".join(conditions)
        query += " ORDER BY r.id DESC"
        rows = self.conn.execute(query, params).fetchall()
        return [dict(r) for r in rows]

    def update_record(self, record_id: int, updates: dict) -> bool:
        if not updates:
            return False
        fields = []
        params = []
        for k, v in updates.items():
            if k in (
                "content",
                "due_date",
                "priority",
                "category",
                "status",
                "confidence",
                "speaker",
                "timestamp",
                "type",
                "user_edited",
            ):
                fields.append(f"{k} = ?")
                params.append(v)
        if not fields:
            return False
        params.append(record_id)
        self.conn.execute(f"UPDATE records SET {', '.join(fields)} WHERE id = ?", params)
        self.conn.commit()
        return True

    def delete_record(self, record_id: int) -> bool:
        self.conn.execute("DELETE FROM records WHERE id = ?", (record_id,))
        self.conn.commit()
        return True

    def get_source_by_hash(self, content_hash: str) -> dict | None:
        row = self.conn.execute(
            "SELECT * FROM sources WHERE content_hash = ?", (content_hash,)
        ).fetchone()
        return dict(row) if row else None

    def get_stats(self) -> dict:
        source_count = self.conn.execute("SELECT COUNT(*) FROM sources").fetchone()[0]
        record_count = self.conn.execute("SELECT COUNT(*) FROM records").fetchone()[0]
        return {
            "sources": source_count,
            "records": record_count,
        }

    def close(self) -> None:
        if self._conn is not None:
            self._conn.close()
            self._conn = None