#!/usr/bin/env python3
"""
Step 1: Data Acquisition - GLP-1 Receptor Agonists Adverse Events from OpenFDA FAERS

This script queries the OpenFDA FAERS (drug/event) endpoint to retrieve adverse event
reports for major GLP-1 receptor agonists: Semaglutide, Tirzepatide, Liraglutide,
Dulaglutide, and Exenatide.
"""

import json
import time
import urllib.request
import urllib.parse
import urllib.error
from pathlib import Path
from datetime import datetime
from collections import defaultdict

# Session directory
SESSION_DIR = Path("/app/sandbox/session_20260203_092044_b80cb683d2ee")
OUTPUT_DIR = SESSION_DIR / "workflow" / "data"
RESULTS_DIR = SESSION_DIR / "results"

# Ensure directories exist
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
RESULTS_DIR.mkdir(parents=True, exist_ok=True)

# OpenFDA API configuration
BASE_URL = "https://api.fda.gov/drug/event.json"

# GLP-1 Receptor Agonists to query (generic names)
GLP1_DRUGS = [
    "semaglutide",
    "tirzepatide",
    "liraglutide",
    "dulaglutide",
    "exenatide"
]

# API limits: OpenFDA allows up to 1000 records per request with skip parameter
# Maximum skip is typically around 25,000 records
LIMIT_PER_REQUEST = 1000
MAX_RECORDS_PER_DRUG = 25000  # OpenFDA API limit on skip parameter

def fetch_events_for_drug(drug_name: str, max_records: int = MAX_RECORDS_PER_DRUG) -> list:
    """
    Fetch adverse event records for a specific drug from OpenFDA.

    Args:
        drug_name: Generic name of the drug
        max_records: Maximum number of records to fetch

    Returns:
        List of adverse event records
    """
    all_records = []
    skip = 0

    # Search in patient.drug.medicinalproduct field (case-insensitive)
    search_query = f'patient.drug.medicinalproduct:"{drug_name}"'

    print(f"\n{'='*60}")
    print(f"Fetching adverse events for: {drug_name.upper()}")
    print(f"{'='*60}")

    while skip < max_records:
        params = {
            "search": search_query,
            "limit": LIMIT_PER_REQUEST,
            "skip": skip
        }

        url = f"{BASE_URL}?{urllib.parse.urlencode(params)}"

        try:
            # Make request with timeout
            req = urllib.request.Request(url)
            req.add_header('User-Agent', 'K-Dense-Research-Agent/1.0')

            with urllib.request.urlopen(req, timeout=60) as response:
                data = json.loads(response.read().decode('utf-8'))

            # Extract results
            results = data.get("results", [])
            if not results:
                print(f"  No more records available at skip={skip}")
                break

            all_records.extend(results)

            # Progress update
            total_available = data.get("meta", {}).get("results", {}).get("total", "unknown")
            print(f"  Fetched {len(results)} records (skip={skip}, total so far: {len(all_records)}, available: {total_available})")

            # Check if we've reached the end
            if len(results) < LIMIT_PER_REQUEST:
                print(f"  Reached end of available records")
                break

            skip += LIMIT_PER_REQUEST

            # Rate limiting - be respectful to the API
            time.sleep(0.5)

        except urllib.error.HTTPError as e:
            if e.code == 404:
                print(f"  No records found for {drug_name}")
                break
            elif e.code == 429:
                print(f"  Rate limited, waiting 10 seconds...")
                time.sleep(10)
                continue
            else:
                print(f"  HTTP Error {e.code}: {e.reason}")
                # Try to continue with what we have
                break
        except urllib.error.URLError as e:
            print(f"  URL Error: {e.reason}")
            break
        except Exception as e:
            print(f"  Error: {str(e)}")
            break

    print(f"  TOTAL: {len(all_records)} records fetched for {drug_name}")
    return all_records


def extract_key_fields(record: dict) -> dict:
    """
    Extract key fields from a FAERS record for analysis.

    Args:
        record: Raw FAERS record

    Returns:
        Dictionary with key fields extracted
    """
    # Basic identifiers
    extracted = {
        "safetyreportid": record.get("safetyreportid"),
        "receivedate": record.get("receivedate"),
        "receiptdate": record.get("receiptdate"),
        "transmissiondate": record.get("transmissiondate"),
        "serious": record.get("serious"),
        "seriousnessdeath": record.get("seriousnessdeath"),
        "seriousnesslifethreatening": record.get("seriousnesslifethreatening"),
        "seriousnesshospitalization": record.get("seriousnesshospitalization"),
        "seriousnessdisabling": record.get("seriousnessdisabling"),
        "seriousnesscongenitalanomali": record.get("seriousnesscongenitalanomali"),
        "seriousnessother": record.get("seriousnessother"),
    }

    # Patient information
    patient = record.get("patient", {})
    extracted["patient_age"] = patient.get("patientonsetage")
    extracted["patient_age_unit"] = patient.get("patientonsetageunit")
    extracted["patient_sex"] = patient.get("patientsex")
    extracted["patient_weight"] = patient.get("patientweight")

    # Reactions/Adverse events
    reactions = patient.get("reaction", [])
    extracted["reactions"] = [
        {
            "reactionmeddrapt": r.get("reactionmeddrapt"),
            "reactionoutcome": r.get("reactionoutcome")
        }
        for r in reactions
    ]
    extracted["reaction_count"] = len(reactions)

    # Drugs involved
    drugs = patient.get("drug", [])
    extracted["drugs"] = [
        {
            "medicinalproduct": d.get("medicinalproduct"),
            "drugindication": d.get("drugindication"),
            "drugcharacterization": d.get("drugcharacterization"),
            "activesubstance": d.get("activesubstance", {}).get("activesubstancename") if d.get("activesubstance") else None
        }
        for d in drugs
    ]
    extracted["drug_count"] = len(drugs)

    # Keep full record for reference
    extracted["_raw"] = record

    return extracted


