import os
import wmill
import yaml
from sqlalchemy import create_engine, text, MetaData, Table, Column, Integer, String, inspect
from sqlalchemy.engine import Engine
import psycopg2

# You can import any PyPi package. 
# See here for more info: https://www.windmill.dev/docs/advanced/dependencies_in_python

# you can use typed resources by doing a type alias to dict
#postgresql = dict

DB_RESOURCE_PATH = 'u/joshuakrzemien/slick_postgresql'
DB_CONFIG_PATH = 'f/CCR_ETL/ccr_db_config'

def create_db_engine(db: dict):
    db_url = f"postgresql+psycopg2://postgres:{db['password']}@{db['host']}:{db['port']}/{db['dbname']}"
    engine = create_engine(db_url)
    engine.connect()
    return engine

def table_exists(engine: Engine, table_name: str) -> bool:
    """Check if a table exists in the database."""
    inspector = inspect(engine)
    return table_name in inspector.get_table_names()

def create_table(engine: Engine, table: dict, strategy: str = "create_if_not_exists"):
    try:
        table_name = table['name']
        columns = table['columns']
        
        # Handle different table strategies
        if strategy == "drop_and_recreate":
            if table_exists(engine, table_name):
                print(f"Dropping existing table: {table_name}")
                with engine.connect() as conn:
                    conn.execute(text(f"DROP TABLE IF EXISTS {table_name} CASCADE"))
                    conn.commit()
        
        elif strategy == "create_if_not_exists":
            if table_exists(engine, table_name):
                print(f"Table {table_name} already exists, skipping creation")
                return
        
        else:
            raise ValueError(f"Unknown table strategy: {strategy}")
        
        # Map config types to SQLAlchemy types
        type_mapping = {
            'integer': Integer,
            'string': String
        }
        
        # Build SQLAlchemy columns
        sqlalchemy_columns = []
        for column in columns:
            col_type = type_mapping.get(column['type'], String)
            sqlalchemy_columns.append(Column(column['name'], col_type, primary_key=column.get('primary_key', False), nullable=column.get('nullable', True), index=column.get('index', False), autoincrement=column.get('autoincrement', False)))
        
        # Create table using SQLAlchemy Core
        metadata = MetaData()
        new_table = Table(table_name, metadata, *sqlalchemy_columns)
        
        # Create the table
        metadata.create_all(engine)
        
        print(f"Successfully created table: {table_name}")
        
    except Exception as e:
        print(f"Error creating table {table_name}: {str(e)}")
        raise

def main():
    db = wmill.client.get_resource(DB_RESOURCE_PATH)
    config_yaml = wmill.get_variable(DB_CONFIG_PATH)
    config = yaml.safe_load(config_yaml)
    engine = create_db_engine(db)
    
    # Get table strategy from config (default to drop_and_recreate)
    table_strategy = config.get('table_strategy', 'drop_and_recreate')
    print(f"Using table strategy: {table_strategy}")
    
    for table in config['schema']['tables']:
        # Allow per-table strategy override
        table_specific_strategy = table.get('strategy', table_strategy)
        create_table(engine, table, table_specific_strategy)
    
    return {"status": "success"}