from __future__ import annotations

import json
import os
import sqlite3
import numpy as np
import sys
from pathlib import Path
from typing import Any

from rich import print


PROJECT_ROOT = Path(__file__).resolve().parents[2]

if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))


from tools.rag.common import (
    CHUNKS_FILE,
    DB_FILE,
    DOCUMENTS_FILE,
    read_json,
)
from tools.rag.embedding_utils import (
    build_embedding_text,
    embed_passages,
    embedding_model_name,
    vector_to_blob,
    model_profile,
    split_passage,
)
from tools.rag.embedding_cache import EmbeddingCache, cache_key


FTS_TOKENIZER = "unicode61 remove_diacritics 2"
FTS_PREFIXES = "3 4 5"


def published_to_db(value: Any) -> int | None:
    if value is True:
        return 1

    if value is False:
        return 0

    return None


def verify_fts5(
    conn: sqlite3.Connection,
) -> None:
    try:
        conn.execute(
            "CREATE VIRTUAL TABLE temp.fts5_check "
            "USING fts5(value)"
        )
        conn.execute(
            "DROP TABLE temp.fts5_check"
        )

    except sqlite3.OperationalError as error:
        raise RuntimeError(
            "Táto inštalácia SQLite nemá dostupné FTS5. "
            "Použi Python/SQLite zostavenie s podporou "
            "SQLITE_ENABLE_FTS5."
        ) from error


def create_tables(
    conn: sqlite3.Connection,
) -> None:
    conn.executescript(
        f"""
        PRAGMA foreign_keys = ON;

        CREATE TABLE documents (
            id INTEGER PRIMARY KEY,
            path TEXT UNIQUE NOT NULL,
            title TEXT,
            author TEXT,
            published INTEGER
                CHECK (
                    published IN (0, 1)
                    OR published IS NULL
                ),
            content_length INTEGER
                NOT NULL DEFAULT 0,
            metadata_json TEXT
                NOT NULL DEFAULT '{{}}'
        );

        CREATE TABLE chunks (
            id INTEGER PRIMARY KEY,
            chunk_id TEXT UNIQUE NOT NULL,
            document_path TEXT NOT NULL,
            title TEXT,
            author TEXT,
            published INTEGER
                CHECK (
                    published IN (0, 1)
                    OR published IS NULL
                ),
            chunk_index INTEGER NOT NULL,
            heading_paths_json TEXT
                NOT NULL DEFAULT '[]',
            text TEXT NOT NULL,
            text_length INTEGER
                NOT NULL DEFAULT 0,
            token_count INTEGER,
            content_hash TEXT,
            FOREIGN KEY(document_path)
                REFERENCES documents(path)
                ON UPDATE CASCADE
                ON DELETE CASCADE
        );

        CREATE TABLE chunk_tags (
            chunk_id TEXT NOT NULL,
            tag TEXT NOT NULL,
            PRIMARY KEY(chunk_id, tag),
            FOREIGN KEY(chunk_id)
                REFERENCES chunks(chunk_id)
                ON UPDATE CASCADE
                ON DELETE CASCADE
        );

        CREATE TABLE chunk_categories (
            chunk_id TEXT NOT NULL,
            category TEXT NOT NULL,
            PRIMARY KEY(chunk_id, category),
            FOREIGN KEY(chunk_id)
                REFERENCES chunks(chunk_id)
                ON UPDATE CASCADE
                ON DELETE CASCADE
        );

        CREATE TABLE chunk_embeddings (
            chunk_id TEXT NOT NULL,
            segment_index INTEGER NOT NULL DEFAULT 0,
            model TEXT NOT NULL,
            dimensions INTEGER NOT NULL,
            embedding BLOB NOT NULL,
            PRIMARY KEY(chunk_id, segment_index),
            FOREIGN KEY(chunk_id)
                REFERENCES chunks(chunk_id)
                ON UPDATE CASCADE
                ON DELETE CASCADE
        );

        CREATE TABLE embedding_model_info (metadata_json TEXT NOT NULL);

        CREATE INDEX idx_documents_path
            ON documents(path);

        CREATE INDEX idx_documents_published
            ON documents(published);

        CREATE INDEX idx_chunks_document_path
            ON chunks(document_path);

        CREATE INDEX idx_chunks_title
            ON chunks(title);

        CREATE INDEX idx_chunks_author
            ON chunks(author);

        CREATE INDEX idx_chunks_published
            ON chunks(published);

        CREATE INDEX idx_chunk_tags_tag
            ON chunk_tags(tag);

        CREATE INDEX idx_chunk_categories_category
            ON chunk_categories(category);

        CREATE INDEX idx_chunk_embeddings_model
            ON chunk_embeddings(model);

        CREATE VIRTUAL TABLE chunks_fts USING fts5(
            chunk_id UNINDEXED,
            title,
            author,
            document_path,
            tags,
            categories,
            text,
            tokenize='{FTS_TOKENIZER}',
            prefix='{FTS_PREFIXES}'
        );
        """
    )


