"""Bounded, local SQLite FTS index for user-selected workspace documents."""

from __future__ import annotations

import re
import sqlite3
import hashlib
import json
import math
import struct
from contextlib import contextmanager
from pathlib import Path
from typing import Any, Iterator

from app import database
from app import workspace

CHUNK_SIZE = 1000
CHUNK_OVERLAP = 120
MAX_CHUNKS_PER_DOCUMENT = 4
MAX_SEARCH_RESULTS = 8
MAX_FILES_PER_INDEX = 20
MAX_TOTAL_CHARS = 2_000_000
SEARCH_STOP_WORDS = {
    "ابحث", "البحث", "بحث", "في", "عن", "فهرس", "الفهرس", "المعرفة", "المحلي",
    "المحلية", "قاعدة", "المرحلة", "مرحلة", "ثم", "لخص", "استشهد", "مسار", "الملف",
    "أي", "ما", "كيف", "أين", "وين", "خدمة", "الذي", "التي", "إلى", "من", "على",
    "the", "and", "for", "from", "with", "search", "knowledge", "index", "indexed",
}
QUERY_SYNONYMS = {
    "تفرغ": ("تفريغ", "التفريغ", "للتفريغ", "transcription"),
    "يفرغ": ("تفريغ", "التفريغ", "للتفريغ", "transcription"),
    "التسجيلات": ("التسجيل", "التسجيلات"),
    "الصوتية": ("الصوت", "الصوتية", "audio"),
    "permission": ("allowed_tools",),
    "permissions": ("allowed_tools",),
    "صلاحية": ("allowed_tools", "مسموح"),
    "صلاحيات": ("allowed_tools", "مسموح"),
}


@contextmanager
def _connect() -> Iterator[sqlite3.Connection]:
    connection = sqlite3.connect(database.DATABASE_PATH, timeout=10)
    try:
        with connection:
            yield connection
    finally:
        connection.close()


def initialize() -> None:
    with _connect() as connection:
        connection.execute("PRAGMA foreign_keys = ON")
        connection.executescript(
            """
            CREATE TABLE IF NOT EXISTS knowledge_documents (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                user_id TEXT NOT NULL,
                workspace_path TEXT NOT NULL,
                relative_path TEXT NOT NULL,
                content_hash TEXT NOT NULL,
                updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP,
                UNIQUE(user_id, workspace_path, relative_path)
            );
            CREATE TABLE IF NOT EXISTS knowledge_chunks (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                document_id INTEGER NOT NULL REFERENCES knowledge_documents(id) ON DELETE CASCADE,
                chunk_index INTEGER NOT NULL,
                body TEXT NOT NULL,
                UNIQUE(document_id, chunk_index)
            );
            CREATE VIRTUAL TABLE IF NOT EXISTS knowledge_fts USING fts5(
                body, tokenize = 'unicode61'
            );
            CREATE TABLE IF NOT EXISTS knowledge_embeddings (
                chunk_id INTEGER PRIMARY KEY REFERENCES knowledge_chunks(id) ON DELETE CASCADE,
                model TEXT NOT NULL,
                dimensions INTEGER NOT NULL,
                vector BLOB NOT NULL
            );
            """
        )


def _chunks(text: str) -> list[str]:
    cleaned = text.strip()
    if not cleaned:
        return []
    step = CHUNK_SIZE - CHUNK_OVERLAP
    result = [cleaned[start : start + CHUNK_SIZE] for start in range(0, len(cleaned), step)]
    return [chunk for chunk in result if chunk.strip()]


