import os
import psycopg2
from psycopg2.extras import execute_values
from dotenv import load_dotenv

load_dotenv()

def get_conn():
    return psycopg2.connect(
        host=os.getenv("DB_HOST"),
        port=os.getenv("DB_PORT"),
        dbname=os.getenv("DB_NAME"),
        user=os.getenv("DB_USER"),
        password=os.getenv("DB_PASSWORD"),
    )

def init_db():
    conn = get_conn()
    cur = conn.cursor()
    cur.execute("CREATE EXTENSION IF NOT EXISTS vector;")
    cur.execute("""
        CREATE TABLE IF NOT EXISTS document_chunks (
            id          SERIAL PRIMARY KEY,
            filename    TEXT NOT NULL,
            page_num    INTEGER,
            chunk_index INTEGER,
            text        TEXT NOT NULL,
            embedding   vector(768),
            nationality TEXT DEFAULT 'unknown',
            corpus      TEXT DEFAULT 'unknown',
            created_at  TIMESTAMPTZ DEFAULT NOW()
        );
    """)
    # Migrate existing tables that predate nationality/corpus columns
    for col, default in [("nationality", "unknown"), ("corpus", "unknown")]:
        cur.execute(f"""
            ALTER TABLE document_chunks
            ADD COLUMN IF NOT EXISTS {col} TEXT DEFAULT '{default}';
        """)
    cur.execute("""
        CREATE INDEX IF NOT EXISTS document_chunks_embedding_idx
        ON document_chunks
        USING ivfflat (embedding vector_cosine_ops)
        WITH (lists = 100);
    """)
    cur.execute("""
        CREATE INDEX IF NOT EXISTS document_chunks_nationality_idx
        ON document_chunks (nationality);
    """)
    cur.execute("""
        CREATE INDEX IF NOT EXISTS document_chunks_corpus_idx
        ON document_chunks (corpus);
    """)
    conn.commit()
    cur.close()
    conn.close()
    print("DB initialized.")

def upsert_chunks(rows: list[dict]):
    conn = get_conn()
    cur = conn.cursor()
    execute_values(
        cur,
        """
        INSERT INTO document_chunks (filename, page_num, chunk_index, text, embedding, nationality, corpus)
        VALUES %s
        """,
        [
            (
                r["filename"], r["page_num"], r["chunk_index"], r["text"], r["embedding"],
                r.get("nationality", "unknown"), r.get("corpus", "unknown"),
            )
            for r in rows
        ],
        template="(%s, %s, %s, %s, %s::vector, %s, %s)"
    )
    conn.commit()
    cur.close()
    conn.close()

def semantic_search(query_embedding: list[float], top_k: int = 5,
                    nationality: str = None, corpus: str = None):
    conn = get_conn()
    cur = conn.cursor()
    filters = []
    params  = []
    if nationality:
        filters.append("nationality = %s")
        params.append(nationality)
    if corpus:
        filters.append("corpus = %s")
        params.append(corpus)
    where = ("WHERE " + " AND ".join(filters)) if filters else ""
    cur.execute(
        f"""
        SELECT filename, page_num, chunk_index, text, nationality, corpus,
               1 - (embedding <=> %s::vector) AS similarity
        FROM document_chunks
        {where}
        ORDER BY embedding <=> %s::vector
        LIMIT %s;
        """,
        [query_embedding] + params + [query_embedding, top_k]
    )
    rows = cur.fetchall()
    cur.close()
    conn.close()
    return rows