def insert_documents(
    conn: sqlite3.Connection,
    documents: list[dict],
) -> None:
    rows = [
        (
            document.get("path"),
            document.get("title"),
            document.get("author"),
            published_to_db(
                document.get("published")
            ),
            int(
                document.get("content_length")
                or 0
            ),
            json.dumps(
                document.get("metadata") or {},
                ensure_ascii=False,
                sort_keys=True,
            ),
        )
        for document in documents
    ]

    conn.executemany(
        """
        INSERT INTO documents (
            path,
            title,
            author,
            published,
            content_length,
            metadata_json
        )
        VALUES (?, ?, ?, ?, ?, ?)
        """,
        rows,
    )


def insert_chunks(
    conn: sqlite3.Connection,
    chunks: list[dict],
) -> None:
    chunk_rows: list[tuple] = []
    tag_rows: list[tuple[str, str]] = []
    category_rows: list[
        tuple[str, str]
    ] = []

    for chunk in chunks:
        chunk_id = str(
            chunk.get("chunk_id") or ""
        ).strip()

        if not chunk_id:
            raise ValueError(
                "Chunk bez chunk_id nie je "
                "možné indexovať"
            )

        text = chunk.get("text") or ""

        chunk_rows.append(
            (
                chunk_id,
                chunk.get("document_path"),
                chunk.get("title"),
                chunk.get("author"),
                published_to_db(
                    chunk.get("published")
                ),
                int(
                    chunk.get("chunk_index")
                    or 0
                ),
                json.dumps(
                    chunk.get(
                        "heading_paths"
                    )
                    or [],
                    ensure_ascii=False,
                ),
                text,
                int(
                    chunk.get("text_length")
                    or len(text)
                ),
                chunk.get("token_count"),
                chunk.get("content_hash"),
            )
        )

        for tag in chunk.get("tags") or []:
            value = str(tag).strip()

            if value:
                tag_rows.append(
                    (
                        chunk_id,
                        value,
                    )
                )

        for category in (
            chunk.get("categories") or []
        ):
            value = str(
                category
            ).strip()

            if value:
                category_rows.append(
                    (
                        chunk_id,
                        value,
                    )
                )

    conn.executemany(
        """
        INSERT INTO chunks (
            chunk_id,
            document_path,
            title,
            author,
            published,
            chunk_index,
            heading_paths_json,
            text,
            text_length,
            token_count,
            content_hash
        )
        VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
        """,
        chunk_rows,
    )

    conn.executemany(
        """
        INSERT OR IGNORE INTO chunk_tags (
            chunk_id,
            tag
        )
        VALUES (?, ?)
        """,
        tag_rows,
    )

    conn.executemany(
        """
        INSERT OR IGNORE INTO chunk_categories (
            chunk_id,
            category
        )
        VALUES (?, ?)
        """,
        category_rows,
    )