def index_document(
    *, user_id: str, workspace_path: Path, relative_path: str, text: str, content_hash: str
) -> dict[str, Any]:
    chunks = _chunks(text)
    if not chunks:
        raise ValueError("الملف لا يحتوي على نص قابل للفهرسة.")
    if sum(map(len, chunks)) > MAX_TOTAL_CHARS + CHUNK_SIZE:
        raise ValueError("نص الملف يتجاوز حد الفهرسة.")

    root = str(workspace_path.resolve(strict=True))
    with _connect() as connection:
        connection.execute("PRAGMA foreign_keys = ON")
        existing = connection.execute(
            "SELECT id FROM knowledge_documents WHERE user_id=? AND workspace_path=? AND relative_path=?",
            (user_id, root, relative_path),
        ).fetchone()
        if existing:
            document_id = int(existing[0])
            rows = connection.execute(
                "SELECT id FROM knowledge_chunks WHERE document_id=?", (document_id,)
            ).fetchall()
            connection.executemany("DELETE FROM knowledge_fts WHERE rowid=?", rows)
            connection.executemany("DELETE FROM knowledge_embeddings WHERE chunk_id=?", rows)
            connection.execute("DELETE FROM knowledge_chunks WHERE document_id=?", (document_id,))
            connection.execute(
                "UPDATE knowledge_documents SET content_hash=?, updated_at=CURRENT_TIMESTAMP WHERE id=?",
                (content_hash, document_id),
            )
        else:
            cursor = connection.execute(
                "INSERT INTO knowledge_documents(user_id,workspace_path,relative_path,content_hash) VALUES(?,?,?,?)",
                (user_id, root, relative_path, content_hash),
            )
            document_id = int(cursor.lastrowid)

        for index, body in enumerate(chunks):
            cursor = connection.execute(
                "INSERT INTO knowledge_chunks(document_id,chunk_index,body) VALUES(?,?,?)",
                (document_id, index, body),
            )
            connection.execute(
                "INSERT INTO knowledge_fts(rowid,body) VALUES(?,?)",
                (int(cursor.lastrowid), body),
            )
    return {"path": relative_path, "chunks": len(chunks), "characters": len(text)}


def store_embeddings(
    *, user_id: str, workspace_path: Path, relative_path: str, model: str, vectors: list[list[float]]
) -> int:
    """Replace a document's vectors, requiring one finite vector per indexed chunk."""
    with _connect() as connection:
        connection.row_factory = sqlite3.Row
        document = connection.execute(
            "SELECT id FROM knowledge_documents WHERE user_id=? AND workspace_path=? AND relative_path=?",
            (user_id, str(workspace_path.resolve()), relative_path),
        ).fetchone()
        if document is None:
            return 0
        chunks = connection.execute(
            "SELECT id FROM knowledge_chunks WHERE document_id=? ORDER BY chunk_index",
            (document["id"],),
        ).fetchall()
        if len(chunks) != len(vectors):
            raise ValueError("يجب أن يقابل كل مقطع مفهرس متجه تضمين واحد.")
        prepared: list[tuple[int, str, int, bytes]] = []
        dimensions: int | None = None
        for chunk, vector in zip(chunks, vectors, strict=True):
            if not vector or (dimensions is not None and len(vector) != dimensions):
                raise ValueError("أبعاد متجهات التضمين غير صالحة أو غير متطابقة.")
            dimensions = len(vector)
            numeric = [float(value) for value in vector]
            if not all(math.isfinite(value) for value in numeric) or not any(numeric):
                raise ValueError("متجه التضمين يحتوي قيمة غير صالحة.")
            packed = struct.pack(f"<{dimensions}f", *numeric)
            prepared.append((int(chunk["id"]), model, dimensions, packed))
        connection.execute(
            "DELETE FROM knowledge_embeddings WHERE chunk_id IN "
            "(SELECT id FROM knowledge_chunks WHERE document_id=?)",
            (document["id"],),
        )
        connection.executemany(
            "INSERT INTO knowledge_embeddings(chunk_id,model,dimensions,vector) VALUES(?,?,?,?)",
            prepared,
        )
    return len(prepared)


def has_embeddings(*, user_id: str, workspace_path: Path, model: str) -> bool:
    with _connect() as connection:
        row = connection.execute(
            """SELECT 1 FROM knowledge_embeddings AS e
               JOIN knowledge_chunks AS c ON c.id=e.chunk_id
               JOIN knowledge_documents AS d ON d.id=c.document_id
               WHERE d.user_id=? AND d.workspace_path=? AND e.model=? LIMIT 1""",
            (user_id, str(workspace_path.resolve()), model),
        ).fetchone()
    return row is not None


