"""Hybrid search — FTS5 + vector with Reciprocal Rank Fusion."""

import json
import re
import struct
import sqlite3


def hybrid_search(conn: sqlite3.Connection, query: str, model_name: str, cfg: dict,
                  top: int = 10, tags: list[str] | None = None,
                  doc_type: str | None = None, fts_only: bool = False,
                  vec_only: bool = False, threshold: float | None = None) -> dict:
    """Run hybrid search and return merged results."""
    candidate_count = top * 3  # Fetch more candidates for RRF

    fts_results = {}
    vec_results = {}

    if not vec_only:
        fts_results = _fts_search(conn, query, candidate_count, tags, doc_type)

    if not fts_only:
        vec_results = _vector_search(conn, query, model_name, cfg, candidate_count, tags, doc_type)

    # Merge via RRF
    rrf_k = cfg.get("search", {}).get("rrf_k", 60)

    if fts_only:
        merged = _single_source_results(fts_results, "fts")
    elif vec_only:
        merged = _single_source_results(vec_results, "vector")
    else:
        merged = _rrf_merge(fts_results, vec_results, rrf_k)

    # Apply threshold
    if threshold is not None:
        merged = [r for r in merged if r["score"] >= threshold]

    # Sort and limit
    merged.sort(key=lambda x: x["score"], reverse=True)
    total = len(merged)
    merged = merged[:top]

    # Enrich with document metadata
    results = []
    for r in merged:
        chunk_id = r["chunk_id"]
        row = conn.execute("""
            SELECT c.id, c.text, c.chunk_index, c.metadata as chunk_meta,
                   d.id as doc_id, d.title, d.source_path, d.doc_type,
                   d.language, d.metadata as doc_meta
            FROM chunks c
            JOIN documents d ON c.document_id = d.id
            WHERE c.id = ?
        """, (chunk_id,)).fetchone()

        if not row:
            continue

        chunk_meta = json.loads(row["chunk_meta"]) if row["chunk_meta"] else {}

        # Get tags for this document
        tag_rows = conn.execute("""
            SELECT t.name FROM tags t
            JOIN document_tags dt ON t.id = dt.tag_id
            WHERE dt.document_id = ?
            ORDER BY t.name
        """, (row["doc_id"],)).fetchall()

        # Count total chunks for this document
        total_chunks = conn.execute(
            "SELECT COUNT(*) FROM chunks WHERE document_id = ?", (row["doc_id"],)
        ).fetchone()[0]

        results.append({
            "chunk_id": row["id"],
            "score": round(r["score"], 6),
            "score_breakdown": r["score_breakdown"],
            "text": row["text"],
            "source": {
                "document_id": row["doc_id"],
                "title": row["title"],
                "path": row["source_path"],
                "type": row["doc_type"],
                "page": chunk_meta.get("page"),
                "section_header": chunk_meta.get("section_header"),
                "chunk_index": row["chunk_index"],
                "total_chunks": total_chunks,
                "tags": [r["name"] for r in tag_rows],
            },
        })

    return {
        "query": query,
        "results": results,
        "total_matches": total,
        "returned": len(results),
    }


def _fts_search(conn: sqlite3.Connection, query: str, limit: int,
                tags: list[str] | None, doc_type: str | None) -> dict[int, float]:
    """Run FTS5 search, return {chunk_id: bm25_score}."""
    escaped = _escape_fts_query(query)
    if not escaped.strip():
        return {}

    sql = """
        SELECT f.rowid as chunk_id, bm25(chunks_fts) as score
        FROM chunks_fts f
    """
    joins = []
    where = [f"chunks_fts MATCH ?"]
    params = [escaped]

    if tags or doc_type:
        joins.append("JOIN chunks c ON f.rowid = c.id")
        joins.append("JOIN documents d ON c.document_id = d.id")

    if doc_type:
        where.append("d.doc_type = ?")
        params.append(doc_type)

    if tags:
        for i, tag in enumerate(tags):
            joins.append(f"JOIN document_tags dt{i} ON d.id = dt{i}.document_id")
            joins.append(f"JOIN tags t{i} ON dt{i}.tag_id = t{i}.id")
            where.append(f"t{i}.name = ?")
            params.append(tag.strip().lower())

    sql += " " + " ".join(joins)
    sql += " WHERE " + " AND ".join(where)
    sql += " ORDER BY score LIMIT ?"
    params.append(limit)

    rows = conn.execute(sql, params).fetchall()

    # BM25 scores are negative (lower = better), normalise to positive
    results = {}
    for row in rows:
        results[row["chunk_id"]] = -row["score"]  # Negate so higher = better

    return results


