import psycopg2
from psycopg2 import pool
from contextlib import contextmanager
from typing import Generator
from app.config import settings
import json

# Connection pool for better performance
connection_pool = pool.SimpleConnectionPool(1, 20, settings.database_url)

@contextmanager
def get_db_connection():
    """Get a database connection from the pool"""
    conn = connection_pool.getconn()
    try:
        yield conn
        conn.commit()
    except Exception as e:
        conn.rollback()
        raise e
    finally:
        connection_pool.putconn(conn)

def init_db():
    """Initialize database tables"""
    with get_db_connection() as conn:
        cur = conn.cursor()
        
        # Create tables
        cur.execute("""
            CREATE TABLE IF NOT EXISTS users (
                id SERIAL PRIMARY KEY,
                email VARCHAR(255) UNIQUE NOT NULL,
                hashed_password VARCHAR(255) NOT NULL,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            );
        """)
        
        cur.execute("""
            CREATE TABLE IF NOT EXISTS profiles (
                id SERIAL PRIMARY KEY,
                user_id INTEGER NOT NULL UNIQUE,
                display_name VARCHAR(255) NOT NULL,
                age INTEGER NOT NULL,
                gender VARCHAR(50) NOT NULL,
                location VARCHAR(255) NOT NULL,
                bio TEXT,
                interests JSONB DEFAULT '[]',
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE
            );
        """)
        
        cur.execute("""
            CREATE TABLE IF NOT EXISTS photos (
                id SERIAL PRIMARY KEY,
                profile_id INTEGER NOT NULL,
                file_path VARCHAR(255) NOT NULL,
                display_order INTEGER NOT NULL,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                FOREIGN KEY (profile_id) REFERENCES profiles(id) ON DELETE CASCADE
            );
        """)
        
        cur.execute("""
            CREATE TABLE IF NOT EXISTS likes (
                id SERIAL PRIMARY KEY,
                liker_id INTEGER NOT NULL,
                liked_id INTEGER NOT NULL,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                UNIQUE(liker_id, liked_id),
                FOREIGN KEY (liker_id) REFERENCES users(id) ON DELETE CASCADE,
                FOREIGN KEY (liked_id) REFERENCES users(id) ON DELETE CASCADE
            );
        """)
        
        cur.execute("""
            CREATE TABLE IF NOT EXISTS conversations (
                id SERIAL PRIMARY KEY,
                user_id_1 INTEGER NOT NULL,
                user_id_2 INTEGER NOT NULL,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                UNIQUE(user_id_1, user_id_2),
                FOREIGN KEY (user_id_1) REFERENCES users(id) ON DELETE CASCADE,
                FOREIGN KEY (user_id_2) REFERENCES users(id) ON DELETE CASCADE
            );
        """)
        
        cur.execute("""
            CREATE TABLE IF NOT EXISTS messages (
                id SERIAL PRIMARY KEY,
                conversation_id INTEGER NOT NULL,
                sender_id INTEGER NOT NULL,
                content TEXT NOT NULL,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                FOREIGN KEY (conversation_id) REFERENCES conversations(id) ON DELETE CASCADE,
                FOREIGN KEY (sender_id) REFERENCES users(id) ON DELETE CASCADE
            );
        """)
        
        # Create indexes for common queries
        cur.execute("CREATE INDEX IF NOT EXISTS idx_profiles_user_id ON profiles(user_id);")
        cur.execute("CREATE INDEX IF NOT EXISTS idx_photos_profile_id ON photos(profile_id);")
        cur.execute("CREATE INDEX IF NOT EXISTS idx_likes_liker_id ON likes(liker_id);")
        cur.execute("CREATE INDEX IF NOT EXISTS idx_likes_liked_id ON likes(liked_id);")
        cur.execute("CREATE INDEX IF NOT EXISTS idx_conversations_users ON conversations(user_id_1, user_id_2);")
        cur.execute("CREATE INDEX IF NOT EXISTS idx_messages_conversation_id ON messages(conversation_id);")
        cur.execute("CREATE INDEX IF NOT EXISTS idx_messages_created_at ON messages(created_at);")
        
        conn.commit()

def close_db():
    """Close all database connections"""
    if connection_pool:
        connection_pool.closeall()