import psycopg2
from psycopg2 import sql
from config import DB_HOST, DB_PORT, DB_NAME, DB_USER, DB_PASSWORD

def get_db_connection():
    """Connect to the PostgreSQL database."""
    return psycopg2.connect(
        host=DB_HOST,
        port=DB_PORT,
        dbname=DB_NAME,
        user=DB_USER,
        password=DB_PASSWORD
    )

def create_tables():
    """Create database tables if they don't exist."""
    commands = [
        """
        CREATE TABLE IF NOT EXISTS trains (
            train_id VARCHAR PRIMARY KEY,
            train_type VARCHAR,
            model VARCHAR,
            capacity INT,
            age_years FLOAT,
            maintenance_history JSONB,
            technical_specs JSONB
        )
        """,
        """
        CREATE TABLE IF NOT EXISTS schedules (
            schedule_id UUID PRIMARY KEY,
            train_id VARCHAR,
            route_id VARCHAR,
            departure_station VARCHAR,
            arrival_station VARCHAR,
            scheduled_departure TIMESTAMP,
            scheduled_arrival TIMESTAMP,
            platform_departure VARCHAR,
            platform_arrival VARCHAR,
            distance_km FLOAT,
            scheduled_stops INT,
            service_type VARCHAR,
            FOREIGN KEY (train_id) REFERENCES trains(train_id)
        )
        """,
        """
        CREATE TABLE IF NOT EXISTS actual_journeys (
            journey_id UUID PRIMARY KEY,
            schedule_id UUID,
            actual_departure TIMESTAMP,
            actual_arrival TIMESTAMP,
            actual_platform_departure VARCHAR,
            actual_platform_arrival VARCHAR,
            delay_departure_minutes INT,
            delay_arrival_minutes INT,
            cancellation_flag BOOLEAN,
            cancellation_reason VARCHAR,
            passenger_count INT,
            load_factor FLOAT,
            FOREIGN KEY (schedule_id) REFERENCES schedules(schedule_id)
        )
        """,
        """
        CREATE TABLE IF NOT EXISTS stations (
            station_id VARCHAR PRIMARY KEY,
            station_name VARCHAR,
            latitude FLOAT,
            longitude FLOAT,
            elevation FLOAT,
            number_of_platforms INT,
            daily_passenger_volume INT,
            station_category INT,
            facilities JSONB,
            connection_types JSONB
        )
        """,
        """
        CREATE TABLE IF NOT EXISTS weather_history (
            weather_id UUID PRIMARY KEY,
            station_id VARCHAR,
            timestamp TIMESTAMP,
            temperature FLOAT,
            precipitation FLOAT,
            wind_speed FLOAT,
            wind_direction INT,
            humidity FLOAT,
            pressure FLOAT,
            visibility FLOAT,
            weather_condition VARCHAR,
            FOREIGN KEY (station_id) REFERENCES stations(station_id)
        )
        """
    ]
    conn = None
    try:
        conn = get_db_connection()
        cur = conn.cursor()
        for command in commands:
            cur.execute(command)
        cur.close()
        conn.commit()
    except Exception as e:
        print(f"Error creating tables: {e}")
    finally:
        if conn is not None:
            conn.close()