try:
    import psycopg2
except ImportError:
    psycopg2 = None

from django.conf import settings
from .gemini_client import GeminiClient
import numpy as np
import json

class RAGService:
    def __init__(self):
        self.client = GeminiClient()
        self.db_config = settings.DATABASES.get('vector_db')
        if self.db_config:
            self.conn_string = f"dbname='{self.db_config['NAME']}' user='{self.db_config['USER']}' host='{self.db_config['HOST']}' password='{self.db_config['PASSWORD']}' port='{self.db_config['PORT']}'"
            self._init_db()
        else:
            print("⚠️ Postgres vector_db no configurada. RAG Service deshabilitado.")
            self.conn_string = None

    def _get_connection(self):
        if not psycopg2 or not self.conn_string:
            raise Exception("Postgres vector_db no está disponible.")
        return psycopg2.connect(self.conn_string)

    def _init_db(self):
        """Initializes the vector database schema if it doesn't exist"""
        try:
            conn = self._get_connection()
            cur = conn.cursor()
            
            # Enable pgvector extension
            cur.execute("CREATE EXTENSION IF NOT EXISTS vector;")
            
            # Create documents table
            cur.execute("""
                CREATE TABLE IF NOT EXISTS documents (
                    id SERIAL PRIMARY KEY,
                    content TEXT,
                    metadata JSONB,
                    embedding vector(768),
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                );
            """)
            
            # Create index for faster search (optional but recommended)
            # cur.execute("CREATE INDEX ON documents USING ivfflat (embedding vector_cosine_ops) WITH (lists = 100);")
            
            conn.commit()
            cur.close()
            conn.close()
        except Exception as e:
            print(f"Error initializing DB: {e}")
            # Don't raise here, we might want to continue even if DB is down temporarily? 
            # Ideally we should raise, but let's log for now.
            
    def add_document(self, text, metadata=None):
        """Adds a document to the knowledge base"""
        embedding = self.client.embed_text(text)
        
        conn = self._get_connection()
        cur = conn.cursor()
        
        cur.execute(
            "INSERT INTO documents (content, metadata, embedding) VALUES (%s, %s, %s)",
            (text, json.dumps(metadata), embedding)
        )
        
        conn.commit()
        cur.close()
        conn.close()

    def search(self, query, limit=5):
        """Searches for similar documents"""
        embedding = self.client.embed_query(query)
        
        conn = self._get_connection()
        cur = conn.cursor()
        
        # Cosine distance search (<=> operator)
        cur.execute("""
            SELECT content, metadata, 1 - (embedding <=> %s::vector) as similarity
            FROM documents
            ORDER BY embedding <=> %s::vector
            LIMIT %s
        """, (embedding, embedding, limit))
        
        results = cur.fetchall()
        cur.close()
        conn.close()
        
        return [{"content": r[0], "metadata": r[1], "score": r[2]} for r in results]