"""
Backfill severity field on all diagnosis documents.

Infers severity from the notes field using keyword matching,
then stores it explicitly as a top-level field.

REVERSIBLE: To undo, run in mongosh:
    db.diagnoses.updateMany({}, {$unset: {severity: ""}})
"""

import sys
import os

sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))

from pymongo import MongoClient
from dotenv import load_dotenv

load_dotenv(os.path.join(os.path.dirname(__file__), '..', '.env'))

MONGODB_URI = os.getenv('MONGODB_URI')
MONGODB_DB_NAME = os.getenv('MONGODB_DB_NAME', 'hipograf')


def infer_severity(notes):
    if not notes:
        return 'mild'
    lower = notes.lower()
    if 'severe' in lower:
        return 'severe'
    if 'moderate' in lower or 'stage' in lower:
        return 'moderate'
    return 'mild'


def main():
    client = MongoClient(MONGODB_URI)
    db = client[MONGODB_DB_NAME]

    diagnoses = list(db.diagnoses.find({}, {'_id': 1, 'notes': 1}))
    print(f"Found {len(diagnoses)} diagnoses to backfill.")

    updated = 0
    for doc in diagnoses:
        severity = infer_severity(doc.get('notes'))
        db.diagnoses.update_one(
            {'_id': doc['_id']},
            {'$set': {'severity': severity}}
        )
        updated += 1

    print(f"Done. Updated {updated} documents.")
    client.close()


if __name__ == '__main__':
    main()
