import asyncio
import json
import sqlite3
import uuid
from pathlib import Path

from axon.config import settings
from axon.models import IndexDocument, IngestDocument, SearchResult

_lock = asyncio.Lock()


def _db_path() -> Path:
    return Path(settings.index_db_path)


def _connect() -> sqlite3.Connection:
    conn = sqlite3.connect(_db_path())
    conn.row_factory = sqlite3.Row
    return conn


async def init_index() -> None:
    def _init() -> None:
        with _connect() as conn:
            conn.execute(
                """
                CREATE TABLE IF NOT EXISTS documents (
                    id TEXT PRIMARY KEY,
                    title TEXT NOT NULL,
                    content TEXT NOT NULL,
                    url TEXT,
                    metadata TEXT NOT NULL DEFAULT '{}',
                    updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP
                )
                """
            )
            conn.execute(
                """
                CREATE VIRTUAL TABLE IF NOT EXISTS documents_fts USING fts5(
                    title,
                    content,
                    content='documents',
                    content_rowid='rowid'
                )
                """
            )
            conn.execute(
                """
                CREATE TRIGGER IF NOT EXISTS documents_ai AFTER INSERT ON documents BEGIN
                    INSERT INTO documents_fts(rowid, title, content)
                    VALUES (new.rowid, new.title, new.content);
                END
                """
            )
            conn.execute(
                """
                CREATE TRIGGER IF NOT EXISTS documents_ad AFTER DELETE ON documents BEGIN
                    INSERT INTO documents_fts(documents_fts, rowid, title, content)
                    VALUES ('delete', old.rowid, old.title, old.content);
                END
                """
            )
            conn.execute(
                """
                CREATE TRIGGER IF NOT EXISTS documents_au AFTER UPDATE ON documents BEGIN
                    INSERT INTO documents_fts(documents_fts, rowid, title, content)
                    VALUES ('delete', old.rowid, old.title, old.content);
                    INSERT INTO documents_fts(rowid, title, content)
                    VALUES (new.rowid, new.title, new.content);
                END
                """
            )

    async with _lock:
        await asyncio.to_thread(_init)


async def ingest_documents(documents: list[IngestDocument]) -> int:
    await init_index()

    def _ingest() -> int:
        with _connect() as conn:
            for doc in documents:
                doc_id = doc.id or str(uuid.uuid4())
                conn.execute(
                    """
                    INSERT INTO documents(id, title, content, url, metadata, updated_at)
                    VALUES (?, ?, ?, ?, ?, CURRENT_TIMESTAMP)
                    ON CONFLICT(id) DO UPDATE SET
                        title=excluded.title,
                        content=excluded.content,
                        url=excluded.url,
                        metadata=excluded.metadata,
                        updated_at=CURRENT_TIMESTAMP
                    """,
                    (
                        doc_id,
                        doc.title,
                        doc.content,
                        doc.url,
                        json.dumps(doc.metadata, sort_keys=True),
                    ),
                )
            return len(documents)

    async with _lock:
        return await asyncio.to_thread(_ingest)


async def search_index(query: str, max_results: int = 10) -> list[IndexDocument]:
    await init_index()

    def _search() -> list[IndexDocument]:
        with _connect() as conn:
            rows = conn.execute(
                """
                SELECT
                    documents.id,
                    documents.title,
                    snippet(documents_fts, 1, '', '', '...', 40) AS content,
                    documents.url,
                    bm25(documents_fts) AS rank
                FROM documents_fts
                JOIN documents ON documents_fts.rowid = documents.rowid
                WHERE documents_fts MATCH ?
                ORDER BY rank
                LIMIT ?
                """,
                (query, max_results),
            ).fetchall()
        results = []
        for row in rows:
            rank = float(row["rank"])
            score = max(0.0, min(1.0, 1.0 / (1.0 + abs(rank))))
            results.append(
                IndexDocument(
                    id=row["id"],
                    title=row["title"],
                    content=row["content"] or "",
                    url=row["url"],
                    score=round(score, 6),
                )
            )
        return results

    async with _lock:
        return await asyncio.to_thread(_search)


async def search_index_as_results(query: str, max_results: int = 5) -> list[SearchResult]:
    docs = await search_index(query, max_results=max_results)
    return [
        SearchResult(
            title=doc.title,
            url=doc.url or f"private://{doc.id}",
            content=doc.content,
            score=doc.score,
        )
        for doc in docs
    ]