variable_reader_engine.py

import json
import paramiko
import psycopg2
import io
import pandas as pd
import subprocess
import os
import time
import warnings
from collections import defaultdict

# Suppress FutureWarning from Pandas
warnings.simplefilter(action='ignore', category=FutureWarning)

# --- CONFIGURATION (MASKED) ---
DB_NAME = "YOUR_DATABASE_NAME"
DB_HOST = "YOUR_DB_HOST_ADDRESS"
DB_PORT = "5432"

def _parse_variables(variable_input):
    """Splits comma-separated variable string into a clean list."""
    if not variable_input: return []
    return list(set([v.strip() for v in variable_input.split(',') if v.strip()]))

def _format_results(found_values, searched_vars):
    """Formats the search results into a readable report."""
    if not found_values:
        return f"No values found for: {', '.join(searched_vars)}"
    
    grouped = defaultdict(list)
    for item in found_values:
        grouped[item['variable']].append(item)

    report_lines = []
    sorted_vars = sorted(list(set(searched_vars)))

    for var in sorted_vars:
        if var not in grouped:
            report_lines.append(f"Variable: {var}\n  - Not found.\n")
            continue

        items = grouped[var]
        report_lines.append(f"Variable: {var}")
        
        unique_items = []
        seen = set()
        
        def sort_key(x):
            return (x.get('date', '0000-00-00'), x.get('event', 'N/A'))

        for item in sorted(items, key=sort_key):
            sig = (item['variable'], item['value'], item.get('event'), item.get('source_file'))
            if sig not in seen:
                seen.add(sig)
                unique_items.append(item)

        for item in unique_items:
            value = item['value']
            details = []
            if 'event' in item and item['event'] != 'N/A':
                details.append(f"Event: {item['event']}")
            if 'date' in item and item['date'] != 'N/A':
                details.append(f"Date: {item['date']}")
            if 'source_file' in item:
                details.append(f"File: {item['source_file']}")
            
            report_lines.append(f"  - {value} ({', '.join(details)})")
        report_lines.append("") 

    return "\n".join(report_lines).strip()

def _process_stream_data(chunk_iter, subject_id, target_vars, filename, found_values):
    """Processes CSV data chunks to find matching variables for a specific subject."""
    # These column names might need adjustment based on your data schema
    event_column = 'redcap_event_name'
    date_column = 'interview_date'

    for df in chunk_iter:
        subject_col = None
        for c in ['subject_id', 'src_subject_id', 'subjectkey']:
            if c in df.columns:
                subject_col = c
                break
        
        if subject_col:
            clean_subject_id = str(subject_id).strip()
            mask = df[subject_col].astype(str).str.strip() == clean_subject_id
            subject_df = df[mask].copy()
        else:
            continue

        if subject_df.empty: continue

        for var_name in target_vars:
            if var_name not in subject_df.columns: continue
            
            for _, row in subject_df.iterrows():
                val = str(row[var_name]).strip()
                if val == "" or val.lower() == "nan": continue
                
                event = row.get(event_column, 'N/A')
                date_str = 'N/A'
                if date_column in row and pd.notna(row[date_column]):
                    try:
                        date_str = pd.to_datetime(row[date_column]).strftime('%Y-%m-%d')
                    except: pass
                
                found_values.append({
                    "variable": var_name,
                    "value": val,
                    "event": event,
                    "date": date_str,
                    "source_file": filename
                })

# --- SSH SEARCH LOGIC ---

def _read_remote_csv_paramiko(ssh_client, remote_path, subject_id, target_vars, found_values):
    filename = os.path.basename(remote_path)
    try:
        stdin, stdout, stderr = ssh_client.exec_command(f"cat \"{remote_path}\"", get_pty=False)
        chunk_iter = pd.read_csv(stdout, chunksize=10000, dtype=str, keep_default_na=False, on_bad_lines='skip')
        _process_stream_data(chunk_iter, subject_id, target_vars, filename, found_values)
    except Exception as e:
        print(f"  [Error] Failed to read {filename}: {e}")