def _vector_search(conn: sqlite3.Connection, query: str, model_name: str,
                   cfg: dict, limit: int, tags: list[str] | None,
                   doc_type: str | None) -> dict[int, float]:
    """Run vector similarity search, return {chunk_id: similarity_score}."""
    from kb_search.embeddings import embed_texts

    prefix = cfg.get("embedding", {}).get("query_prefix", "")
    query_emb = embed_texts(model_name, [query], prefix=prefix)[0]
    blob = struct.pack(f"{len(query_emb)}f", *query_emb)

    # sqlite-vec returns results ordered by distance (lower = more similar)
    rows = conn.execute("""
        SELECT chunk_id, distance
        FROM chunks_vec
        WHERE embedding MATCH ?
        ORDER BY distance
        LIMIT ?
    """, (blob, limit)).fetchall()

    results = {}
    for row in rows:
        # Convert distance to similarity (1 - distance for cosine)
        similarity = max(0, 1 - row["distance"])
        chunk_id = row["chunk_id"]

        # Apply filters post-hoc for vector search
        if tags or doc_type:
            check = conn.execute("""
                SELECT 1 FROM chunks c
                JOIN documents d ON c.document_id = d.id
                WHERE c.id = ?
            """ + (" AND d.doc_type = ?" if doc_type else ""),
                (chunk_id,) + ((doc_type,) if doc_type else ())
            ).fetchone()
            if not check:
                continue

            if tags:
                tag_count = conn.execute("""
                    SELECT COUNT(*) FROM chunks c
                    JOIN documents d ON c.document_id = d.id
                    JOIN document_tags dt ON d.id = dt.document_id
                    JOIN tags t ON dt.tag_id = t.id
                    WHERE c.id = ? AND t.name IN ({})
                """.format(",".join("?" * len(tags))),
                    (chunk_id, *[t.strip().lower() for t in tags])
                ).fetchone()[0]
                if tag_count < len(tags):
                    continue

        results[chunk_id] = similarity

    return results


def _rrf_merge(fts_results: dict[int, float], vec_results: dict[int, float],
               k: int = 60) -> list[dict]:
    """Merge two result sets using Reciprocal Rank Fusion."""
    # Rank each result set
    fts_ranked = _rank_results(fts_results)
    vec_ranked = _rank_results(vec_results)

    all_ids = set(fts_ranked.keys()) | set(vec_ranked.keys())

    merged = []
    for chunk_id in all_ids:
        fts_rank = fts_ranked.get(chunk_id)
        vec_rank = vec_ranked.get(chunk_id)

        score = 0
        if fts_rank is not None:
            score += 1 / (k + fts_rank)
        if vec_rank is not None:
            score += 1 / (k + vec_rank)

        fts_score = round(1 / (k + fts_rank), 6) if fts_rank is not None else None
        vec_score = round(1 / (k + vec_rank), 6) if vec_rank is not None else None

        merged.append({
            "chunk_id": chunk_id,
            "score": score,
            "score_breakdown": {"fts": fts_score, "vector": vec_score},
        })

    return merged


def _single_source_results(results: dict[int, float], source: str) -> list[dict]:
    """Convert single-source results to merged format."""
    ranked = _rank_results(results)
    merged = []
    for chunk_id, rank in ranked.items():
        score = results[chunk_id]
        breakdown = {"fts": None, "vector": None}
        breakdown[source] = round(score, 6)
        merged.append({
            "chunk_id": chunk_id,
            "score": score,
            "score_breakdown": breakdown,
        })
    return merged


def _rank_results(results: dict[int, float]) -> dict[int, int]:
    """Rank results by score (1-indexed, higher score = lower rank number)."""
    sorted_ids = sorted(results.keys(), key=lambda x: results[x], reverse=True)
    return {chunk_id: rank + 1 for rank, chunk_id in enumerate(sorted_ids)}


def _escape_fts_query(query: str) -> str:
    """Escape special FTS5 characters in a query."""
    # Remove FTS5 operators that could cause syntax errors
    query = re.sub(r'["\(\)\*\:\^]', " ", query)
    # Collapse multiple spaces
    query = re.sub(r"\s+", " ", query).strip()
    return query
