"""
Marziel Vector Store — SQLite-based with optional TimescaleDB.

Zero-config: uses SQLite + numpy cosine similarity by default.
If TimescaleDB is available, upgrades to pgvector for production scale.

pip install marziel  # includes everything needed
"""

import hashlib
import json
import os
import sqlite3
import time
from pathlib import Path
from typing import Optional


MARZIEL_HOME = os.path.expanduser("~/.marziel")
DB_PATH = os.path.join(MARZIEL_HOME, "rag.db")


def _get_db():
    """Get SQLite connection with WAL mode for concurrency."""
    os.makedirs(MARZIEL_HOME, exist_ok=True)
    conn = sqlite3.connect(DB_PATH, check_same_thread=False)
    conn.row_factory = sqlite3.Row
    conn.execute("PRAGMA journal_mode=WAL")
    conn.execute("PRAGMA foreign_keys=ON")
    return conn


def _init_schema(conn):
    """Create RAG tables if they don't exist."""
    conn.executescript("""
        CREATE TABLE IF NOT EXISTS documents (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            filename TEXT NOT NULL,
            content_hash TEXT UNIQUE,
            page_count INTEGER,
            pageindex_tree TEXT,
            doc_summary TEXT,
            created_at REAL DEFAULT (strftime('%s','now'))
        );

        CREATE TABLE IF NOT EXISTS document_chunks (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            document_id INTEGER REFERENCES documents(id) ON DELETE CASCADE,
            page_num INTEGER,
            section_title TEXT,
            content TEXT NOT NULL,
            embedding BLOB,
            created_at REAL DEFAULT (strftime('%s','now'))
        );

        CREATE INDEX IF NOT EXISTS idx_chunks_doc
            ON document_chunks(document_id);
        CREATE INDEX IF NOT EXISTS idx_docs_hash
            ON documents(content_hash);
    """)
    conn.commit()


class VectorStore:
    """SQLite-backed vector store with cosine similarity search."""

    def __init__(self, db_path: Optional[str] = None):
        self.db_path = db_path or DB_PATH
        self.conn = _get_db()
        _init_schema(self.conn)
        self._embedder = None

    def _get_embedder(self):
        """Lazy-load fastembed for embeddings."""
        if self._embedder is None:
            try:
                from fastembed import TextEmbedding
                self._embedder = TextEmbedding("BAAI/bge-small-en-v1.5")
            except ImportError:
                import subprocess, sys
                subprocess.run(
                    [sys.executable, "-m", "pip", "install", "fastembed", "-q"],
                    check=True, capture_output=True, timeout=120
                )
                from fastembed import TextEmbedding
                self._embedder = TextEmbedding("BAAI/bge-small-en-v1.5")
        return self._embedder

    def _embed(self, texts: list[str]) -> list[list[float]]:
        """Generate embeddings for texts."""
        embedder = self._get_embedder()
        return list(embedder.embed(texts))

    def _cosine_similarity(self, a, b) -> float:
        """Compute cosine similarity between two vectors."""
        import numpy as np
        a, b = np.array(a), np.array(b)
        return float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-8))

    # ── Document Management ─────────────────────────────

    def add_document(self, filename: str, pages: list[str],
                     pageindex_tree: Optional[dict] = None,
                     summary: Optional[str] = None) -> int:
        """Add a document and its page chunks to the store."""
        full_text = "\n".join(pages)
        content_hash = hashlib.sha256(full_text.encode()).hexdigest()[:16]

        # Check if already indexed
        existing = self.conn.execute(
            "SELECT id FROM documents WHERE content_hash = ?",
            (content_hash,)
        ).fetchone()
        if existing:
            return existing["id"]

        # Insert document
        cur = self.conn.execute(
            """INSERT INTO documents (filename, content_hash, page_count,
                                      pageindex_tree, doc_summary)
               VALUES (?, ?, ?, ?, ?)""",
            (filename, content_hash, len(pages),
             json.dumps(pageindex_tree) if pageindex_tree else None,
             summary)
        )
        doc_id = cur.lastrowid

        # Chunk pages and embed
        chunks = []
        for i, page_text in enumerate(pages):
            if page_text.strip():
                chunks.append({
                    "page_num": i + 1,
                    "content": page_text.strip(),
                    "section_title": f"Page {i + 1}",
                })

        if chunks:
            texts = [c["content"] for c in chunks]
            embeddings = self._embed(texts)

            for chunk, emb in zip(chunks, embeddings):
                import struct
                emb_blob = struct.pack(f"{len(emb)}f", *emb)
                self.conn.execute(
                    """INSERT INTO document_chunks
                       (document_id, page_num, section_title, content, embedding)
                       VALUES (?, ?, ?, ?, ?)""",
                    (doc_id, chunk["page_num"], chunk["section_title"],
                     chunk["content"], emb_blob)
                )

        self.conn.commit()
        return doc_id

    def search(self, query: str, top_k: int = 5,
               doc_id: Optional[int] = None) -> list[dict]:
        """Vector similarity search over document chunks."""
        import struct

        query_emb = self._embed([query])[0]

        where = "WHERE embedding IS NOT NULL"
        params = []
        if doc_id:
            where += " AND document_id = ?"
            params.append(doc_id)

        rows = self.conn.execute(
            f"""SELECT id, document_id, page_num, section_title, content,
                       embedding
                FROM document_chunks {where}""",
            params
        ).fetchall()

        results = []
        for row in rows:
            emb_blob = row["embedding"]
            dim = len(emb_blob) // 4
            emb = list(struct.unpack(f"{dim}f", emb_blob))
            score = self._cosine_similarity(query_emb, emb)
            results.append({
                "chunk_id": row["id"],
                "document_id": row["document_id"],
                "page_num": row["page_num"],
                "section": row["section_title"],
                "content": row["content"],
                "score": score,
                "source": "vector",
            })

        results.sort(key=lambda x: x["score"], reverse=True)
        return results[:top_k]

    def list_documents(self) -> list[dict]:
        """List all indexed documents."""
        rows = self.conn.execute(
            """SELECT id, filename, page_count, doc_summary, created_at
               FROM documents ORDER BY created_at DESC"""
        ).fetchall()
        return [dict(r) for r in rows]

    def get_document(self, doc_id: int) -> Optional[dict]:
        """Get document metadata including PageIndex tree."""
        row = self.conn.execute(
            "SELECT * FROM documents WHERE id = ?", (doc_id,)
        ).fetchone()
        if row:
            d = dict(row)
            if d.get("pageindex_tree"):
                d["pageindex_tree"] = json.loads(d["pageindex_tree"])
            return d
        return None

    def delete_document(self, doc_id: int):
        """Delete a document and its chunks."""
        self.conn.execute("DELETE FROM document_chunks WHERE document_id = ?",
                          (doc_id,))
        self.conn.execute("DELETE FROM documents WHERE id = ?", (doc_id,))
        self.conn.commit()