def insert_embeddings(conn: sqlite3.Connection, chunks: list[dict], *, cache_file: Path) -> dict:
    if not chunks:
        return {'embedding_reused': 0, 'embedding_computed': 0}
    model_name = embedding_model_name()
    profile = model_profile(model_name)
    conn.execute('INSERT INTO embedding_model_info VALUES(?)', (json.dumps(profile, sort_keys=True),))
    segments = []
    for chunk in chunks:
        if not chunk.get('chunk_id'):
            raise ValueError('Chunk bez chunk_id nemôže mať embedding')
        text = build_embedding_text(chunk)
        for index, part in enumerate(split_passage(text, model_name=model_name)):
            segments.append((chunk['chunk_id'], index, part, cache_key(profile, part)))
    cache = EmbeddingCache(cache_file)
    found = cache.load([row[3] for row in segments])
    missing = {row[3]: row[2] for row in segments if row[3] not in found}
    computed = {}
    if missing:
        vectors = embed_passages(list(missing.values()))
        if len(vectors) != len(missing):
            raise RuntimeError('Počet embeddingov sa nezhoduje s počtom vstupov')
        for key, vector in zip(missing, vectors):
            vector = np.asarray(vector, dtype=np.float32)
            if vector.ndim != 1 or not len(vector) or not np.isfinite(vector).all() or np.linalg.norm(vector) == 0:
                raise RuntimeError('Embedding nie je platný nenulový vektor')
            computed[key] = vector
        cache.save(computed)
    all_vectors = {**found, **computed}
    dimensions = {len(vector) for vector in all_vectors.values()}
    if len(dimensions) != 1:
        raise RuntimeError('Embeddingy majú rozdielne rozmery')
    conn.executemany('''INSERT INTO chunk_embeddings
        (chunk_id,segment_index,model,dimensions,embedding) VALUES(?,?,?,?,?)''',
        [(chunk_id, index, model_name, len(all_vectors[key]), vector_to_blob(all_vectors[key]))
         for chunk_id, index, _, key in segments])
    return {'embedding_reused': sum(row[3] in found for row in segments),
            'embedding_computed': len(computed)}


def build_fts_index(
    conn: sqlite3.Connection,
) -> None:
    conn.execute(
        """
        INSERT INTO chunks_fts (
            rowid,
            chunk_id,
            title,
            author,
            document_path,
            tags,
            categories,
            text
        )
        SELECT
            chunks.id,
            chunks.chunk_id,
            COALESCE(
                chunks.title,
                ''
            ),
            COALESCE(
                chunks.author,
                ''
            ),
            chunks.document_path,
            COALESCE(
                tags.values_text,
                ''
            ),
            COALESCE(
                categories.values_text,
                ''
            ),
            chunks.text
        FROM chunks
        LEFT JOIN (
            SELECT
                chunk_id,
                GROUP_CONCAT(
                    tag,
                    ' '
                ) AS values_text
            FROM chunk_tags
            GROUP BY chunk_id
        ) AS tags
            ON tags.chunk_id
            = chunks.chunk_id
        LEFT JOIN (
            SELECT
                chunk_id,
                GROUP_CONCAT(
                    category,
                    ' '
                ) AS values_text
            FROM chunk_categories
            GROUP BY chunk_id
        ) AS categories
            ON categories.chunk_id
            = chunks.chunk_id
        ORDER BY chunks.id
        """
    )

    conn.execute(
        "INSERT INTO chunks_fts"
        "(chunks_fts) "
        "VALUES('optimize')"
    )


def validate_database(
    conn: sqlite3.Connection,
) -> None:
    integrity = conn.execute(
        "PRAGMA integrity_check"
    ).fetchone()[0]

    if integrity != "ok":
        raise RuntimeError(
            "SQLite integrity check "
            f"zlyhal: {integrity}"
        )

    foreign_key_errors = (
        conn.execute(
            "PRAGMA foreign_key_check"
        ).fetchall()
    )

    if foreign_key_errors:
        raise RuntimeError(
            "Databáza obsahuje chyby "
            "cudzích kľúčov: "
            f"{foreign_key_errors[:5]}"
        )

    conn.execute(
        "INSERT INTO chunks_fts"
        "(chunks_fts) "
        "VALUES('integrity-check')"
    )

    chunk_count = conn.execute(
        "SELECT COUNT(*) FROM chunks"
    ).fetchone()[0]

    fts_count = conn.execute(
        "SELECT COUNT(*) FROM chunks_fts"
    ).fetchone()[0]

    embedding_count = conn.execute(
        """
        SELECT COUNT(DISTINCT chunk_id)
        FROM chunk_embeddings
        """
    ).fetchone()[0]

    if chunk_count != fts_count:
        raise RuntimeError(
            "Počet záznamov v chunks "
            "a chunks_fts sa nezhoduje: "
            f"{chunk_count} != {fts_count}"
        )

    if chunk_count != embedding_count:
        raise RuntimeError(
            "Počet chunkov a embeddingov "
            "sa nezhoduje: "
            f"{chunk_count} != "
            f"{embedding_count}"
        )

    model_count = conn.execute(
        """
        SELECT COUNT(
            DISTINCT model
        )
        FROM chunk_embeddings
        """
    ).fetchone()[0]

    if (
        embedding_count > 0
        and model_count != 1
    ):
        raise RuntimeError(
            "Embedding index obsahuje "
            "viac modelov"
        )


