#!/usr/bin/env python3
"""
Hybrid Jaccard matcher using pg_trgm + pgvector RRF.
Computes concept embeddings, then queries hybrid_search().

Usage:
  nix-shell /tmp/embed_env.nix --run "python3 /tmp/jaccard_hybrid.py"
"""

import json
import subprocess
import re
from pathlib import Path

CONCEPTS_PATH = "/tmp/cornfield_concepts.json"
MODEL_NAME = "sentence-transformers/static-retrieval-mrl-en-v1"

def psql(sql, timeout=60):
    r = subprocess.run(
        ['podman', 'exec', 'arxiv-pg', 'psql', '-U', 'postgres', '-d', 'arxiv',
         '-t', '-A', '-c', sql],
        capture_output=True, text=True, timeout=timeout
    )
    return r.stdout.strip(), r.stderr.strip()

def extract_query(concept):
    parts = [concept.get("name", "")]
    parts.extend(concept.get("tags", []))
    desc = concept.get("description", "")
    if desc:
        parts.append(desc[:200])
    novelty = concept.get("novelty_statement", "")
    if novelty:
        parts.append(novelty[:200])
    return " ".join(filter(None, parts))

def main():
    from sentence_transformers import SentenceTransformer

    print(f"Loading model: {MODEL_NAME}")
    model = SentenceTransformer(MODEL_NAME)
    print("Model loaded")

    # hybrid_search function already installed in DB
    print("Using existing hybrid_search function...")

    # Load concepts
    data = json.loads(Path(CONCEPTS_PATH).read_text())
    concepts = [c for c in data["concepts"] if c.get("id", "").startswith("cf_")]
    print(f"Processing {len(concepts)} cornfield concepts")

    results = []
    for i, concept in enumerate(concepts):
        cid = concept["id"]
        query_text = extract_query(concept)
        query_text_escaped = query_text.replace("'", "''")

        # Embed the concept query
        embedding = model.encode(query_text, normalize_embeddings=True)
        emb_str = '[' + ','.join(f'{x:.6f}' for x in embedding) + ']'

        # Hybrid search
        sql = f"""
        SELECT paper_id, title, trigram_rank, vector_rank, rrf_score
        FROM hybrid_search('{query_text_escaped}', '{emb_str}'::vector(1024), 10)
        """
        out, err = psql(sql, timeout=30)

        if err and "ERROR" in err:
            print(f"  [{i+1}/{len(concepts)}] {cid}: ERROR - {err[:80]}")
            continue

        matches = []
        for line in out.split("\n"):
            if not line.strip():
                continue
            parts = line.split("|")
            if len(parts) >= 5:
                matches.append({
                    "paper_id": parts[0],
                    "title": parts[1][:200],
                    "trigram_rank": parts[2],
                    "vector_rank": parts[3],
                    "rrf_score": float(parts[4])
                })

        if matches:
            print(f"  [{i+1}/{len(concepts)}] {cid}: {len(matches)} matches (best={matches[0]['rrf_score']:.6f})")
            results.append({
                "concept_id": cid,
                "query": query_text[:100],
                "top_matches": matches[:5]
            })
        else:
            print(f"  [{i+1}/{len(concepts)}] {cid}: no matches")

    # Save
    out_path = Path("/tmp/hybrid_matches.json")
    out_path.write_text(json.dumps(results, indent=2))
    total = sum(len(r["top_matches"]) for r in results)
    print(f"\nDone: {len(results)} concepts, {total} citations → {out_path}")

if __name__ == "__main__":
    main()