import pymysql
import psycopg2
from psycopg2.extras import RealDictCursor, execute_values
import sys
from datetime import datetime, date
import hashlib
import re

# Import config and de-identification modules
from config import (
    SOURCE_MYSQL_CONFIG,
    SOURCE_POSTGRES_CONFIG,
    TARGET_OMOP_CONFIG,
    MIGRATION_CUTOFF_DATE,
    SALT
)
from deidentify import (
    get_deterministic_hash_int,
    get_patient_date_shift_days,
    shift_date,
    scrub_text
)

def get_mysql_connection():
    return pymysql.connect(
        host=SOURCE_MYSQL_CONFIG["host"],
        port=SOURCE_MYSQL_CONFIG["port"],
        user=SOURCE_MYSQL_CONFIG["user"],
        password=SOURCE_MYSQL_CONFIG["password"],
        database=SOURCE_MYSQL_CONFIG["database"],
        cursorclass=pymysql.cursors.DictCursor
    )

def get_postgres_source_connection():
    return psycopg2.connect(
        host=SOURCE_POSTGRES_CONFIG["host"],
        port=SOURCE_POSTGRES_CONFIG["port"],
        user=SOURCE_POSTGRES_CONFIG["user"],
        password=SOURCE_POSTGRES_CONFIG["password"],
        database=SOURCE_POSTGRES_CONFIG["database"]
    )

def get_postgres_target_connection():
    return 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"]
    )

def to_datetime(val):
    if val is None:
        return None
    if isinstance(val, datetime):
        return val
    if isinstance(val, date):
        return datetime(val.year, val.month, val.day)
    return val

def clean_date_val(val):
    if val is None:
        return None
    if isinstance(val, str):
        val_str = val.strip()
        if val_str == "" or val_str.startswith("0000") or "0000-00-00" in val_str:
            return None
        if re.match(r"^0\d{3}", val_str):
            return None
        
        # Parse and fix zero month/day patterns: YYYY-MM-DD
        match = re.match(r"^(\d{4})[-/](\d{2})[-/](\d{2})(.*)$", val_str)
        if match:
            year, month, day, rest = match.groups()
            if month == "00":
                month = "01"
            if day == "00":
                day = "01"
            val_str = f"{year}-{month}-{day}{rest}"
            
        return val_str
    if hasattr(val, "year"):
        if val.year <= 1000 or val.year > 2100:
            return None
    return val

def format_date_and_datetime(shifted_dt):
    if shifted_dt is None:
        return None, None
    if isinstance(shifted_dt, str):
        parts = shifted_dt.split()
        date_str = parts[0] if parts else shifted_dt
        return date_str, shifted_dt
    else:
        date_str = shifted_dt.strftime("%Y-%m-%d")
        if hasattr(shifted_dt, "hour"):
            datetime_str = shifted_dt.strftime("%Y-%m-%d %H:%M:%S")
        else:
            datetime_str = shifted_dt.strftime("%Y-%m-%d 00:00:00")
        return date_str, datetime_str

# Note parsers
def parse_diagnosis_note(note_text):
    if not note_text:
        return [], ""
        
    icd_codes = []
    
    # Normalise newlines
    note_norm = note_text.replace("<br>", "\n").replace("<br/>", "\n").replace("<div>", "\n").replace("</div>", "\n")
    
    # Locate "ICD Codes" section
    icd_match = re.search(r"ICD Codes:.*", note_norm, re.IGNORECASE | re.DOTALL)
    if icd_match:
        icd_section = icd_match.group(0)
        icd_codes = re.findall(r"\b[A-Z]\d{2}(?:\.\d+)?\b", icd_section)
        clinical_section = note_norm[:icd_match.start()].strip()
    else:
        clinical_section = note_norm.strip()
        icd_codes = re.findall(r"\b[A-Z]\d{2}(?:\.\d+)?\b", note_norm)
        
    # Clean up clinical section for free-text diagnosis
    diagnosis_desc = re.sub(r"<[^>]+>", " ", clinical_section)
    diagnosis_desc = re.sub(r"\s+", " ", diagnosis_desc).strip()
    
    return icd_codes, diagnosis_desc

def parse_prescription_note(note_text):
    if not note_text:
        return []
        
    prescriptions = []
    
    # Normalise newlines
    note_norm = note_text.replace("<br>", "\n").replace("<br/>", "\n").replace("<div>", "\n").replace("</div>", "\n")
    
    lines = note_norm.split("\n")
    for line in lines:
        line_clean = re.sub(r"<[^>]+>", "", line).strip()
        # Clean up tags and check if line matches drug keywords
        if line_clean and re.search(r"\b(tabs|caps|inj|syr|im|iv|mg|ml|g|bd|tds|stat|qds|amoxiclav|pcm|diclofenac|coartem|amoxicillin|flagyl|vitc|vasoprin|lisionopril|amlodipine|tabs\b|caps\b|dly|daily|x\d+|x\s*\d+|mls|ml\w*)\b", line_clean, re.IGNORECASE):
            line_clean = line_clean.lstrip(";:,.- ").strip()
            if line_clean:
                prescriptions.append(line_clean)
    return prescriptions

def parse_numeric(val):
    if val is None:
        return None
    if isinstance(val, (int, float)):
        return float(val)
    val_str = str(val).strip()
    match = re.search(r"^\d+(?:\.\d+)?", val_str)
    if match:
        try:
            return float(match.group(0))
        except ValueError:
            return None
    match = re.search(r"\d+(?:\.\d+)?", val_str)
    if match:
        try:
            return float(match.group(0))
        except ValueError:
            return None
    return None

def parse_bp(bp_str):
    if not bp_str:
        return None, None
    bp_clean = str(bp_str).strip()
    match = re.search(r"(\d+)\s*/\s*(\d+)", bp_clean)
    if match:
        try:
            return float(match.group(1)), float(match.group(2))
        except ValueError:
            return None, None
    return None, None