def search_by_embedding(
    vector: list[float], *, user_id: str, workspace_path: Path, model: str,
    limit: int = MAX_SEARCH_RESULTS,
) -> list[dict[str, Any]]:
    if not vector or not all(math.isfinite(float(value)) for value in vector):
        return []
    query_norm = math.sqrt(sum(float(value) ** 2 for value in vector))
    if query_norm == 0:
        return []
    with _connect() as connection:
        connection.row_factory = sqlite3.Row
        rows = connection.execute(
            """SELECT d.id AS document_id, d.relative_path, d.content_hash,
                      c.chunk_index, c.body, e.dimensions, e.vector
               FROM knowledge_embeddings AS e
               JOIN knowledge_chunks AS c ON c.id=e.chunk_id
               JOIN knowledge_documents AS d ON d.id=c.document_id
               WHERE d.user_id=? AND d.workspace_path=? AND e.model=?""",
            (user_id, str(workspace_path.resolve()), model),
        ).fetchall()
    ranked: list[tuple[float, sqlite3.Row]] = []
    for row in rows:
        dimensions = int(row["dimensions"])
        if dimensions != len(vector) or len(row["vector"]) != dimensions * 4:
            continue
        values = struct.unpack(f"<{dimensions}f", row["vector"])
        norm = math.sqrt(sum(value * value for value in values))
        if norm == 0:
            continue
        similarity = sum(float(left) * right for left, right in zip(vector, values, strict=True)) / (query_norm * norm)
        ranked.append((similarity, row))
    ranked.sort(key=lambda item: (-item[0], item[1]["relative_path"], item[1]["chunk_index"]))
    bounded_limit = max(1, min(limit, MAX_SEARCH_RESULTS))
    selected: list[tuple[float, sqlite3.Row]] = []
    per_document: dict[int, int] = {}
    for score, row in ranked:
        document_id = int(row["document_id"])
        if per_document.get(document_id, 0) >= MAX_CHUNKS_PER_DOCUMENT:
            continue
        per_document[document_id] = per_document.get(document_id, 0) + 1
        selected.append((score, row))
        if len(selected) >= bounded_limit:
            break
    valid: dict[int, bool] = {}
    results: list[dict[str, Any]] = []
    for score, row in selected:
        document_id = int(row["document_id"])
        if document_id not in valid:
            try:
                path = workspace.relative_knowledge_file(workspace_path, str(row["relative_path"]))
                valid[document_id] = hashlib.sha256(path.read_bytes()).hexdigest() == row["content_hash"]
            except (OSError, ValueError):
                valid[document_id] = False
            if not valid[document_id]:
                delete_document(
                    user_id=user_id,
                    workspace_path=workspace_path,
                    relative_path=str(row["relative_path"]),
                )
        if valid[document_id]:
            body = str(row["body"])
            results.append(
                {
                    "path": row["relative_path"],
                    "chunk": row["chunk_index"],
                    "text": body,
                    "excerpt": body[:220],
                    "similarity": round(score, 4),
                }
            )
    return results


def merge_search_results(
    lexical: list[dict[str, Any]], semantic: list[dict[str, Any]],
    *, limit: int = MAX_SEARCH_RESULTS,
) -> list[dict[str, Any]]:
    """Combine lexical and semantic ranks with reciprocal-rank fusion."""
    combined: dict[tuple[str, int], dict[str, Any]] = {}
    scores: dict[tuple[str, int], float] = {}
    for rank, result in enumerate(lexical, 1):
        key = (str(result["path"]), int(result["chunk"]))
        combined[key] = dict(result)
        scores[key] = scores.get(key, 0.0) + 1 / (60 + rank)
    for rank, result in enumerate(semantic, 1):
        key = (str(result["path"]), int(result["chunk"]))
        if key not in combined:
            combined[key] = dict(result)
        else:
            combined[key]["similarity"] = result.get("similarity")
        scores[key] = scores.get(key, 0.0) + 1 / (60 + rank)
    ordered = sorted(combined, key=lambda key: (-scores[key], key[0], key[1]))
    return [combined[key] for key in ordered[: max(1, min(limit, MAX_SEARCH_RESULTS))]]


