import os
import psycopg2
from psycopg2.extras import RealDictCursor


class PGVectorService:
    """Service for managing embeddings with PostgreSQL pgvector"""

    def __init__(self):
        self.conn = psycopg2.connect(
            host=os.getenv("POSTGRES_HOST", "postgres"),
            user=os.getenv("POSTGRES_USER"),
            password=os.getenv("POSTGRES_PASSWORD"),
            dbname=os.getenv("POSTGRES_DB"),
        )

    def insert_embedding(self, id: int, embedding: list):
        """Insert or update an embedding"""
        with self.conn.cursor() as cur:
            cur.execute(
                """
                INSERT INTO embeddings_table (id, embedding)
                VALUES (%s, %s)
                ON CONFLICT (id) DO UPDATE SET embedding = EXCLUDED.embedding;
                """,
                (id, embedding),
            )
            self.conn.commit()

    def search_nearest(self, query_embedding: list, limit: int = 3):
        """Search nearest embeddings using cosine distance (<-> operator)"""
        with self.conn.cursor(cursor_factory=RealDictCursor) as cur:
            cur.execute(
                """
                SELECT id, embedding, embedding <-> %s AS distance
                FROM embeddings_table
                ORDER BY embedding <-> %s
                LIMIT %s;
                """,
                (query_embedding, query_embedding, limit),
            )
            return cur.fetchall()

    def close(self):
        if self.conn:
            self.conn.close()
