#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import sqlite3
from pathlib import Path
from typing import Any

import numpy as np
from sentence_transformers import SentenceTransformer
import os
from pathlib import Path

DB_PATH = Path("cheat.db")

LOCAL_CACHE_DIR = Path("models/hf")
os.environ.setdefault("HF_HOME", str(LOCAL_CACHE_DIR.resolve()))
os.environ.setdefault("SENTENCE_TRANSFORMERS_HOME", str(LOCAL_CACHE_DIR.resolve()))

def deserialize_embedding(blob: bytes, dim: int) -> np.ndarray:
    vec = np.frombuffer(blob, dtype=np.float32)
    if vec.shape[0] != dim:
        raise ValueError(f"Embedding length mismatch: expected {dim}, got {vec.shape[0]}")
    return vec


def load_index(conn: sqlite3.Connection) -> tuple[list[dict[str, Any]], np.ndarray, str]:
    rows = conn.execute("""
    SELECT
        c.id,
        c.command,
        c.explanation,
        c.intent_json,
        c.alternatives_json,
        c.requires_json,
        c.packages_json,
        c.tags_json,
        c.platform_json,
        c.shell_json,
        c.safety,
        e.model_name,
        e.embedding_blob,
        e.embedding_dim
    FROM cards c
    JOIN card_embeddings e ON c.id = e.card_id
    ORDER BY c.id
    """).fetchall()

    if not rows:
        raise RuntimeError("No indexed cards found. Run build_index.py first.")

    cards: list[dict[str, Any]] = []
    vectors: list[np.ndarray] = []
    model_name: str | None = None

    for row in rows:
        (
            card_id,
            command,
            explanation,
            intent_json,
            alternatives_json,
            requires_json,
            packages_json,
            tags_json,
            platform_json,
            shell_json,
            safety,
            row_model_name,
            embedding_blob,
            embedding_dim,
        ) = row

        if model_name is None:
            model_name = row_model_name
        elif model_name != row_model_name:
            raise RuntimeError("Mixed embedding models found in the index.")

        cards.append({
            "id": card_id,
            "command": command,
            "explanation": explanation,
            "intent": json.loads(intent_json),
            "alternatives": json.loads(alternatives_json),
            "requires": json.loads(requires_json),
            "packages": json.loads(packages_json),
            "tags": json.loads(tags_json),
            "platform": json.loads(platform_json),
            "shell": json.loads(shell_json),
            "safety": safety,
        })
        vectors.append(deserialize_embedding(embedding_blob, embedding_dim))

    matrix = np.vstack(vectors)
    return cards, matrix, model_name


def search(
    query: str,
    top_k: int = 5,
) -> list[dict[str, Any]]:
    conn = sqlite3.connect(DB_PATH)
    try:
        cards, matrix, model_name = load_index(conn)
    finally:
        conn.close()

    model = SentenceTransformer(
        model_name,
        cache_folder=str(LOCAL_CACHE_DIR.resolve()),
        local_files_only=True,
    )

    qvec = model.encode([query], normalize_embeddings=True, convert_to_numpy=True)[0]

    scores = matrix @ qvec
    top_indices = np.argsort(scores)[::-1][:top_k]

    results: list[dict[str, Any]] = []
    for idx in top_indices:
        card = dict(cards[idx])
        card["score"] = float(scores[idx])
        results.append(card)

    return results


def main() -> None:
    parser = argparse.ArgumentParser(description="Query the local command card index.")
    parser.add_argument("query", type=str, help="Natural language query")
    parser.add_argument("--top-k", type=int, default=5, help="Number of results to return")
    args = parser.parse_args()

    results = search(args.query, top_k=args.top_k)

    for i, result in enumerate(results, start=1):
        print(f"[{i}] score={result['score']:.4f} id={result['id']}")
        print(f"    command: {result['command']}")
        print(f"    explanation: {result['explanation']}")
        if result["alternatives"]:
            print(f"    alternatives: {', '.join(result['alternatives'])}")
        print(f"    intent: {', '.join(result['intent'][:3])}")
        print()


if __name__ == "__main__":
    main()