import argparse
import json
import sys
from pathlib import Path

import numpy as np
from src.memory.embedder import embed_texts


def _cosine_sim(a: np.ndarray, b: np.ndarray) -> float:
    norm_a = np.linalg.norm(a)
    norm_b = np.linalg.norm(b)
    if norm_a == 0 or norm_b == 0:
        return 0.0
    return float(np.dot(a, b) / (norm_a * norm_b))


def _load_index(index_path: Path) -> list[dict[str, object]]:
    if not index_path.exists():
        return []
    index = json.loads(index_path.read_text(encoding="utf-8"))
    files = index.get("files", []) if isinstance(index, dict) else []
    return files if isinstance(files, list) else []


def _rank_by_similarity(
    query_vec: np.ndarray, files: list[dict[str, object]]
) -> list[dict[str, str | float]]:
    results: list[dict[str, str | float]] = []
    for f in files:
        file_vec = np.array(f["embedding"])
        sim = _cosine_sim(query_vec, file_vec)
        results.append(
            {
                "source": str(f["source"]),
                "score": sim,
                "text": str(f["text"]),
            }
        )
    results.sort(key=lambda r: r["score"], reverse=True)
    return results


def _dedup_top(results: list[dict[str, str | float]], k: int) -> list[dict[str, str | float]]:
    seen: set[str] = set()
    top: list[dict[str, str | float]] = []
    for r in results:
        src = str(r["source"])
        if src in seen:
            continue
        seen.add(src)
        top.append(r)
        if len(top) >= k:
            break
    return top


def run_search(args: argparse.Namespace) -> None:
    index_path = Path(args.index_dir) / "index.json"
    files = _load_index(index_path)
    if not files:
        json.dump([], sys.stdout)
        return

    query_emb = embed_texts([args.query])
    if query_emb is None:
        json.dump([], sys.stdout)
        return
    query_vec = np.array(query_emb[0])

    results = _rank_by_similarity(query_vec, files)
    top = _dedup_top(results, args.k)
    print(json.dumps(top, ensure_ascii=False))