def main():
    """Main execution function."""
    print("=" * 70)
    print("GLP-1 Receptor Agonists Adverse Events - Data Acquisition")
    print(f"Started at: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print("=" * 70)

    # Collect all data
    all_events = []
    drug_counts = {}

    for drug in GLP1_DRUGS:
        print(f"\nProcessing: {drug}")
        records = fetch_events_for_drug(drug)
        drug_counts[drug] = len(records)

        # Extract key fields and tag with drug name
        for record in records:
            extracted = extract_key_fields(record)
            extracted["query_drug"] = drug
            all_events.append(extracted)

        print(f"Cumulative total events: {len(all_events)}")

    # Save raw data to JSON
    output_file = OUTPUT_DIR / "glp1_adverse_events.json"
    print(f"\nSaving data to: {output_file}")

    with open(output_file, 'w') as f:
        json.dump(all_events, f, indent=2)

    file_size_mb = output_file.stat().st_size / (1024 * 1024)
    print(f"Saved {len(all_events)} records ({file_size_mb:.2f} MB)")

    # Generate summary
    summary_lines = [
        "=" * 70,
        "GLP-1 Receptor Agonists Adverse Events - Data Acquisition Summary",
        f"Generated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}",
        "=" * 70,
        "",
        "RECORD COUNTS BY DRUG:",
        "-" * 40
    ]

    total_records = 0
    for drug, count in sorted(drug_counts.items(), key=lambda x: -x[1]):
        summary_lines.append(f"  {drug.upper():20} : {count:,} records")
        total_records += count

    summary_lines.extend([
        "-" * 40,
        f"  TOTAL                 : {total_records:,} records",
        "",
        "DATA FIELDS AVAILABLE:",
        "-" * 40,
        "  - safetyreportid: Unique report identifier",
        "  - receivedate: Date report was received by FDA",
        "  - serious: Whether event was serious (1/2)",
        "  - seriousnessdeath: Death occurred (1 if yes)",
        "  - seriousnesslifethreatening: Life-threatening (1 if yes)",
        "  - seriousnesshospitalization: Hospitalization (1 if yes)",
        "  - patient_age: Patient age at event onset",
        "  - patient_sex: Patient sex (1=male, 2=female)",
        "  - reactions: List of adverse reactions (MedDRA terms)",
        "  - drugs: List of concomitant drugs",
        "  - query_drug: The GLP-1 drug that matched the query",
        "",
        "OUTPUT FILES:",
        "-" * 40,
        f"  Raw data: {output_file}",
        f"  File size: {file_size_mb:.2f} MB",
        "",
        "VERIFICATION:",
        "-" * 40
    ])

    # Verify data fields
    if all_events:
        sample = all_events[0]
        has_reactions = sample.get("reaction_count", 0) > 0
        has_date = sample.get("receivedate") is not None
        has_outcome_info = sample.get("serious") is not None

        summary_lines.extend([
            f"  Sample record has reactions: {'Yes' if has_reactions else 'No'}",
            f"  Sample record has date: {'Yes' if has_date else 'No'}",
            f"  Sample record has outcome info: {'Yes' if has_outcome_info else 'No'}",
            f"  Total unique report IDs: {len(set(e.get('safetyreportid') for e in all_events if e.get('safetyreportid')))}",
        ])

    summary_lines.extend([
        "",
        "STATUS: Data acquisition SUCCESSFUL",
        "=" * 70
    ])

    summary_text = "\n".join(summary_lines)

    # Save summary
    summary_file = RESULTS_DIR / "data_acquisition_summary.txt"
    with open(summary_file, 'w') as f:
        f.write(summary_text)

    print("\n" + summary_text)
    print(f"\nSummary saved to: {summary_file}")

    return len(all_events), drug_counts


if __name__ == "__main__":
    total, counts = main()
    print(f"\n[COMPLETE] Fetched {total:,} total adverse event records")