def _search_combined_csv(ssh_client, base_path, file_pattern, subject_id, variable_input):
    target_vars = _parse_variables(variable_input)
    found_values = []
    search_glob = f"{base_path.rstrip('/')}/{file_pattern}"
    cmd = f"grep -l '{subject_id}' {search_glob} 2>/dev/null"
    
    try:
        stdin, stdout, stderr = ssh_client.exec_command(cmd, timeout=30)
        target_files = [f.strip() for f in stdout.read().decode('utf-8').split('\n') if f.strip()]
        
        for filepath in target_files:
            _read_remote_csv_paramiko(ssh_client, filepath, subject_id, target_vars, found_values)
    except Exception as e:
        return f"Error: {e}"

    return _format_results(found_values, target_vars)

def _read_surveys_ssh(ssh_client, path, subject_id, variable_input):
    target_vars = _parse_variables(variable_input)
    found_values = []
    
    try:
        cmd = f"ls {path}*.json {path}*.csv 2>/dev/null"
        stdin, stdout, stderr = ssh_client.exec_command(cmd, timeout=30)
        files = [f.strip() for f in stdout.read().decode('utf-8').split('\n') if f.strip()]
        
        for filepath in files:
            filename = os.path.basename(filepath)
            if filename.endswith('.json'):
                stdin, stdout, stderr = ssh_client.exec_command(f"cat \"{filepath}\"")
                data = json.loads(stdout.read().decode('utf-8'))
                data_list = data if isinstance(data, list) else [data]
                for entry in data_list:
                    for var in target_vars:
                        val = entry.get(var)
                        if val:
                            found_values.append({"variable": var, "value": str(val), "source_file": filename})
            elif filename.endswith('.csv'):
                _read_remote_csv_paramiko(ssh_client, filepath, subject_id, target_vars, found_values)
    except Exception as e:
        return f"Error: {e}"

    return _format_results(found_values, target_vars)

# --- DISPATCHER ---

def read_from_ssh_server(source, host, user, password, site, subject_id, variable_input, search_type):
    try:
        ssh = paramiko.SSHClient()
        ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        
        if source == "NetworkA":
            # MASKED PATH
            path = f"/path/to/network_a/data/{site}/raw/{subject_id}/surveys/"
            ssh.connect(host, username=user, timeout=15)
            result = _read_surveys_ssh(ssh, path, subject_id, variable_input)
            ssh.close()
            return result
            
        elif source == "Internal":
            ssh.connect(host, username=user, password=password, timeout=15)
            if search_type == "Surveys":
                # MASKED PATH
                p = f"/path/to/internal/surveys/{site}/{subject_id}/"
                result = _read_surveys_ssh(ssh, p, subject_id, variable_input)
            else:
                # MASKED PATH
                base = "/path/to/internal/combined_data/"
                pat = "*.csv"
                result = _search_combined_csv(ssh, base, pat, subject_id, variable_input)
            ssh.close()
            return result
    except Exception as e: return f"Connection Error: {e}"

def read_from_database(db_user, db_password, subject_id, variable_input):
    target_vars = _parse_variables(variable_input)
    try:
        conn = psycopg2.connect(dbname=DB_NAME, user=db_user, password=db_password, host=DB_HOST, port=DB_PORT)
        cursor = conn.cursor()
        # Ensure 'forms' schema and 'subject_id' column match your DB
        cursor.execute("SELECT form_data, event_name FROM data_table WHERE subject_id = %s", (subject_id,))
        results = cursor.fetchall()
        cursor.close(); conn.close()
        
        found_values = []
        for row in results:
            form_data, event_name = row
            if isinstance(form_data, dict):
                for var in target_vars:
                    val = form_data.get(var)
                    if val is not None:
                        found_values.append({"variable": var, "value": str(val), "event": event_name})
        return _format_results(found_values, target_vars)
    except Exception as e: return f"DB Error: {e}"