#!/usr/bin/env python3

import csv
import sys
import xml.etree.ElementTree as ET
from pathlib import Path


WRIST_TEMP_TYPE = "HKQuantityTypeIdentifierAppleSleepingWristTemperature"
SLEEP_TYPE = "HKCategoryTypeIdentifierSleepAnalysis"


def main():
    if len(sys.argv) != 2:
        print("Usage: python3 extract_health_data.py /path/to/export.xml")
        sys.exit(1)

    input_path = Path(sys.argv[1]).expanduser()

    if not input_path.exists():
        print(f"File not found: {input_path}")
        sys.exit(1)

    output_dir = input_path.parent
    wrist_temp_csv = output_dir / "wrist_temperature.csv"
    sleep_csv = output_dir / "sleep_analysis.csv"

    wrist_rows = []
    sleep_rows = []

    print(f"Reading: {input_path}")
    print("This may take a little while for a large Apple Health export...")

    for event, elem in ET.iterparse(input_path, events=("end",)):
        if elem.tag != "Record":
            continue

        record_type = elem.attrib.get("type")

        if record_type == WRIST_TEMP_TYPE:
            wrist_rows.append({
                "sourceName": elem.attrib.get("sourceName", ""),
                "unit": elem.attrib.get("unit", ""),
                "value": elem.attrib.get("value", ""),
                "startDate": elem.attrib.get("startDate", ""),
                "endDate": elem.attrib.get("endDate", ""),
                "creationDate": elem.attrib.get("creationDate", ""),
            })

        elif record_type == SLEEP_TYPE:
            sleep_rows.append({
                "sourceName": elem.attrib.get("sourceName", ""),
                "value": elem.attrib.get("value", ""),
                "startDate": elem.attrib.get("startDate", ""),
                "endDate": elem.attrib.get("endDate", ""),
                "creationDate": elem.attrib.get("creationDate", ""),
            })

        # Important for large XML files:
        # discard processed elements to keep memory use low.
        elem.clear()

    with wrist_temp_csv.open("w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(
            f,
            fieldnames=[
                "sourceName",
                "unit",
                "value",
                "startDate",
                "endDate",
                "creationDate",
            ],
        )
        writer.writeheader()
        writer.writerows(wrist_rows)

    with sleep_csv.open("w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(
            f,
            fieldnames=[
                "sourceName",
                "value",
                "startDate",
                "endDate",
                "creationDate",
            ],
        )
        writer.writeheader()
        writer.writerows(sleep_rows)

    print()
    print("Done.")
    print(f"Wrist temperature records: {len(wrist_rows):,}")
    print(f"Sleep records:             {len(sleep_rows):,}")
    print()
    print(f"Created:")
    print(f"  {wrist_temp_csv}")
    print(f"  {sleep_csv}")


if __name__ == "__main__":
    main()