def execute_batched_insert(cursor, query, data_list, batch_size=5000):
    if not data_list:
        return
    total = len(data_list)
    print(f"  - Progress: 0 / {total} records inserted...", end="", flush=True)
    for i in range(0, total, batch_size):
        batch = data_list[i:i+batch_size]
        execute_values(cursor, query, batch)
        if (i + batch_size) < total:
            print(f"\r  - Progress: {i + len(batch)} / {total} records inserted...", end="", flush=True)
        else:
            print(f"\r  - Progress: {total} / {total} records inserted... done!", flush=True)

def main():
    print("=== STARTING OMOP CDM ETL PIPELINE ===")
    
    # 1. Connect to databases
    try:
        mysql_conn = get_mysql_connection()
        print("Connected to Source MySQL (Legacy DB)")
    except Exception as e:
        print(f"Warning: Cannot connect to MySQL legacy source: {e}. Skipping MySQL data extraction.")
        mysql_conn = None

    try:
        pg_src_conn = get_postgres_source_connection()
        print("Connected to Source PostgreSQL (Supabase DB)")
    except Exception as e:
        print(f"Fatal Error: Cannot connect to PostgreSQL source database: {e}")
        sys.exit(1)

    try:
        pg_tgt_conn = get_postgres_target_connection()
        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)

    pg_src_cursor = pg_src_conn.cursor(cursor_factory=RealDictCursor)
    pg_tgt_cursor = pg_tgt_conn.cursor()
    
    cutoff_dt = datetime.strptime(MIGRATION_CUTOFF_DATE, "%Y-%m-%d %H:%M:%S")
    
    processed_patient_keys = set()
    processed_visit_keys = set()
    processed_note_keys = set()
    processed_condition_keys = set()
    processed_drug_keys = set()
    processed_measurement_keys = set()
    
    loaded_counts = {
        "person": 0,
        "visit_occurrence": 0,
        "note": 0,
        "condition_occurrence": 0,
        "drug_exposure": 0,
        "measurement": 0
    }

    # ==========================================
    # STEP 1: MIGRATE PATIENTS (PERSON table)
    # ==========================================
    print("\n[ETL Step 1] Migrating patients to PERSON table...")
    person_inserts = []
    
    # A. Extract from MySQL
    if mysql_conn:
        try:
            mysql_cursor = mysql_conn.cursor()
            mysql_cursor.execute("""
                SELECT p.patient_id, u.gender, u.dob 
                FROM patients p 
                LEFT JOIN user_logins u ON p.login_id = u.login_id;
            """)
            for row in mysql_cursor.fetchall():
                orig_patient_id = str(row["patient_id"])
                person_id = get_deterministic_hash_int(orig_patient_id)
                if person_id in processed_patient_keys:
                    continue
                processed_patient_keys.add(person_id)
                
                gender_str = (row["gender"] or "").lower()
                dob = clean_date_val(row["dob"])
                shift_days = get_patient_date_shift_days(orig_patient_id)
                shifted_dob = shift_date(dob, shift_days)
                dob_date_str, dob_datetime_str = format_date_and_datetime(shifted_dob)
                
                year_of_birth = 1900
                month_of_birth = None
                day_of_birth = None
                if shifted_dob:
                    if hasattr(shifted_dob, "year"):
                        year_of_birth = shifted_dob.year
                        month_of_birth = shifted_dob.month
                        day_of_birth = shifted_dob.day
                
                gender_concept_id = 8551
                if "f" in gender_str:
                    gender_concept_id = 8532
                elif "m" in gender_str:
                    gender_concept_id = 8507
                    
                person_source_value = hashlib.sha256(f"{orig_patient_id}{SALT}".encode()).hexdigest()[:20]
                person_inserts.append((
                    person_id, gender_concept_id, year_of_birth, month_of_birth, day_of_birth,
                    dob_datetime_str, 8552, 0, person_source_value, gender_str[:50], 0, "Black / African", 0, "Not Hispanic", 0
                ))
            mysql_cursor.close()
        except Exception as e:
            print(f"Error reading patients from MySQL: {e}")
            sys.exit(1)

    # B. Extract from Postgres source
    try:
        pg_src_cursor.execute("""
            SELECT p.patient_id, u.gender, u.dob 
            FROM hms.patients p
            LEFT JOIN hms.user_logins u ON p.login_id = u.login_id;
        """)
        for row in pg_src_cursor.fetchall():
            orig_patient_id = str(row["patient_id"])
            person_id = get_deterministic_hash_int(orig_patient_id)
            if person_id in processed_patient_keys:
                continue
            processed_patient_keys.add(person_id)
            
            gender_str = (row["gender"] or "").lower()
            dob = clean_date_val(row["dob"])
            shift_days = get_patient_date_shift_days(orig_patient_id)
            shifted_dob = shift_date(dob, shift_days)
            dob_date_str, dob_datetime_str = format_date_and_datetime(shifted_dob)
            
            year_of_birth = 1900
            month_of_birth = None
            day_of_birth = None
            if shifted_dob:
                if hasattr(shifted_dob, "year"):
                    year_of_birth = shifted_dob.year
                    month_of_birth = shifted_dob.month
                    day_of_birth = shifted_dob.day
            
            gender_concept_id = 8551
            if "f" in gender_str:
                gender_concept_id = 8532
            elif "m" in gender_str:
                gender_concept_id = 8507
                
            person_source_value = hashlib.sha256(f"{orig_patient_id}{SALT}".encode()).hexdigest()[:20]
            person_inserts.append((
                person_id, gender_concept_id, year_of_birth, month_of_birth, day_of_birth,
                dob_datetime_str, 8552, 0, person_source_value, gender_str[:50], 0, "Black / African", 0, "Not Hispanic", 0
            ))
    except Exception as e:
        print(f"Error reading patients from Postgres source: {e}")
        sys.exit(1)

    if person_inserts:
        print(f"Inserting {len(person_inserts)} patients into PERSON...")
        execute_batched_insert(pg_tgt_cursor, """
            INSERT INTO person (
                person_id, gender_concept_id, year_of_birth, month_of_birth, day_of_birth,
                birth_datetime, race_concept_id, ethnicity_concept_id, person_source_value,
                gender_source_value, gender_source_concept_id, race_source_value, race_source_concept_id,
                ethnicity_source_value, ethnicity_source_concept_id
            ) VALUES %s
            ON CONFLICT (person_id) DO NOTHING;
        """, person_inserts)
        loaded_counts["person"] = len(person_inserts)
    
    pg_tgt_conn.commit()
    print(f"Done. Current person keys processed: {len(processed_patient_keys)}")

    # ==========================================
    # STEP 2: MIGRATE VISITS (VISIT_OCCURRENCE table)
    # ==========================================
    print("\n[ETL Step 2] Migrating admissions to VISIT_OCCURRENCE...")
    visit_inserts = []

    def process_admission_row(adm_id, pat_id, startdate, enddate, date_created, is_history=False, is_mysql=False, adm_type="", action=""):
        actual_start = startdate if startdate else date_created
        if not actual_start:
            return None
        actual_end = enddate if enddate else actual_start
        
        if is_mysql:
            compare_val = to_datetime(actual_start)
            if compare_val and compare_val >= cutoff_dt:
                return None
                
        person_id = get_deterministic_hash_int(pat_id)
        if person_id not in processed_patient_keys:
            return None
            
        visit_key = f"hist_{adm_id}" if is_history else f"adm_{adm_id}"
        visit_occurrence_id = get_deterministic_hash_int(visit_key)
        if visit_occurrence_id in processed_visit_keys:
            return None
        processed_visit_keys.add(visit_occurrence_id)
        
        shift_days = get_patient_date_shift_days(pat_id)
        start_date = shift_date(actual_start, shift_days)
        end_date = shift_date(actual_end, shift_days)
        
        start_date_only, start_datetime = format_date_and_datetime(start_date)
        end_date_only, end_datetime = format_date_and_datetime(end_date)
        
        return (
            visit_occurrence_id, person_id, 9201, start_date_only, start_datetime,
            end_date_only, end_datetime, 44818518, 0, 0, visit_key[:50], 0, 0, adm_type[:50], 0, action[:50], None
        )

    # A. MySQL Admissions
    if mysql_conn:
        try:
            mysql_cursor = mysql_conn.cursor()
            mysql_cursor.execute("SELECT admission_id, patient_id, startdate, enddate, date_created FROM admissions WHERE patient_id IS NOT NULL")
            for row in mysql_cursor.fetchall():
                res = process_admission_row(
                    str(row["admission_id"]), str(row["patient_id"]), 
                    clean_date_val(row["startdate"]), clean_date_val(row["enddate"]), 
                    clean_date_val(row["date_created"]), is_mysql=True
                )
                if res: visit_inserts.append(res)
                
            mysql_cursor.execute("SELECT admissions_history_id, patient_id, date_admitted, date_discharged, created_at, admission_type, action FROM admissions_history WHERE patient_id IS NOT NULL")
            for row in mysql_cursor.fetchall():
                res = process_admission_row(
                    str(row["admissions_history_id"]), str(row["patient_id"]),
                    clean_date_val(row["date_admitted"]), clean_date_val(row["date_discharged"]),
                    clean_date_val(row["created_at"]), is_history=True, is_mysql=True,
                    adm_type=row["admission_type"] or "", action=row["action"] or ""
                )
                if res: visit_inserts.append(res)
            mysql_cursor.close()
        except Exception as e:
            print(f"Error reading visits from MySQL: {e}")
            sys.exit(1)

    # B. Postgres Admissions
    try:
        pg_src_cursor.execute("SELECT admission_id, patient_id, startdate, enddate, date_created FROM hms.admissions WHERE patient_id IS NOT NULL")
        for row in pg_src_cursor.fetchall():
            res = process_admission_row(
                str(row["admission_id"]), str(row["patient_id"]), 
                clean_date_val(row["startdate"]), clean_date_val(row["enddate"]), 
                clean_date_val(row["date_created"])
            )
            if res: visit_inserts.append(res)
            
        pg_src_cursor.execute("SELECT admissions_history_id, patient_id, date_admitted, date_discharged, created_at, admission_type, action FROM hms.admissions_history WHERE patient_id IS NOT NULL")
        for row in pg_src_cursor.fetchall():
            res = process_admission_row(
                str(row["admissions_history_id"]), str(row["patient_id"]),
                clean_date_val(row["date_admitted"]), clean_date_val(row["date_discharged"]),
                clean_date_val(row["created_at"]), is_history=True,
                adm_type=row["admission_type"] or "", action=row["action"] or ""
            )
            if res: visit_inserts.append(res)
    except Exception as e:
        print(f"Error reading visits from Postgres source: {e}")
        sys.exit(1)

    if visit_inserts:
        print(f"Inserting {len(visit_inserts)} visits into VISIT_OCCURRENCE...")
        execute_batched_insert(pg_tgt_cursor, """
            INSERT INTO visit_occurrence (
                visit_occurrence_id, person_id, visit_concept_id, visit_start_date, visit_start_datetime,
                visit_end_date, visit_end_datetime, visit_type_concept_id, provider_id, care_site_id,
                visit_source_value, visit_source_concept_id, admitted_from_concept_id, 
                admitted_from_source_value, discharged_to_concept_id, discharged_to_source_value,
                preceding_visit_occurrence_id
            ) VALUES %s
            ON CONFLICT (visit_occurrence_id) DO NOTHING;
        """, visit_inserts)
        loaded_counts["visit_occurrence"] = len(visit_inserts)
        
    pg_tgt_conn.commit()
    print("Done.")

    # ==========================================
    # STEP 3: MIGRATE NOTES (NOTE table)
    # ==========================================
    print("\n[ETL Step 3] Migrating clinical notes to NOTE table (with NER sanitization)...")
    note_inserts = []

    def process_note_row(note_id_raw, pat_id, date_created, note_content, note_type, table_prefix, is_mysql=False):
        if not note_content or str(note_content).strip() == "":
            return None
        if not date_created:
            return None
        compare_val = to_datetime(date_created)
        if is_mysql and compare_val and compare_val >= cutoff_dt:
            return None
            
        person_id = get_deterministic_hash_int(pat_id)
        if person_id not in processed_patient_keys:
            return None
            
        note_key = f"{table_prefix}_{note_id_raw}"
        note_id = get_deterministic_hash_int(note_key)
        if note_id in processed_note_keys:
            return None
        processed_note_keys.add(note_id)
        
        shift_days = get_patient_date_shift_days(pat_id)
        note_date_shifted = shift_date(date_created, shift_days)
        note_date_only, note_datetime = format_date_and_datetime(note_date_shifted)
        
        scrubbed_note = scrub_text(note_content)
        
        note_class_concept_id = 44814647 # General
        if "doctor" in table_prefix:
            note_class_concept_id = 44814645
        elif "nurse" in table_prefix:
            note_class_concept_id = 44814646
            
        return (
            note_id, person_id, note_date_only, note_datetime, note_class_concept_id,
            note_class_concept_id, note_type[:249], scrubbed_note, 0, 4180186, None, None, None, note_key[:50], None, None
        )

    # A. MySQL Notes
    if mysql_conn:
        try:
            mysql_cursor = mysql_conn.cursor()
            # doctor_Notes
            mysql_cursor.execute("SELECT id, patient_id, date_created, note, note_type FROM doctor_Notes WHERE patient_id IS NOT NULL")
            for row in mysql_cursor.fetchall():
                res = process_note_row(str(row["id"]), str(row["patient_id"]), clean_date_val(row["date_created"]), row["note"], row["note_type"] or "Doctor Note", "doctor", is_mysql=True)
                if res: note_inserts.append(res)
            # nurse_Notes
            mysql_cursor.execute("SELECT id, patient_id, date_created, note, diagnosis, allergies, dispensed_medication FROM nurse_Notes WHERE patient_id IS NOT NULL")
            for row in mysql_cursor.fetchall():
                combined = f"Note: {row['note'] or ''}\nDiagnosis: {row['diagnosis'] or ''}\nAllergies: {row['allergies'] or ''}\nDispensed: {row['dispensed_medication'] or ''}"
                res = process_note_row(str(row["id"]), str(row["patient_id"]), clean_date_val(row["date_created"]), combined, "Nurse Note", "nurse", is_mysql=True)
                if res: note_inserts.append(res)
            # general notes
            mysql_cursor.execute("SELECT id, patient_id, date_created, note, note_type FROM notes WHERE patient_id IS NOT NULL")
            for row in mysql_cursor.fetchall():
                res = process_note_row(str(row["id"]), str(row["patient_id"]), clean_date_val(row["date_created"]), row["note"], row["note_type"] or "General Note", "notes", is_mysql=True)
                if res: note_inserts.append(res)
            mysql_cursor.close()
        except Exception as e:
            print(f"Error reading notes from MySQL: {e}")
            sys.exit(1)

    # B. Postgres Notes
    try:
        # doctor_notes
        pg_src_cursor.execute("SELECT id, patient_id, date_created, note, note_type FROM hms.doctor_notes WHERE patient_id IS NOT NULL")
        for row in pg_src_cursor.fetchall():
            res = process_note_row(str(row["id"]), str(row["patient_id"]), clean_date_val(row["date_created"]), row["note"], row["note_type"] or "Doctor Note", "doctor")
            if res: note_inserts.append(res)
        # nurse_notes
        pg_src_cursor.execute("SELECT id, patient_id, date_created, note, diagnosis, allergies, dispensed_medication FROM hms.nurse_notes WHERE patient_id IS NOT NULL")
        for row in pg_src_cursor.fetchall():
            combined = f"Note: {row['note'] or ''}\nDiagnosis: {row['diagnosis'] or ''}\nAllergies: {row['allergies'] or ''}\nDispensed: {row['dispensed_medication'] or ''}"
            res = process_note_row(str(row["id"]), str(row["patient_id"]), clean_date_val(row["date_created"]), combined, "Nurse Note", "nurse")
            if res: note_inserts.append(res)
        # general notes
        pg_src_cursor.execute("SELECT id, patient_id, date_created, note, note_type FROM hms.notes WHERE patient_id IS NOT NULL")
        for row in pg_src_cursor.fetchall():
            res = process_note_row(str(row["id"]), str(row["patient_id"]), clean_date_val(row["date_created"]), row["note"], row["note_type"] or "General Note", "notes")
            if res: note_inserts.append(res)
    except Exception as e:
        print(f"Error reading notes from Postgres source: {e}")
        sys.exit(1)

    if note_inserts:
        print(f"Inserting {len(note_inserts)} notes into NOTE...")
        execute_batched_insert(pg_tgt_cursor, """
            INSERT INTO note (
                note_id, person_id, note_date, note_datetime, note_type_concept_id,
                note_class_concept_id, note_title, note_text, encoding_concept_id,
                language_concept_id, provider_id, visit_occurrence_id, visit_detail_id,
                note_source_value, note_event_id, note_event_field_concept_id
            ) VALUES %s
            ON CONFLICT (note_id) DO NOTHING;
        """, note_inserts, batch_size=500)
        loaded_counts["note"] = len(note_inserts)
        
    pg_tgt_conn.commit()
    print("Done.")

    # ==========================================
    # STEP 4: MIGRATE DIAGNOSES (CONDITION_OCCURRENCE)
    # ==========================================
    print("\n[ETL Step 4] Migrating diagnoses to CONDITION_OCCURRENCE...")
    condition_inserts = []

    def process_diag_row(diag_id, pat_id, diagnosis_text, code_str, date_created, is_mysql=False):
        if not date_created:
            return
        compare_val = to_datetime(date_created)
        if is_mysql and compare_val and compare_val >= cutoff_dt:
            return
            
        person_id = get_deterministic_hash_int(pat_id)
        if person_id not in processed_patient_keys:
            return
            
        shift_days = get_patient_date_shift_days(pat_id)
        shifted_dt = shift_date(date_created, shift_days)
        start_date_only, start_datetime = format_date_and_datetime(shifted_dt)
        
        # Combine diagnosis and code into source value
        source_val = f"{diagnosis_text or ''} [{code_str or ''}]".strip()
        if source_val == "":
            source_val = "EHR Diagnosis"
            
        cond_key = f"struct_{diag_id}"
        cond_id = get_deterministic_hash_int(cond_key)
        if cond_id not in processed_condition_keys:
            processed_condition_keys.add(cond_id)
            condition_inserts.append((
                cond_id, person_id, 0, start_date_only, start_datetime, None, None, 32817, None, None, None, None, None, source_val[:50], 0, None
            ))

    def process_note_diag(note_id, pat_id, note_text, date_created, is_mysql=False):
        if not date_created:
            return
        compare_val = to_datetime(date_created)
        if is_mysql and compare_val and compare_val >= cutoff_dt:
            return
            
        person_id = get_deterministic_hash_int(pat_id)
        if person_id not in processed_patient_keys:
            return
            
        shift_days = get_patient_date_shift_days(pat_id)
        shifted_dt = shift_date(date_created, shift_days)
        start_date_only, start_datetime = format_date_and_datetime(shifted_dt)
        
        icd_codes, diagnosis_desc = parse_diagnosis_note(note_text)
        
        # Parse ICD codes
        for code in icd_codes:
            cond_key = f"note_icd_{note_id}_{code}"
            cond_id = get_deterministic_hash_int(cond_key)
            if cond_id not in processed_condition_keys:
                processed_condition_keys.add(cond_id)
                condition_inserts.append((
                    cond_id, person_id, 0, start_date_only, start_datetime, None, None, 32817, None, None, None, None, None, code[:50], 0, None
                ))
                
        # Parse free text description if no ICD code or as secondary
        if diagnosis_desc and diagnosis_desc != "":
            cond_key = f"note_text_{note_id}"
            cond_id = get_deterministic_hash_int(cond_key)
            if cond_id not in processed_condition_keys:
                processed_condition_keys.add(cond_id)
                condition_inserts.append((
                    cond_id, person_id, 0, start_date_only, start_datetime, None, None, 32817, None, None, None, None, None, diagnosis_desc[:50], 0, None
                ))

    # A. MySQL Diagnoses
    if mysql_conn:
        try:
            mysql_cursor = mysql_conn.cursor()
            # Structured
            mysql_cursor.execute("SELECT diagnosis_id, patient_id, diagnosis, diagnosis_codes, date_created FROM medikal_patients_diagnosis WHERE patient_id IS NOT NULL")
            for row in mysql_cursor.fetchall():
                process_diag_row(
                    str(row["diagnosis_id"]), str(row["patient_id"]), 
                    row["diagnosis"], row["diagnosis_codes"], 
                    clean_date_val(row["date_created"]), is_mysql=True
                )
            # Unstructured notes (where note_type = 'diagnosis')
            mysql_cursor.execute("SELECT id, patient_id, note, date_created FROM notes WHERE patient_id IS NOT NULL AND note_type = 'diagnosis'")
            for row in mysql_cursor.fetchall():
                process_note_diag(str(row["id"]), str(row["patient_id"]), row["note"], clean_date_val(row["date_created"]), is_mysql=True)
            mysql_cursor.close()
        except Exception as e:
            print(f"Error reading diagnoses from MySQL: {e}")
            sys.exit(1)

    # B. Postgres Diagnoses
    try:
        # Structured
        pg_src_cursor.execute("SELECT diagnosis_id, patient_id, diagnosis, diagnosis_codes, date_created FROM hms.medikal_patients_diagnosis WHERE patient_id IS NOT NULL")
        for row in pg_src_cursor.fetchall():
            process_diag_row(
                str(row["diagnosis_id"]), str(row["patient_id"]), 
                row["diagnosis"], row["diagnosis_codes"], 
                clean_date_val(row["date_created"])
            )
        # Unstructured notes
        pg_src_cursor.execute("SELECT id, patient_id, note, date_created FROM hms.notes WHERE patient_id IS NOT NULL AND note_type = 'diagnosis'")
        for row in pg_src_cursor.fetchall():
            process_note_diag(str(row["id"]), str(row["patient_id"]), row["note"], clean_date_val(row["date_created"]))
    except Exception as e:
        print(f"Error reading diagnoses from Postgres source: {e}")
        sys.exit(1)

    if condition_inserts:
        print(f"Inserting {len(condition_inserts)} diagnoses into CONDITION_OCCURRENCE...")
        execute_batched_insert(pg_tgt_cursor, """
            INSERT INTO condition_occurrence (
                condition_occurrence_id, person_id, condition_concept_id, condition_start_date, condition_start_datetime,
                condition_end_date, condition_end_datetime, condition_type_concept_id, condition_status_concept_id,
                stop_reason, provider_id, visit_occurrence_id, visit_detail_id, condition_source_value,
                condition_source_concept_id, condition_status_source_value
            ) VALUES %s
            ON CONFLICT (condition_occurrence_id) DO NOTHING;
        """, condition_inserts)
        loaded_counts["condition_occurrence"] = len(condition_inserts)
        
    pg_tgt_conn.commit()
    print("Done.")

    # ==========================================
    # STEP 5: MIGRATE PRESCRIPTIONS (DRUG_EXPOSURE)
    # ==========================================
    print("\n[ETL Step 5] Migrating prescriptions to DRUG_EXPOSURE...")
    drug_inserts = []

    def process_details_row(det_id, pat_id, drug_name, dosage_val, measure_str, startdate, enddate, date_created, is_mysql=False):
        actual_start = startdate if startdate else date_created
        if not actual_start:
            return
        compare_val = to_datetime(actual_start)
        if is_mysql and compare_val and compare_val >= cutoff_dt:
            return
            
        person_id = get_deterministic_hash_int(pat_id)
        if person_id not in processed_patient_keys:
            return
            
        shift_days = get_patient_date_shift_days(pat_id)
        start_shifted = shift_date(actual_start, shift_days)
        end_shifted = shift_date(enddate if enddate else actual_start, shift_days)
        
        start_date_only, start_datetime = format_date_and_datetime(start_shifted)
        end_date_only, end_datetime = format_date_and_datetime(end_shifted)
        
        qty = parse_numeric(dosage_val)
        drug_key = f"details_{det_id}"
        drug_id = get_deterministic_hash_int(drug_key)
        
        if drug_id not in processed_drug_keys:
            processed_drug_keys.add(drug_id)
            drug_inserts.append((
                drug_id, person_id, 0, start_date_only, start_datetime, end_date_only, end_datetime, None,
                32838, None, None, qty, None, f"{dosage_val or ''} {measure_str or ''}", None, None, None, None, None,
                drug_name[:50], 0, None, measure_str[:50] if measure_str else None
            ))

    def process_schedule_row(sch_id, pat_id, drug_name, qty_val, dosage_val, measure_str, entry_date, note_str, is_mysql=False):
        if not entry_date:
            return
        compare_val = to_datetime(entry_date)
        if is_mysql and compare_val and compare_val >= cutoff_dt:
            return
            
        person_id = get_deterministic_hash_int(pat_id)
        if person_id not in processed_patient_keys:
            return
            
        shift_days = get_patient_date_shift_days(pat_id)
        shifted_dt = shift_date(entry_date, shift_days)
        date_only, datetime_str = format_date_and_datetime(shifted_dt)
        
        qty = parse_numeric(qty_val)
        sig = note_str or f"Dosage: {dosage_val or ''} {measure_str or ''}"
        
        drug_key = f"schedule_{sch_id}"
        drug_id = get_deterministic_hash_int(drug_key)
        
        if drug_id not in processed_drug_keys:
            processed_drug_keys.add(drug_id)
            drug_inserts.append((
                drug_id, person_id, 0, date_only, datetime_str, date_only, datetime_str, None,
                32818, None, None, qty, None, sig, None, None, None, None, None,
                drug_name[:50] if drug_name else "Unknown Drug", 0, None, measure_str[:50] if measure_str else None
            ))

    def process_presc_note(presc_id, pat_id, note_text, prescription_date, is_mysql=False):
        if not prescription_date:
            return
        compare_val = to_datetime(prescription_date)
        if is_mysql and compare_val and compare_val >= cutoff_dt:
            return
            
        person_id = get_deterministic_hash_int(pat_id)
        if person_id not in processed_patient_keys:
            return
            
        shift_days = get_patient_date_shift_days(pat_id)
        shifted_dt = shift_date(prescription_date, shift_days)
        date_only, datetime_str = format_date_and_datetime(shifted_dt)
        
        rx_lines = parse_prescription_note(note_text)
        for idx, line in enumerate(rx_lines):
            drug_key = f"presc_note_{presc_id}_{idx}"
            drug_id = get_deterministic_hash_int(drug_key)
            if drug_id not in processed_drug_keys:
                processed_drug_keys.add(drug_id)
                drug_inserts.append((
                    drug_id, person_id, 0, date_only, datetime_str, date_only, datetime_str, None,
                    32838, None, None, None, None, line, None, None, None, None, None,
                    line[:50], 0, None, None
                ))

    # A. MySQL Prescriptions
    if mysql_conn:
        try:
            mysql_cursor = mysql_conn.cursor()
            # Structured Details
            mysql_cursor.execute("""
                SELECT d.prescription_details_id, p.patient_id, d.drug_code, d.dosage, d.measure, d.startdate, d.enddate, d.date_created 
                FROM medikal_prescription_details d 
                JOIN medikal_prescriptions p ON d.prescription_id = p.prescription_id 
                WHERE p.patient_id IS NOT NULL;
            """)
            for row in mysql_cursor.fetchall():
                process_details_row(
                    str(row["prescription_details_id"]), str(row["patient_id"]), 
                    row["drug_code"] or "Unknown", row["dosage"], row["measure"], 
                    clean_date_val(row["startdate"]), clean_date_val(row["enddate"]), 
                    clean_date_val(row["date_created"]), is_mysql=True
                )
            # Structured Schedules
            mysql_cursor.execute("SELECT id, patient_id, item, quantity, dosage, measure, entry_date, note FROM medikal_prescription_schedule WHERE patient_id IS NOT NULL")
            for row in mysql_cursor.fetchall():
                process_schedule_row(
                    str(row["id"]), str(row["patient_id"]), row["item"], 
                    row["quantity"], row["dosage"], row["measure"], 
                    clean_date_val(row["entry_date"]), row["note"], is_mysql=True
                )
            # Free-text prescription notes
            mysql_cursor.execute("SELECT prescription_id, patient_id, prescription_date, note FROM medikal_prescriptions WHERE patient_id IS NOT NULL AND note IS NOT NULL AND note != ''")
            for row in mysql_cursor.fetchall():
                process_presc_note(
                    str(row["prescription_id"]), str(row["patient_id"]), 
                    row["note"], clean_date_val(row["prescription_date"]), is_mysql=True
                )
            mysql_cursor.close()
        except Exception as e:
            print(f"Error reading prescriptions from MySQL: {e}")
            sys.exit(1)

    # B. Postgres Prescriptions
    try:
        # Structured Details
        pg_src_cursor.execute("""
            SELECT d.prescription_details_id, p.patient_id, d.drug_code, d.dosage, d.measure, d.startdate, d.enddate, d.date_created 
            FROM hms.medikal_prescription_details d 
            JOIN hms.medikal_prescriptions p ON d.prescription_id = p.prescription_id 
            WHERE p.patient_id IS NOT NULL;
        """)
        for row in pg_src_cursor.fetchall():
            process_details_row(
                str(row["prescription_details_id"]), str(row["patient_id"]), 
                row["drug_code"] or "Unknown", row["dosage"], row["measure"], 
                clean_date_val(row["startdate"]), clean_date_val(row["enddate"]), 
                clean_date_val(row["date_created"])
            )
        # Structured Schedules
        pg_src_cursor.execute("SELECT id, patient_id, item, quantity, dosage, measure, entry_date, note FROM hms.medikal_prescription_schedule WHERE patient_id IS NOT NULL")
        for row in pg_src_cursor.fetchall():
            process_schedule_row(
                str(row["id"]), str(row["patient_id"]), row["item"], 
                row["quantity"], row["dosage"], row["measure"], 
                clean_date_val(row["entry_date"]), row["note"]
            )
        # Free-text prescription notes
        pg_src_cursor.execute("SELECT prescription_id, patient_id, prescription_date, note FROM hms.medikal_prescriptions WHERE patient_id IS NOT NULL AND note IS NOT NULL AND note != ''")
        for row in pg_src_cursor.fetchall():
            process_presc_note(
                str(row["prescription_id"]), str(row["patient_id"]), 
                row["note"], clean_date_val(row["prescription_date"])
            )
    except Exception as e:
        print(f"Error reading prescriptions from Postgres source: {e}")
        sys.exit(1)

    if drug_inserts:
        print(f"Inserting {len(drug_inserts)} prescriptions into DRUG_EXPOSURE...")
        execute_batched_insert(pg_tgt_cursor, """
            INSERT INTO drug_exposure (
                drug_exposure_id, person_id, drug_concept_id, drug_exposure_start_date, drug_exposure_start_datetime,
                drug_exposure_end_date, drug_exposure_end_datetime, verbatim_end_date, drug_type_concept_id,
                stop_reason, refills, quantity, days_supply, sig, route_concept_id, lot_number,
                provider_id, visit_occurrence_id, visit_detail_id, drug_source_value, drug_source_concept_id,
                route_source_value, dose_unit_source_value
            ) VALUES %s
            ON CONFLICT (drug_exposure_id) DO NOTHING;
        """, drug_inserts)
        loaded_counts["drug_exposure"] = len(drug_inserts)
        
    pg_tgt_conn.commit()
    print("Done.")

    # ==========================================
    # STEP 6: MIGRATE VITALS & LABS (MEASUREMENT)
    # ==========================================
    print("\n[ETL Step 6] Migrating vitals and labs to MEASUREMENT...")
    measurement_inserts = []

    # A. Vitals
    def process_vital_reading(vital_id, pat_id, reading_date, concept_id, value_raw, param_name, is_mysql=False):
        if not value_raw or str(value_raw).strip() == "":
            return
        if not reading_date:
            return
            
        compare_val = to_datetime(reading_date)
        if is_mysql and compare_val and compare_val >= cutoff_dt:
            return
            
        person_id = get_deterministic_hash_int(pat_id)
        if person_id not in processed_patient_keys:
            return
            
        val_numeric = parse_numeric(value_raw)
        
        shift_days = get_patient_date_shift_days(pat_id)
        shifted_dt = shift_date(reading_date, shift_days)
        date_only, datetime_str = format_date_and_datetime(shifted_dt)
        
        meas_key = f"vital_{vital_id}_{param_name}"
        meas_id = get_deterministic_hash_int(meas_key)
        
        meas_val_str = str(value_raw)[:50] if value_raw is not None else ""
        if meas_id not in processed_measurement_keys:
            processed_measurement_keys.add(meas_id)
            measurement_inserts.append((
                meas_id, person_id, concept_id, date_only, datetime_str, None, 44818702, None,
                val_numeric, None, None, None, None, None, None, None,
                meas_val_str, 0, None, None, meas_val_str, None, None
            ))

    def process_vital_row(vital_id, pat_id, reading_date, temp, hgt, wgt, bp, pulse, resp, o2, sugar, is_mysql=False):
        # Temp
        process_vital_reading(vital_id, pat_id, reading_date, 3020891, temp, "temp", is_mysql)
        # Height
        process_vital_reading(vital_id, pat_id, reading_date, 3036277, hgt, "height", is_mysql)
        # Weight
        process_vital_reading(vital_id, pat_id, reading_date, 3025315, wgt, "weight", is_mysql)
        # Pulse
        process_vital_reading(vital_id, pat_id, reading_date, 3027018, pulse, "pulse", is_mysql)
        # Respiration
        process_vital_reading(vital_id, pat_id, reading_date, 3024171, resp, "resp", is_mysql)
        # Oxygen
        process_vital_reading(vital_id, pat_id, reading_date, 40762636, o2, "o2", is_mysql)
        # Blood Sugar
        process_vital_reading(vital_id, pat_id, reading_date, 3004501, sugar, "sugar", is_mysql)
        # BP
        sbp_val, dbp_val = parse_bp(bp)
        if sbp_val is not None:
            process_vital_reading(vital_id, pat_id, reading_date, 3004249, sbp_val, "sbp", is_mysql)
        if dbp_val is not None:
            process_vital_reading(vital_id, pat_id, reading_date, 3012888, dbp_val, "dbp", is_mysql)

    # A1. MySQL Vitals
    if mysql_conn:
        try:
            mysql_cursor = mysql_conn.cursor()
            mysql_cursor.execute("SELECT id, patient_id, reading_date, temperature, height, weight, bp, pulse, respiration, oxygen, blood_sugar FROM vitalsigns WHERE patient_id IS NOT NULL")
            print("Processing MySQL vitals...")
            for idx, row in enumerate(mysql_cursor.fetchall()):
                if idx > 0 and idx % 50000 == 0:
                    print(f" - Processed {idx} vitals rows...")
                process_vital_row(
                    str(row["id"]), str(row["patient_id"]), clean_date_val(row["reading_date"]),
                    row["temperature"], row["height"], row["weight"], row["bp"], row["pulse"],
                    row["respiration"], row["oxygen"], row["blood_sugar"], is_mysql=True
                )
            mysql_cursor.close()
        except Exception as e:
            print(f"Error reading vitals from MySQL: {e}")
            sys.exit(1)

    # A2. Postgres Vitals
    try:
        pg_src_cursor.execute("SELECT id, patient_id, reading_date, temperature, height, weight, bp, pulse, respiration, oxygen, blood_sugar FROM hms.vitalsigns WHERE patient_id IS NOT NULL")
        for row in pg_src_cursor.fetchall():
            process_vital_row(
                str(row["id"]), str(row["patient_id"]), clean_date_val(row["reading_date"]),
                row["temperature"], row["height"], row["weight"], row["bp"], row["pulse"],
                row["respiration"], row["oxygen"], row["blood_sugar"]
            )
    except Exception as e:
        print(f"Error reading vitals from Postgres source: {e}")
        sys.exit(1)

    # B. Labs
    def process_lab_row(det_id, pat_id, invest, result, units, date_created, trans_date, is_mysql=False):
        actual_date = date_created if date_created else trans_date
        if not actual_date:
            return
        compare_val = to_datetime(actual_date)
        if is_mysql and compare_val and compare_val >= cutoff_dt:
            return
            
        person_id = get_deterministic_hash_int(pat_id)
        if person_id not in processed_patient_keys:
            return
            
        val_numeric = parse_numeric(result)
        
        shift_days = get_patient_date_shift_days(pat_id)
        shifted_dt = shift_date(actual_date, shift_days)
        date_only, datetime_str = format_date_and_datetime(shifted_dt)
        
        meas_key = f"lab_{det_id}"
        meas_id = get_deterministic_hash_int(meas_key)
        
        if meas_id not in processed_measurement_keys:
            processed_measurement_keys.add(meas_id)
            measurement_inserts.append((
                meas_id, person_id, 0, date_only, datetime_str, None, 32856, None,
                val_numeric, None, None, None, None, None, None, None,
                invest[:50] if invest else "Lab", 0, units[:50] if units else None, None, result[:50] if result else None, None, None
            ))

    # B1. MySQL Labs
    if mysql_conn:
        try:
            mysql_cursor = mysql_conn.cursor()
            mysql_cursor.execute("""
                SELECT d.labdata_details_id, d.patient_id, d.investigation, d.result, d.units, d.date_created, l.trans_date 
                FROM labdata_details d 
                JOIN labdata l ON d.lab_id = l.lab_id 
                WHERE d.patient_id IS NOT NULL;
            """)
            print("Processing MySQL labs...")
            for idx, row in enumerate(mysql_cursor.fetchall()):
                if idx > 0 and idx % 100000 == 0:
                    print(f" - Processed {idx} lab details...")
                process_lab_row(
                    str(row["labdata_details_id"]), str(row["patient_id"]), 
                    row["investigation"], row["result"], row["units"], 
                    clean_date_val(row["date_created"]), clean_date_val(row["trans_date"]), is_mysql=True
                )
            mysql_cursor.close()
        except Exception as e:
            print(f"Error reading labs from MySQL: {e}")
            sys.exit(1)

    # B2. Postgres Labs
    try:
        pg_src_cursor.execute("""
            SELECT d.labdata_details_id, d.patient_id, d.investigation, d.result, d.units, d.date_created, l.trans_date 
            FROM hms.labdata_details d 
            JOIN hms.labdata l ON d.lab_id = l.lab_id 
            WHERE d.patient_id IS NOT NULL;
        """)
        for row in pg_src_cursor.fetchall():
            process_lab_row(
                str(row["labdata_details_id"]), str(row["patient_id"]), 
                row["investigation"], row["result"], row["units"], 
                clean_date_val(row["date_created"]), clean_date_val(row["trans_date"])
            )
    except Exception as e:
        print(f"Error reading labs from Postgres source: {e}")
        sys.exit(1)

    if measurement_inserts:
        print(f"Inserting {len(measurement_inserts)} measurements into MEASUREMENT...")
        execute_batched_insert(pg_tgt_cursor, """
            INSERT INTO measurement (
                measurement_id, person_id, measurement_concept_id, measurement_date, measurement_datetime,
                measurement_time, measurement_type_concept_id, operator_concept_id, value_as_number,
                value_as_concept_id, unit_concept_id, range_low, range_high, provider_id,
                visit_occurrence_id, visit_detail_id, measurement_source_value, measurement_source_concept_id,
                unit_source_value, unit_source_concept_id, value_source_value, measurement_event_id,
                meas_event_field_concept_id
            ) VALUES %s
            ON CONFLICT (measurement_id) DO NOTHING;
        """, measurement_inserts, batch_size=10000)
        loaded_counts["measurement"] = len(measurement_inserts)
        
    pg_tgt_conn.commit()
    print("Done.")

    # 7. Print summary metrics
    print("\n=== ETL PIPELINE RUN COMPLETE ===")
    print("Summary of records migrated to local OMOP database:")
    for table, count in loaded_counts.items():
        print(f" - {table.upper()}: {count} records loaded.")

    pg_src_cursor.close()
    pg_tgt_cursor.close()
    pg_src_conn.close()
    pg_tgt_conn.close()
    if mysql_conn:
        mysql_conn.close()

if __name__ == "__main__":
    main()