from __future__ import annotations

import json
import os
import sqlite3
import sys
from pathlib import Path
from typing import Any

from rich import print


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

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


from scripts.common import (
    CHUNKS_FILE,
    DB_FILE,
    DOCUMENTS_FILE,
    read_json,
)
from scripts.embedding_utils import (
    build_embedding_text,
    embed_passages,
    embedding_model_name,
    vector_to_blob,
)


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 PRIMARY KEY,
            model TEXT NOT NULL,
            dimensions INTEGER NOT NULL,
            embedding BLOB NOT NULL,
            FOREIGN KEY(chunk_id)
                REFERENCES chunks(chunk_id)
                ON UPDATE CASCADE
                ON DELETE CASCADE
        );

        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],
) -> None:
    if not chunks:
        return

    chunk_ids: list[str] = []
    texts: list[str] = []

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

        if not chunk_id:
            raise ValueError(
                "Chunk bez chunk_id nemôže "
                "mať embedding"
            )

        chunk_ids.append(
            chunk_id
        )

        texts.append(
            build_embedding_text(
                chunk
            )
        )

    vectors = embed_passages(
        texts
    )

    if len(vectors) != len(chunk_ids):
        raise RuntimeError(
            "Počet embeddingov sa "
            "nezhoduje s počtom chunkov"
        )

    model_name = (
        embedding_model_name()
    )

    rows: list[tuple] = []

    for chunk_id, vector in zip(
        chunk_ids,
        vectors,
    ):
        rows.append(
            (
                chunk_id,
                model_name,
                int(vector.shape[0]),
                vector_to_blob(vector),
            )
        )

    conn.executemany(
        """
        INSERT INTO chunk_embeddings (
            chunk_id,
            model,
            dimensions,
            embedding
        )
        VALUES (?, ?, ?, ?)
        """,
        rows,
    )


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(*)
        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 {
        "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(*)
            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() -> dict[str, int]:
    documents = read_json(
        DOCUMENTS_FILE
    )

    chunks = read_json(
        CHUNKS_FILE
    )

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

    temporary_file = (
        temporary_database_path(
            DB_FILE
        )
    )

    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
                )

                insert_embeddings(
                    conn,
                    chunks,
                )

            validate_database(
                conn
            )

            counts = get_counts(
                conn
            )

        os.replace(
            temporary_file,
            DB_FILE,
        )

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

        raise

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

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