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()