"""
SQLiteDocumentRepository for encrypted chunk storage.

Design decisions:
- Fernet encryption/decryption happens in this repository layer.
- Thread safety: a per-instance Lock guards every sqlite3 call. Connections are opened
  per-operation (safest pattern for multi-threaded use without a connection pool).
- Schema uses INSERT OR REPLACE (UPSERT) so ingest remains idempotent.
"""

import json
import logging
import sqlite3
import threading

from src.core.exceptions import InfrastructureError

logger = logging.getLogger(__name__)

_SCHEMA = """
CREATE TABLE IF NOT EXISTS document_chunks (
    chunk_id       TEXT PRIMARY KEY,
    encrypted_text TEXT NOT NULL,
    source         TEXT DEFAULT 'unknown',
    chunk_index    INTEGER DEFAULT 0,
    extra_metadata TEXT DEFAULT '{}'
)
"""


class SQLiteDocumentRepository:
    def __init__(self, db_path: str, encryption_manager):
        self.db_path = db_path
        self.encryption_manager = encryption_manager
        self._lock = threading.Lock()
        self._init_db()

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _init_db(self) -> None:
        try:
            with self._lock, sqlite3.connect(self.db_path) as conn:
                conn.execute(_SCHEMA)
                conn.commit()
            logger.info("SQLite document store ready at '%s'.", self.db_path)
        except Exception as e:
            raise InfrastructureError(f"SQLite initialization failed: {e}") from e

    def _conn(self) -> sqlite3.Connection:
        conn = sqlite3.connect(self.db_path)
        conn.row_factory = sqlite3.Row
        return conn

    # ------------------------------------------------------------------
    # Public interface (mirrors DocumentRepository)
    # ------------------------------------------------------------------

    def save_document_chunk(self, chunk_id: str, original_text: str, metadata: dict) -> bool:
        """Encrypt and upsert chunk payload into SQLite."""
        try:
            encrypted_text = self.encryption_manager.encrypt_data(original_text)
            source = metadata.get("source", "unknown")
            chunk_index = metadata.get("chunk_index", 0)
            extra = {
                k: v
                for k, v in metadata.items()
                if k not in {
                    "chunk_id", "source", "chunk_index",
                    "encrypted_text", "decrypted_text", "text",
                }
                and v is not None
            }
            with self._lock, self._conn() as conn:
                conn.execute(
                    """
                    INSERT INTO document_chunks
                        (chunk_id, encrypted_text, source, chunk_index, extra_metadata)
                    VALUES (?, ?, ?, ?, ?)
                    ON CONFLICT(chunk_id) DO UPDATE SET
                        encrypted_text = excluded.encrypted_text,
                        source         = excluded.source,
                        chunk_index    = excluded.chunk_index,
                        extra_metadata = excluded.extra_metadata
                    """,
                    (chunk_id, encrypted_text, source, chunk_index, json.dumps(extra)),
                )
                conn.commit()
            logger.info("Chunk '%s' encrypted and saved.", chunk_id)
            return True
        except Exception as e:
            logger.error("Chunk save failed (%s): %s", chunk_id, e)
            return False

    def get_document_chunks(self, chunk_ids: list[str]) -> list[dict]:
        if not chunk_ids:
            return []

        placeholders = ",".join("?" for _ in chunk_ids)
        try:
            with self._lock, self._conn() as conn:
                rows = conn.execute(
                    f"SELECT * FROM document_chunks WHERE chunk_id IN ({placeholders})",
                    chunk_ids,
                ).fetchall()

            by_id = {}
            for row in rows:
                result = dict(row)
                result["decrypted_text"] = self.encryption_manager.decrypt_data(
                    result.pop("encrypted_text")
                )
                extra = json.loads(result.pop("extra_metadata", "{}"))
                result.update(extra)
                by_id[result["chunk_id"]] = result

            return [by_id[cid] for cid in chunk_ids if cid in by_id]
        except Exception as e:
            logger.error("Batch chunk read failed: %s", e)
            raise InfrastructureError(f"SQLite batch read failed: {e}") from e

    def has_document_chunk(self, chunk_id: str) -> bool:
        try:
            with self._lock, self._conn() as conn:
                row = conn.execute(
                    "SELECT 1 FROM document_chunks WHERE chunk_id = ? LIMIT 1", (chunk_id,)
                ).fetchone()
            return row is not None
        except Exception as e:
            logger.error("Chunk existence check failed (%s): %s", chunk_id, e)
            return False

    def delete_document_chunk(self, chunk_id: str) -> None:
        try:
            with self._lock, self._conn() as conn:
                conn.execute(
                    "DELETE FROM document_chunks WHERE chunk_id = ?", (chunk_id,)
                )
                conn.commit()
        except Exception as e:
            raise InfrastructureError(f"SQLite delete failed for {chunk_id}: {e}") from e

    def delete_chunk_ids(self, chunk_ids: list[str]) -> None:
        if not chunk_ids:
            return
        placeholders = ",".join("?" for _ in chunk_ids)
        try:
            with self._lock, self._conn() as conn:
                conn.execute(
                    f"DELETE FROM document_chunks WHERE chunk_id IN ({placeholders})",
                    chunk_ids,
                )
                conn.commit()
        except Exception as e:
            raise InfrastructureError(f"SQLite bulk delete failed: {e}") from e

    def delete_all_chunks(self) -> None:
        try:
            with self._lock, self._conn() as conn:
                conn.execute("DELETE FROM document_chunks")
                conn.commit()
        except Exception as e:
            raise InfrastructureError(f"SQLite full reset failed: {e}") from e

    def list_chunk_ids(self) -> set[str]:
        try:
            with self._lock, self._conn() as conn:
                rows = conn.execute(
                    "SELECT chunk_id FROM document_chunks"
                ).fetchall()
            return {row[0] for row in rows}
        except Exception as e:
            raise InfrastructureError(f"SQLite list IDs failed: {e}") from e

    def ping(self) -> bool:
        """Lightweight liveness check — verifies the DB file is accessible and schema exists."""
        try:
            with self._lock, self._conn() as conn:
                conn.execute("SELECT 1 FROM document_chunks LIMIT 1")
            return True
        except Exception as e:
            logger.error("SQLite ping failed: %s", e)
            return False