import os
import psycopg2
import sys
# Add etl directory to path to import config
sys.path.append(os.path.join(os.path.dirname(__file__), "..", "etl"))
from config import TARGET_OMOP_CONFIG, VOCABULARY_DIR

def main():
    print("=== STARTING OMOP VOCABULARY LOADING ===")
    
    # Check if vocabulary directory exists
    vocab_dir = VOCABULARY_DIR
    if not os.path.exists(vocab_dir):
        print(f"Error: Vocabulary directory '{vocab_dir}' does not exist.")
        print("Please download vocabularies from Athena and place them in the vocabularies/ directory.")
        sys.exit(1)
        
    # List of vocabulary tables and their Athena files
    vocab_files = {
        "concept": "CONCEPT.csv",
        "vocabulary": "VOCABULARY.csv",
        "concept_relationship": "CONCEPT_RELATIONSHIP.csv",
        "relationship": "RELATIONSHIP.csv",
        "concept_synonym": "CONCEPT_SYNONYM.csv",
        "concept_ancestor": "CONCEPT_ANCESTOR.csv",
        "concept_class": "CONCEPT_CLASS.csv",
        "domain": "DOMAIN.csv",
        "drug_strength": "DRUG_STRENGTH.csv"
    }
    
    # Connect to Target OMOP PostgreSQL database
    try:
        conn = psycopg2.connect(
            host=TARGET_OMOP_CONFIG["host"],
            port=TARGET_OMOP_CONFIG["port"],
            user=TARGET_OMOP_CONFIG["user"],
            password=TARGET_OMOP_CONFIG["password"],
            database=TARGET_OMOP_CONFIG["database"]
        )
        cursor = conn.cursor()
        print("Connected to Target OMOP PostgreSQL database.")
    except Exception as e:
        print(f"Fatal Error: Cannot connect to PostgreSQL target database: {e}")
        sys.exit(1)
        
    # Disable triggers and constraints temporarily to allow loading in any order
    try:
        cursor.execute("SET session_replication_role = 'replica';")
        print("Session replication role set to 'replica' (triggers and constraints disabled).")
    except Exception as e:
        print(f"Warning: Could not disable constraints: {e}")
        
    for table, filename in vocab_files.items():
        filepath = os.path.join(vocab_dir, filename)
        if not os.path.exists(filepath):
            # Check lowercase as well
            filepath = os.path.join(vocab_dir, filename.lower())
            if not os.path.exists(filepath):
                print(f"Warning: File '{filename}' not found in vocabulary directory. Skipping.")
                continue
                
        print(f"Loading table '{table}' from {filename}...")
        
        # Truncate table first to prevent duplicate entries
        cursor.execute(f"TRUNCATE {table} CASCADE;")
        
        # Load using COPY command (copy_expert)
        # OHDSI Athena files are tab-separated, have headers, and use UTF-8
        # We specify CSV HEADER DELIMITER E'\t' QUOTE E'\b' to disable standard quoting
        sql = f"""
            COPY {table} 
            FROM STDIN 
            WITH (
                FORMAT csv, 
                HEADER true, 
                DELIMITER E'\t', 
                QUOTE E'\b', 
                ENCODING 'utf8'
            );
        """
        try:
            with open(filepath, 'r', encoding='utf-8', errors='ignore') as f:
                cursor.copy_expert(sql, f)
            conn.commit()
            print(f" - Successfully loaded '{table}'.")
        except Exception as e:
            print(f" - Error loading '{table}': {e}")
            conn.rollback()
            # Reset replica status in case rollback cleared it
            try:
                cursor.execute("SET session_replication_role = 'replica';")
                conn.commit()
            except Exception:
                pass
            
    # Reset triggers and constraints
    try:
        cursor.execute("SET session_replication_role = 'origin';")
        print("Session replication role reset to 'origin'.")
    except Exception as e:
        pass
        
    conn.commit()
    cursor.close()
    conn.close()
    print("=== VOCABULARY LOADING RUN COMPLETE ===")

if __name__ == "__main__":
    main()