def get_counts(
    conn: sqlite3.Connection,
) -> dict[str, int]:
    return {
        "embedding_segments": conn.execute("SELECT COUNT(*) FROM chunk_embeddings").fetchone()[0],
        "documents": conn.execute(
            """
            SELECT COUNT(*)
            FROM documents
            """
        ).fetchone()[0],
        "chunks": conn.execute(
            """
            SELECT COUNT(*)
            FROM chunks
            """
        ).fetchone()[0],
        "fts_chunks": conn.execute(
            """
            SELECT COUNT(*)
            FROM chunks_fts
            """
        ).fetchone()[0],
        "embedding_chunks": conn.execute(
            """
            SELECT COUNT(DISTINCT chunk_id)
            FROM chunk_embeddings
            """
        ).fetchone()[0],
        "tags": conn.execute(
            """
            SELECT COUNT(*)
            FROM chunk_tags
            """
        ).fetchone()[0],
        "categories": conn.execute(
            """
            SELECT COUNT(*)
            FROM chunk_categories
            """
        ).fetchone()[0],
    }


def temporary_database_path(
    db_file: Path,
) -> Path:
    return db_file.with_name(
        f".{db_file.name}.tmp"
    )


def build_database(*, documents_file: Path | None = None, chunks_file: Path | None = None,
                   db_file: Path | None = None, cache_file: Path | None = None) -> dict[str, int]:
    target = db_file or DB_FILE
    documents = read_json(
        documents_file or DOCUMENTS_FILE
    )

    chunks = read_json(
        chunks_file or CHUNKS_FILE
    )

    target.parent.mkdir(
        parents=True,
        exist_ok=True,
    )

    temporary_file = (
        temporary_database_path(
            target
        )
    )

    if temporary_file.exists():
        temporary_file.unlink()

    try:
        with sqlite3.connect(
            temporary_file
        ) as conn:
            conn.execute(
                "PRAGMA foreign_keys = ON"
            )
            conn.execute(
                "PRAGMA temp_store = MEMORY"
            )

            verify_fts5(
                conn
            )

            with conn:
                create_tables(
                    conn
                )

                insert_documents(
                    conn,
                    documents,
                )

                insert_chunks(
                    conn,
                    chunks,
                )

                build_fts_index(
                    conn
                )

                embedding_stats = insert_embeddings(conn, chunks, cache_file=cache_file or target.parent / "embedding_cache.sqlite")

            validate_database(
                conn
            )

            counts = {**get_counts(conn), **embedding_stats}

        os.replace(
            temporary_file,
            target,
        )

    except Exception:
        if temporary_file.exists():
            temporary_file.unlink()

        raise

    print(
        "[green]SQLite index "
        "vytvorený:[/green] "
        f"{target}"
    )

    print(
        f"Dokumentov: "
        f"{counts['documents']}"
    )

    print(
        f"Chunkov: "
        f"{counts['chunks']}"
    )

    print(
        f"FTS5 chunkov: "
        f"{counts['fts_chunks']}"
    )

    print(
        f"Embedding chunkov: "
        f"{counts['embedding_chunks']}"
    )

    print(
        f"Tag záznamov: "
        f"{counts['tags']}"
    )

    print(
        "Kategória záznamov: "
        f"{counts['categories']}"
    )

    return counts


def main() -> None:
    build_database()


if __name__ == "__main__":
    main()