#!/usr/bin/env python3
"""
Database setup script for Tech4biz Wireframe Generator
This script creates the database and runs the schema files
"""

import os
import psycopg2
from psycopg2.extensions import ISOLATION_LEVEL_AUTOCOMMIT
from dotenv import load_dotenv

def setup_database():
    """Setup the database and create tables"""
    
    # Load environment variables
    load_dotenv()
    
    # Database connection details
    db_host = os.getenv('POSTGRES_HOST', 'localhost')
    db_user = os.getenv('POSTGRES_USER', 'pipeline_admin')
    db_password = os.getenv('POSTGRES_PASSWORD', 'secure_pipeline_2024')
    db_port = os.getenv('POSTGRES_PORT', '5432')  # Changed to 5432 for Docker
    db_name = os.getenv('POSTGRES_DB', 'dev_pipeline')
    
    # First connect to postgres to create database
    try:
        conn = psycopg2.connect(
            host=db_host,
            user=db_user,
            password=db_password,
            port=db_port,
            database='postgres'  # Connect to default postgres database first
        )
        conn.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
        
        with conn.cursor() as cur:
            # Check if database exists
            cur.execute("SELECT 1 FROM pg_database WHERE datname = %s", (db_name,))
            if not cur.fetchone():
                print(f"Creating database '{db_name}'...")
                cur.execute(f"CREATE DATABASE {db_name}")
                print("Database created successfully!")
            else:
                print(f"Database '{db_name}' already exists")
        
        conn.close()
        
    except Exception as e:
        print(f"Error creating database: {e}")
        return False
    
    # Now connect to the new database and run schema files
    try:
        conn = psycopg2.connect(
            host=db_host,
            user=db_user,
            password=db_password,
            port=db_port,
            database=db_name
        )
        
        with conn.cursor() as cur:
            print("Running user authentication schema...")
            schema_file = os.path.join(os.path.dirname(__file__), 'sql', '001_user_auth_schema.sql')
            with open(schema_file, 'r') as f:
                schema_sql = f.read()
                cur.execute(schema_sql)
            
            print("Running wireframe schema...")
            schema_file = os.path.join(os.path.dirname(__file__), 'sql', '002_wireframe_schema.sql')
            with open(schema_file, 'r') as f:
                schema_sql = f.read()
                cur.execute(schema_sql)
            
            conn.commit()
            print("Database setup completed successfully!")
            
        conn.close()
        return True
        
    except Exception as e:
        print(f"Error setting up database: {e}")
        return False

if __name__ == "__main__":
    print("Setting up Tech4biz Wireframe Generator database...")
    success = setup_database()
    
    if success:
        print("\n✅ Database setup completed successfully!")
        print("\nNext steps:")
        print("1. Make sure your .env file has the correct database credentials")
        print("2. Start the backend server with: python app.py")
        print("3. The wireframes will now be automatically saved and loaded!")
    else:
        print("\n❌ Database setup failed!")
        print("Please check your database connection settings and try again.")
