try:
    from airflow import DAG
    from datetime import timedelta
    from datetime import datetime

    from airflow import settings
    from airflow.models import Connection, Variable
    from airflow.operators.python_operator import PythonOperator
    from airflow.hooks.base import BaseHook
except Exception as e:
    print("Error {} ".format(e))

# Postgres RDS Connection Configuration
conns = [
    {
        "conn_id": "postgres_rds",
        "conn_type": "postgres",
        "host": "drowdatewarehouse.crlwwhgepgi7.ap-east-1.rds.amazonaws.com",
        "schema": "drowDateWareHouse",
        "login": "dRowAdmin",
        "password": "drowsuper",
        "port": 5432
    }
]

# Airflow Variables Configuration
variables = [
    # Example:
    # {
    #     "key": "example_var",
    #     "value": "example_value",
    #     "description": "Example variable description"
    # }
]

def clear_error_connections():
    """Clear connections that have encryption errors"""
    session = settings.Session()
    try:
        print("Checking for erroneous connections...")
        error_conns = session.query(Connection).filter(Connection.is_encrypted == True).all()
        print(f"Found {len(error_conns)} connections with encryption flag")
        
        for conn in error_conns:
            session.delete(conn)
            print(f"Deleted erroneous connection: {conn.conn_id}")
        session.commit()
        print("Error connections cleared.")
    except Exception as e:
        print(f"Error clearing connections: {e}")
        session.rollback()
    finally:
        session.close()

def validate_connections():
    """Test all connections and identify invalid ones"""
    session = settings.Session()
    invalid_connections = []
    valid_connections = []
    
    try:
        print("=" * 80)
        print("VALIDATING ALL AIRFLOW CONNECTIONS")
        print("=" * 80)
        
        all_connections = session.query(Connection).all()
        print(f"\nTotal connections found: {len(all_connections)}\n")
        
        for conn in all_connections:
            conn_id = conn.conn_id
            conn_type = conn.conn_type
            
            print(f"Testing connection: '{conn_id}' (type: {conn_type})")
            
            try:
                # Try to get the connection using BaseHook
                hook_conn = BaseHook.get_connection(conn_id)
                
                # Attempt to test the connection based on type
                if conn_type == 'postgres':
                    # Try to import and test PostgreSQL connection
                    try:
                        import psycopg2
                        conn_uri = hook_conn.get_uri()
                        test_conn = psycopg2.connect(
                            host=hook_conn.host,
                            port=hook_conn.port,
                            user=hook_conn.login,
                            password=hook_conn.password,
                            database=hook_conn.schema,
                            connect_timeout=10
                        )
                        test_conn.close()
                        print(f"  ✓ Connection '{conn_id}' is VALID")
                        valid_connections.append({
                            'conn_id': conn_id,
                            'conn_type': conn_type,
                            'host': hook_conn.host,
                            'port': hook_conn.port,
                            'schema': hook_conn.schema
                        })
                    except ImportError:
                        print(f"  ⚠ Cannot test '{conn_id}': psycopg2 not installed")
                        valid_connections.append({
                            'conn_id': conn_id,
                            'conn_type': conn_type,
                            'host': hook_conn.host,
                            'status': 'untested - missing driver'
                        })
                    except Exception as e:
                        print(f"  ✗ Connection '{conn_id}' is INVALID: {str(e)}")
                        invalid_connections.append({
                            'conn_id': conn_id,
                            'conn_type': conn_type,
                            'host': hook_conn.host,
                            'port': hook_conn.port,
                            'schema': hook_conn.schema,
                            'error': str(e)
                        })
                        
                elif conn_type == 'mysql':
                    try:
                        import pymysql
                        test_conn = pymysql.connect(
                            host=hook_conn.host,
                            port=hook_conn.port or 3306,
                            user=hook_conn.login,
                            password=hook_conn.password,
                            database=hook_conn.schema,
                            connect_timeout=10
                        )
                        test_conn.close()
                        print(f"  ✓ Connection '{conn_id}' is VALID")
                        valid_connections.append({
                            'conn_id': conn_id,
                            'conn_type': conn_type,
                            'host': hook_conn.host,
                            'port': hook_conn.port,
                            'schema': hook_conn.schema
                        })
                    except ImportError:
                        print(f"  ⚠ Cannot test '{conn_id}': pymysql not installed")
                        valid_connections.append({
                            'conn_id': conn_id,
                            'conn_type': conn_type,
                            'status': 'untested - missing driver'
                        })
                    except Exception as e:
                        print(f"  ✗ Connection '{conn_id}' is INVALID: {str(e)}")
                        invalid_connections.append({
                            'conn_id': conn_id,
                            'conn_type': conn_type,
                            'host': hook_conn.host,
                            'error': str(e)
                        })
                        
                else:
                    # For other connection types, just check if they can be retrieved
                    print(f"  ⚠ Connection '{conn_id}' type '{conn_type}' - basic validation only")
                    valid_connections.append({
                        'conn_id': conn_id,
                        'conn_type': conn_type,
                        'status': 'exists - detailed test not implemented'
                    })
                    
            except Exception as e:
                print(f"  ✗ Connection '{conn_id}' ERROR: {str(e)}")
                invalid_connections.append({
                    'conn_id': conn_id,
                    'conn_type': conn_type,
                    'error': str(e)
                })
        
        # Print summary
        print("\n" + "=" * 80)
        print("VALIDATION SUMMARY")
        print("=" * 80)
        print(f"\n✓ Valid/Testable Connections: {len(valid_connections)}")
        for conn in valid_connections:
            status = conn.get('status', 'VALID')
            print(f"  - {conn['conn_id']} ({conn['conn_type']}) - {status}")
        
        print(f"\n✗ Invalid/Problematic Connections: {len(invalid_connections)}")
        if invalid_connections:
            print("\n⚠️  ATTENTION: The following connections have issues:\n")
            for conn in invalid_connections:
                print(f"  - Connection ID: {conn['conn_id']}")
                print(f"    Type: {conn['conn_type']}")
                if 'host' in conn:
                    print(f"    Host: {conn.get('host', 'N/A')}")
                print(f"    Error: {conn['error']}")
                print()
        else:
            print("  (None found)")
        
        print("=" * 80)
        
        # Store invalid connections for future reference
        if invalid_connections:
            Variable.set(
                "invalid_connections_list", 
                str(invalid_connections),
                description="List of connections that failed validation"
            )
            print(f"\n📝 Invalid connections list saved to Airflow Variable 'invalid_connections_list'")
        
    except Exception as e:
        print(f"Error in validate_connections: {e}")
    finally:
        session.close()