def delete_document(*, user_id: str, workspace_path: Path, relative_path: str) -> bool:
    root = str(workspace_path.resolve())
    with _connect() as connection:
        connection.execute("PRAGMA foreign_keys = ON")
        row = connection.execute(
            "SELECT id FROM knowledge_documents WHERE user_id=? AND workspace_path=? AND relative_path=?",
            (user_id, root, relative_path),
        ).fetchone()
        if row is None:
            return False
        document_id = int(row[0])
        chunks = connection.execute(
            "SELECT id FROM knowledge_chunks WHERE document_id=?", (document_id,)
        ).fetchall()
        connection.executemany("DELETE FROM knowledge_fts WHERE rowid=?", chunks)
        connection.executemany("DELETE FROM knowledge_embeddings WHERE chunk_id=?", chunks)
        connection.execute("DELETE FROM knowledge_documents WHERE id=?", (document_id,))
    return True


def search(
    query: str,
    *,
    user_id: str,
    workspace_path: Path,
    limit: int = MAX_SEARCH_RESULTS,
) -> list[dict[str, Any]]:
    base_terms = [
        term
        for term in dict.fromkeys(re.findall(r"[\w\u0600-\u06ff]{2,}", query.casefold()))
        if term not in SEARCH_STOP_WORDS
    ]
    terms = list(
        dict.fromkeys(
            term
            for base in base_terms
            for term in (base, *QUERY_SYNONYMS.get(base, ()))
        )
    )[:20]
    if not terms:
        return []
    match_query = " OR ".join('"' + term.replace('"', '""') + '"' for term in terms)
    bounded_limit = max(1, min(limit, MAX_SEARCH_RESULTS))
    with _connect() as connection:
        connection.row_factory = sqlite3.Row
        rows = connection.execute(
            """
            SELECT d.id AS document_id, d.relative_path, d.content_hash,
                   c.chunk_index, c.body,
                   snippet(knowledge_fts, 0, '[', ']', ' … ', 24) AS excerpt,
                   bm25(knowledge_fts) AS rank
            FROM knowledge_fts
            JOIN knowledge_chunks AS c ON c.id=knowledge_fts.rowid
            JOIN knowledge_documents AS d ON d.id=c.document_id
            WHERE knowledge_fts MATCH ? AND d.user_id=? AND d.workspace_path=?
            ORDER BY rank, d.relative_path, c.chunk_index
            LIMIT ?
            """,
            (match_query, user_id, str(workspace_path.resolve()), bounded_limit * 8),
        ).fetchall()
    selected_rows: list[sqlite3.Row] = []
    chunks_per_document: dict[int, int] = {}
    for row in rows:
        document_id = int(row["document_id"])
        if chunks_per_document.get(document_id, 0) >= MAX_CHUNKS_PER_DOCUMENT:
            continue
        chunks_per_document[document_id] = chunks_per_document.get(document_id, 0) + 1
        selected_rows.append(row)
        if len(selected_rows) >= bounded_limit:
            break
    current_documents: dict[int, bool] = {}
    for row in selected_rows:
        document_id = int(row["document_id"])
        if document_id in current_documents:
            continue
        relative_path = str(row["relative_path"])
        try:
            source = workspace.relative_knowledge_file(workspace_path, relative_path)
            actual_hash = hashlib.sha256(source.read_bytes()).hexdigest()
            current_documents[document_id] = actual_hash == row["content_hash"]
        except (OSError, ValueError):
            current_documents[document_id] = False
        if not current_documents[document_id]:
            delete_document(
                user_id=user_id,
                workspace_path=workspace_path,
                relative_path=relative_path,
            )
    return [
        {"path": row["relative_path"], "chunk": row["chunk_index"], "text": row["body"], "excerpt": row["excerpt"]}
        for row in selected_rows
        if current_documents.get(int(row["document_id"]), False)
    ]


initialize()