from sqlalchemy import text
from sqlalchemy.exc import SQLAlchemyError

def ensure_schema(engine) -> bool:
    with engine.connect().execution_options(isolation_level="AUTOCOMMIT") as conn:
        conn.execute(text("CREATE EXTENSION IF NOT EXISTS vector"))
    
    with engine.begin() as conn:
        conn.execute(text("""
            CREATE TABLE IF NOT EXISTS domains (
                id SERIAL PRIMARY KEY,
                name TEXT UNIQUE NOT NULL,
                embedding vector(384) NOT NULL
            )
        """))
        conn.execute(text("""
            CREATE TABLE IF NOT EXISTS universal_dimensions (
                id SERIAL PRIMARY KEY,
                name TEXT UNIQUE NOT NULL,
                embedding vector(384) NOT NULL
            )
        """))
        conn.execute(text("""
            CREATE TABLE IF NOT EXISTS dimension_groups (
                id SERIAL PRIMARY KEY,
                domain_id INTEGER NOT NULL REFERENCES domains(id) ON DELETE CASCADE,
                name TEXT NOT NULL,
                UNIQUE(domain_id, name)
            )
        """))
        conn.execute(text("""
            CREATE TABLE IF NOT EXISTS attributes (
                id SERIAL PRIMARY KEY,
                group_id INTEGER REFERENCES dimension_groups(id) ON DELETE CASCADE,
                name TEXT NOT NULL,
                embedding vector(384) NOT NULL,
                UNIQUE(group_id, name)
            )
        """))
        conn.execute(text("ALTER TABLE attributes DROP CONSTRAINT IF EXISTS attributes_name_key"))
        
        conn.execute(text("""
        CREATE TABLE IF NOT EXISTS attribute_feedback (
        id SERIAL PRIMARY KEY,
        query TEXT NOT NULL,
        attribute_name TEXT NOT NULL,
        domain TEXT,
        click_count INTEGER DEFAULT 1,
        last_clicked TIMESTAMP DEFAULT NOW(),
        UNIQUE(query, attribute_name)
            )
        """))
        
        conn.execute(text("""
            ALTER TABLE domains 
            ADD COLUMN IF NOT EXISTS quality_score FLOAT DEFAULT 1.0,
            ADD COLUMN IF NOT EXISTS created_at TIMESTAMP DEFAULT NOW(),
            ADD COLUMN IF NOT EXISTS query_count INTEGER DEFAULT 0
        """))
    
    return True