def create_or_update_connections():
    """Create new connections or update existing ones based on configuration"""
    session = settings.Session()
    try:
        for conn in conns:
            try:
                existing_conn = session.query(Connection).filter(
                    Connection.conn_id == conn["conn_id"]
                ).first()
                
                if not existing_conn:
                    # Create new connection
                    new_conn = Connection(
                        conn_id=conn["conn_id"],
                        conn_type=conn["conn_type"],
                        host=conn.get("host"),
                        schema=conn.get("schema"),
                        login=conn.get("login"),
                        password=conn.get("password"),
                        port=conn.get("port"),
                        extra=conn.get("extra")
                    )
                    session.add(new_conn)
                    print(f"✓ Connection '{conn['conn_id']}' created.")
                else:
                    # Update existing connection
                    existing_conn.conn_type = conn["conn_type"]
                    existing_conn.host = conn.get("host")
                    existing_conn.schema = conn.get("schema")
                    existing_conn.login = conn.get("login")
                    existing_conn.password = conn.get("password")
                    existing_conn.port = conn.get("port")
                    if "extra" in conn:
                        existing_conn.extra = conn.get("extra")
                    print(f"✓ Connection '{conn['conn_id']}' updated.")
                    
            except Exception as e:
                print(f"✗ Error managing connection '{conn.get('conn_id', 'unknown')}': {e}")
                session.rollback()
                continue
                
        session.commit()
        print(f"Successfully processed {len(conns)} connection(s).")
    except Exception as e:
        print(f"Error in create_or_update_connections: {e}")
        session.rollback()
    finally:
        session.close()

def create_or_update_variables():
    """Create new variables or update existing ones based on configuration"""
    try:
        for var in variables:
            try:
                key = var["key"]
                value = var["value"]
                description = var.get("description", "")
                
                # Check if variable exists
                existing_var = Variable.get(key, default_var=None)
                
                if existing_var is None:
                    # Create new variable
                    Variable.set(key, value, description=description)
                    print(f"✓ Variable '{key}' created.")
                else:
                    # Update existing variable
                    Variable.set(key, value, description=description)
                    print(f"✓ Variable '{key}' updated.")
                    
            except Exception as e:
                print(f"✗ Error managing variable '{var.get('key', 'unknown')}': {e}")
                continue
                
        print(f"Successfully processed {len(variables)} variable(s).")
    except Exception as e:
        print(f"Error in create_or_update_variables: {e}")

# Execute the function once everyday
with DAG(
    dag_id='manage_connections_and_variables', 
    schedule_interval='@daily',
    default_args={
        "owner": "airflow",
        "retries": 1,
        "retry_delay": timedelta(minutes=5),
        "start_date": datetime(2023, 1, 17)
    },
    catchup=False,
    description='Manage Airflow Connections and Variables from configuration'
) as dag:

    clear_error_connections_task = PythonOperator(
        task_id='clear_error_connections',
        python_callable=clear_error_connections,
        op_kwargs={},
    )

    validate_connections_task = PythonOperator(
        task_id='validate_connections',
        python_callable=validate_connections,
        op_kwargs={},
    )

    create_or_update_connections_task = PythonOperator(
        task_id='create_or_update_connections',
        python_callable=create_or_update_connections,
        op_kwargs={},
    )

    create_or_update_variables_task = PythonOperator(
        task_id='create_or_update_variables',
        python_callable=create_or_update_variables,
        op_kwargs={},
    )

    # Define task dependencies
    # First clear errors, then create/update connections and variables in parallel, then validate
    clear_error_connections_task >> [create_or_update_connections_task, create_or_update_variables_task] >> validate_connections_task