# -----------------------------------------------
# prepare_data.py
# Loads HuggingFace medical dataset
# Combines selected splits
# Saves to data/medical_knowledge.txt for RAG
# Run once before build_faiss.py
# -----------------------------------------------

from datasets import load_dataset
from config import (
    DATASET_NAME,
    DATASET_SPLITS,
    MEDICAL_KNOWLEDGE_FILE
)
import os

# -----------------------------------------------
# Load one split from HuggingFace
# -----------------------------------------------
def load_split(split_name):
    try:
        print(f"  Loading '{split_name}' split...")
        dataset = load_dataset(DATASET_NAME, split=split_name)
        print(f"  ✅ Loaded {len(dataset)} rows")
        return dataset
    except Exception as e:
        print(f"  ⚠️  Failed to load '{split_name}': {e}")
        return None

# -----------------------------------------------
# Convert one dataset row to structured text
# -----------------------------------------------
def row_to_text(row, split_name):
    disease   = str(row.get("Disease",   "")).strip()
    symptoms  = str(row.get("Symptoms",  "")).strip()
    treatment = str(row.get("Treatments","")).strip()

    # Skip empty or useless rows
    if not disease or not symptoms:
        return None
    if len(symptoms) < 10:
        return None

    # Build structured text entry
    text = f"DISEASE: {disease}\n"
    text += f"SYMPTOMS: {symptoms}\n"
    if treatment:
        text += f"TREATMENTS: {treatment}\n"
    text += f"SOURCE_SPLIT: {split_name}\n"

    return text

# -----------------------------------------------
# Save all entries to txt file
# -----------------------------------------------
def save_knowledge_base(entries):
    os.makedirs("data", exist_ok=True)

    with open(MEDICAL_KNOWLEDGE_FILE, "w", encoding="utf-8") as f:
        for entry in entries:
            f.write(entry)
            f.write("\n---\n\n")

    print(f"\n✅ Saved {len(entries)} entries to {MEDICAL_KNOWLEDGE_FILE}")

# -----------------------------------------------
# Main
# -----------------------------------------------
if __name__ == "__main__":
    print("=== Medical Report Analyzer — Data Preparation ===\n")

    all_entries  = []
    seen_diseases = set()  # track duplicates

    # Step 1 — Load each split
    print("Step 1: Loading dataset splits...")
    for split_name in DATASET_SPLITS:
        dataset = load_split(split_name)
        if dataset is None:
            continue

        split_entries = 0
        for row in dataset:
            text = row_to_text(row, split_name)
            if text is None:
                continue

            # Deduplicate by disease+symptoms combo
            key = row.get("Disease","").strip().lower()
            symptoms_key = row.get("Symptoms","").strip().lower()[:50]
            unique_key = f"{key}_{symptoms_key}"

            if unique_key not in seen_diseases:
                seen_diseases.add(unique_key)
                all_entries.append(text)
                split_entries += 1

        print(f"  ✅ {split_name}: added {split_entries} unique entries")

    print(f"\nTotal unique entries: {len(all_entries)}")

    if not all_entries:
        print("❌ No entries collected. Check dataset loading.")
        exit()

    # Step 2 — Save
    print("\nStep 2: Saving knowledge base...")
    save_knowledge_base(all_entries)

    # Step 3 — Verify
    print("\n=== VERIFICATION ===")
    file_size = os.path.getsize(MEDICAL_KNOWLEDGE_FILE)
    print(f"✅ File exists: {os.path.exists(MEDICAL_KNOWLEDGE_FILE)}")
    print(f"✅ File size: {file_size:,} bytes")
    print(f"✅ Total entries: {len(all_entries)}")

    # Step 4 — Preview first 3 entries
    print("\n=== PREVIEW (first 3 entries) ===")
    for i, entry in enumerate(all_entries[:3]):
        print(f"--- Entry {i+1} ---")
        print(entry)

    print("✅ Data preparation complete!")
    print("Next step: Run build_faiss.py")