#!/opt/homebrew/bin/python3.11
"""
Natural Language Vector Search Comparison: PostgreSQL pgvector vs Qdrant
"""

import time
import requests
import psycopg2


VIDEO_UUID = "39567a0eb16f39fd"

POSTGRES_CONFIG = {
    "host": "localhost",
    "port": 5432,
    "user": "accusys",
    "password": "Test3200",
    "database": "momentry",
}


# Natural language test queries
NATURAL_LANGUAGE_QUERIES = [
    # Scene descriptions
    "a person talking",
    "someone speaking on camera",
    "outdoor scene",
    "indoor setting",
    # Actions
    "walking or moving",
    "dialogue or conversation",
    "looking at something",
    # Emotions/tone
    "happy or joyful",
    "serious or dramatic",
    "comedy or funny",
    # Objects
    "wearing a tie",
    "holding an object",
    "sitting on a chair",
    # Locations
    "city or urban",
    "building or room",
    "open space",
]


def get_embedding(text):
    """Get embedding from Ollama"""
    resp = requests.post(
        "http://localhost:11434/api/embeddings",
        json={"model": "nomic-embed-text", "prompt": text},
    )
    return resp.json()["embedding"]


def test_qdrant(queries):
    """Test Qdrant vector search"""
    results = {}

    for query in queries:
        # Get embedding
        embedding = get_embedding(query)

        # Search in Qdrant
        start = time.time()
        resp = requests.post(
            "http://localhost:6333/collections/AccusysDB/points/search",
            headers={"api-key": "Test3200Test3200Test3200"},
            json={"vector": embedding, "limit": 5},
        )
        elapsed = (time.time() - start) * 1000

        data = resp.json()
        results[query] = {
            "ms": round(elapsed, 2),
            "scores": [r["score"] for r in data.get("result", [])],
        }
        print(
            f"Qdrant: '{query}' -> {elapsed:.2f}ms, top score: {results[query]['scores'][0]:.3f}"
            if results[query]["scores"]
            else f"Qdrant: '{query}' -> {elapsed:.2f}ms, no results"
        )

    return results


def test_pgvector(queries):
    """Test PostgreSQL pgvector"""
    results = {}
    conn = psycopg2.connect(**POSTGRES_CONFIG)
    cur = conn.cursor()

    for query in queries:
        # Get embedding
        embedding = get_embedding(query)
        vector_str = "[" + ",".join(str(x) for x in embedding) + "]"

        # Search in PostgreSQL
        start = time.time()
        cur.execute(
            """
            SELECT chunk_id, (embedding_vector <=> %s::vector) as distance
            FROM chunk_vectors
            WHERE embedding_vector IS NOT NULL
            ORDER BY embedding_vector <=> %s::vector
            LIMIT 5
        """,
            (vector_str, vector_str),
        )

        rows = cur.fetchall()
        elapsed = (time.time() - start) * 1000

        # Convert distance to similarity score (1 - distance for cosine)
        scores = [1 - r[1] for r in rows]

        results[query] = {
            "ms": round(elapsed, 2),
            "scores": scores,
            "chunk_ids": [r[0] for r in rows],
        }
        print(
            f"pgvector: '{query}' -> {elapsed:.2f}ms, top score: {scores[0]:.3f}"
            if scores
            else f"pgvector: '{query}' -> {elapsed:.2f}ms, no results"
        )

    cur.close()
    conn.close()
    return results


def get_chunk_text(chunk_id):
    """Get chunk text content"""
    conn = psycopg2.connect(**POSTGRES_CONFIG)
    cur = conn.cursor()
    cur.execute("SELECT content->>'text' FROM chunks WHERE chunk_id = %s", (chunk_id,))
    result = cur.fetchone()
    cur.close()
    conn.close()
    return result[0][:100] if result else ""


def main():
    print("=" * 70)
    print("NATURAL LANGUAGE VECTOR SEARCH COMPARISON")
    print("=" * 70)
    print(f"\nVideo UUID: {VIDEO_UUID}")
    print(f"Testing {len(NATURAL_LANGUAGE_QUERIES)} natural language queries\n")

    print("=" * 70)
    print("QDRANT SEARCH")
    print("=" * 70)
    qdrant_results = test_qdrant(NATURAL_LANGUAGE_QUERIES)

    print("\n" + "=" * 70)
    print("POSTGRESQL PGVECTOR SEARCH")
    print("=" * 70)
    pgvector_results = test_pgvector(NATURAL_LANGUAGE_QUERIES)

    # Calculate averages
    qdrant_avg = sum(r["ms"] for r in qdrant_results.values()) / len(qdrant_results)
    pgvector_avg = sum(r["ms"] for r in pgvector_results.values()) / len(
        pgvector_results
    )

    print("\n" + "=" * 70)
    print("SUMMARY")
    print("=" * 70)
    print(f"\n{'Query':<30} | {'Qdrant':<15} | {'pgvector':<15}")
    print("-" * 65)

    for query in NATURAL_LANGUAGE_QUERIES:
        qd = qdrant_results[query]
        pg = pgvector_results[query]

        qd_score = f"{qd['scores'][0]:.3f}" if qd["scores"] else "N/A"
        pg_score = f"{pg['scores'][0]:.3f}" if pg["scores"] else "N/A"

        print(
            f"{query:<30} | {qd['ms']:>5.1f}ms {qd_score:<7} | {pg['ms']:>5.1f}ms {pg_score}"
        )

    print("\n" + "=" * 70)
    print("AVERAGE RESPONSE TIME")
    print("=" * 70)
    print(f"  Qdrant:     {qdrant_avg:.2f}ms")
    print(f"  pgvector:   {pgvector_avg:.2f}ms")
    print(
        f"  Difference: {abs(qdrant_avg - pgvector_avg):.2f}ms ({'Qdrant faster' if qdrant_avg < pgvector_avg else 'pgvector faster'})"
    )

    # Show sample results
    print("\n" + "=" * 70)
    print("SAMPLE RESULTS")
    print("=" * 70)

    sample_query = "a person talking"
    print(f"\nQuery: '{sample_query}'")

    print("\nQdrant results:")
    for i, r in enumerate(qdrant_results[sample_query]["scores"][:3]):
        chunk_id = f"sentence_{(i * 3):04d}"  # Approximate
        print(f"  {i + 1}. score: {r:.3f}")

    print("\npgvector results:")
    for i, (chunk_id, score) in enumerate(
        zip(
            pgvector_results[sample_query]["chunk_ids"][:3],
            pgvector_results[sample_query]["scores"][:3],
        )
    ):
        text = get_chunk_text(chunk_id)
        print(f"  {i + 1}. {chunk_id} (score: {score:.3f})")
        print(f'     "{text}..."')


if __name__ == "__main__":
    main()
