#!/usr/bin/env python3
"""
Python bridge for Apple iCloud Reminders using pyicloud.
Called from Node.js via subprocess.

Uses CloudKit Web Services API (ckdatabasews) for modern iOS 13+ reminders,
with fallback to legacy /rd/startup API for older accounts.

Usage:
  python3 icloud-reminders.py init <email> <password> [--session-dir <dir>]
  python3 icloud-reminders.py verify <email> <code> [--session-dir <dir>]
  python3 icloud-reminders.py lists <email> <password> [--session-dir <dir>]
  python3 icloud-reminders.py reminders <email> <password> [--collection-guid <guid>] [--session-dir <dir>]

Output: JSON to stdout
"""

import sys
import json
import os
import argparse
from urllib.parse import urlencode
from datetime import datetime


def get_session_dir(args_session_dir=None):
    return args_session_dir or os.environ.get(
        'ICLOUD_SESSION_DIR',
        os.path.join(os.getcwd(), 'data', 'icloud-sessions')
    )


def get_api(email, password, session_dir):
    """Create a PyiCloudService instance with session persistence."""
    from pyicloud import PyiCloudService

    os.makedirs(session_dir, exist_ok=True)

    api = PyiCloudService(
        apple_id=email,
        password=password,
        cookie_directory=session_dir,
    )
    return api


# ---------------------------------------------------------------------------
# CloudKit helpers
# ---------------------------------------------------------------------------

def get_ck_base(api):
    """Get the CloudKit base URL for the reminders container."""
    ck_url = api.data["webservices"]["ckdatabasews"]["url"]
    return f"{ck_url}/database/1/com.apple.reminders/production/private"


def ck_request(api, path, body=None):
    """Make a CloudKit request using the authenticated session."""
    base = get_ck_base(api)
    params = dict(api.params)
    url = f"{base}{path}?{urlencode(params)}"
    data = json.dumps(body) if body else "{}"
    resp = api.session.post(url, data=data, headers={"Content-type": "text/plain"})
    return resp.json()


def ck_list_zones(api):
    """List all zones in the reminders container."""
    return ck_request(api, "/zones/list")


def ck_query_records(api, record_type, zone_id, filter_by=None, results_limit=200):
    """Query records by type in a specific zone."""
    body = {
        "query": {"recordType": record_type},
        "zoneID": zone_id,
        "resultsLimit": results_limit,
    }
    if filter_by:
        body["query"]["filterBy"] = filter_by
    return ck_request(api, "/records/query", body)


def ck_get_zone_changes(api, zone_id, sync_token=None):
    """Get all changes in a zone (used to discover all record types)."""
    body = {"zoneID": zone_id}
    if sync_token:
        body["syncToken"] = sync_token
    return ck_request(api, "/changes/zone", body)


# ---------------------------------------------------------------------------
# CloudKit-based list/reminder fetching
# ---------------------------------------------------------------------------

# Common record type names to try for reminder lists
LIST_RECORD_TYPES = ["Checklist", "List", "REMCDList", "Collection", "ReminderList"]
# Common record type names to try for individual reminders
TASK_RECORD_TYPES = ["Task", "Reminder", "REMCDReminder", "Item"]


def discover_schema(api, zone_id):
    """Discover record types by fetching zone changes."""
    try:
        data = ck_get_zone_changes(api, zone_id)
        records = data.get("records", [])
        record_types = {}
        for r in records:
            rt = r.get("recordType", "unknown")
            if rt not in record_types:
                record_types[rt] = {
                    "count": 0,
                    "fields": list(r.get("fields", {}).keys()) if r.get("fields") else [],
                    "sample": r,
                }
            record_types[rt]["count"] += 1
        return record_types, records
    except Exception:
        return {}, []


def find_list_record_type(record_types):
    """Find the record type used for reminder lists."""
    # Check known names
    for rt in LIST_RECORD_TYPES:
        if rt in record_types:
            return rt
    # Heuristic: look for record types with a 'title' field
    for rt, info in record_types.items():
        fields = info.get("fields", [])
        if "title" in fields and "parentList" not in fields:
            return rt
    return None


def find_task_record_type(record_types):
    """Find the record type used for individual reminders."""
    for rt in TASK_RECORD_TYPES:
        if rt in record_types:
            return rt
    # Heuristic: look for record types with 'title' AND 'parentList'
    for rt, info in record_types.items():
        fields = info.get("fields", [])
        if "title" in fields and ("parentList" in fields or "parent" in fields):
            return rt
    return None


def extract_field_value(field):
    """Extract the value from a CloudKit field dict."""
    if isinstance(field, dict):
        return field.get("value")
    return field


def ck_timestamp_to_iso(ts):
    """Convert CloudKit timestamp (ms since epoch) to ISO string."""
    if ts is None:
        return None
    try:
        if isinstance(ts, (int, float)):
            return datetime.utcfromtimestamp(ts / 1000).isoformat()
        if isinstance(ts, dict):
            return datetime.utcfromtimestamp(ts.get("value", 0) / 1000).isoformat()
    except Exception:
        pass
    return None


def fetch_lists_via_cloudkit(api):
    """Fetch reminder lists using the CloudKit API. Returns (lists, error)."""
    try:
        # Step 1: Discover zones
        zones_data = ck_list_zones(api)
        zones = zones_data.get("zones", [])

        if not zones:
            return [], "No zones found in reminders container"

        all_lists = []

        for zone in zones:
            zone_id = zone.get("zoneID", {})
            zone_name = zone_id.get("zoneName", "unknown")

            # Step 2: Discover schema via zone changes
            record_types, all_records = discover_schema(api, zone_id)

            if not record_types:
                # Try querying common record types directly
                for rt in LIST_RECORD_TYPES:
                    try:
                        result = ck_query_records(api, rt, zone_id)
                        records = result.get("records", [])
                        if records:
                            record_types[rt] = {
                                "count": len(records),
                                "fields": list(records[0].get("fields", {}).keys()),
                            }
                            break
                    except Exception:
                        continue

            list_type = find_list_record_type(record_types)

            if list_type:
                # Query for lists using the discovered record type
                result = ck_query_records(api, list_type, zone_id)
                records = result.get("records", [])
                for r in records:
                    fields = r.get("fields", {})
                    title = extract_field_value(fields.get("title"))
                    if not title:
                        title = extract_field_value(fields.get("name"))
                    if not title:
                        title = "Untitled"

                    all_lists.append({
                        "guid": r.get("recordName", ""),
                        "title": title,
                        "zone": zone_name,
                        "color": extract_field_value(fields.get("color")),
                    })
            else:
                # Fallback: extract lists from zone changes data
                # Look for any record that looks like a list
                for r in all_records:
                    fields = r.get("fields", {})
                    title = extract_field_value(fields.get("title"))
                    if title and "parentList" not in fields and "parent" not in fields:
                        rt = r.get("recordType", "")
                        # Skip obvious task/reminder records
                        if rt.lower() not in ("task", "reminder", "item"):
                            all_lists.append({
                                "guid": r.get("recordName", ""),
                                "title": title,
                                "zone": zone_name,
                            })

        return all_lists, None

    except Exception as e:
        return [], str(e)


def fetch_reminders_via_cloudkit(api, collection_guid=None):
    """Fetch reminders using the CloudKit API. Returns (reminders, error)."""
    try:
        zones_data = ck_list_zones(api)
        zones = zones_data.get("zones", [])

        if not zones:
            return [], "No zones found in reminders container"

        all_reminders = []

        for zone in zones:
            zone_id = zone.get("zoneID", {})

            # Discover schema
            record_types, all_records = discover_schema(api, zone_id)
            task_type = find_task_record_type(record_types)

            if task_type:
                # Build filter if collection_guid specified
                filter_by = None
                if collection_guid:
                    filter_by = [{
                        "fieldName": "parentList",
                        "comparator": "EQUALS",
                        "fieldValue": {
                            "value": {
                                "recordName": collection_guid,
                                "action": "NONE",
                            },
                            "type": "REFERENCE",
                        },
                    }]

                result = ck_query_records(api, task_type, zone_id, filter_by=filter_by)
                records = result.get("records", [])
            else:
                # Use zone changes data and filter for task-like records
                records = []
                for r in all_records:
                    fields = r.get("fields", {})
                    if "parentList" in fields or "parent" in fields:
                        if collection_guid:
                            parent_ref = extract_field_value(fields.get("parentList") or fields.get("parent"))
                            if isinstance(parent_ref, dict):
                                parent_name = parent_ref.get("recordName", "")
                            else:
                                parent_name = str(parent_ref) if parent_ref else ""
                            if parent_name != collection_guid:
                                continue
                        records.append(r)

            for r in records:
                fields = r.get("fields", {})
                title = extract_field_value(fields.get("title"))
                if not title:
                    continue

                # Get parent list reference
                parent_ref = extract_field_value(fields.get("parentList") or fields.get("parent"))
                p_guid = ""
                if isinstance(parent_ref, dict):
                    p_guid = parent_ref.get("recordName", "")
                elif parent_ref:
                    p_guid = str(parent_ref)

                # Parse completion status
                is_completed = False
                completed_val = extract_field_value(fields.get("isCompleted") or fields.get("completed"))
                if completed_val is not None:
                    is_completed = bool(completed_val)

                # Parse dates
                due_date = None
                due_val = extract_field_value(fields.get("dueDate") or fields.get("due"))
                if due_val:
                    due_date = ck_timestamp_to_iso(due_val)

                completed_date = None
                cd_val = extract_field_value(fields.get("completionDate") or fields.get("completedDate"))
                if cd_val:
                    completed_date = ck_timestamp_to_iso(cd_val)

                all_reminders.append({
                    "guid": r.get("recordName", ""),
                    "pGuid": p_guid,
                    "title": title,
                    "description": extract_field_value(fields.get("notes") or fields.get("description")) or "",
                    "priority": extract_field_value(fields.get("priority")) or 0,
                    "isCompleted": is_completed,
                    "dueDate": due_date,
                    "completedDate": completed_date,
                })

        return all_reminders, None

    except Exception as e:
        return [], str(e)


# ---------------------------------------------------------------------------
# Legacy API fallback
# ---------------------------------------------------------------------------

def fetch_lists_via_legacy(api):
    """Fetch lists using the legacy /rd/startup API."""
    reminders_service = api.reminders
    reminders_service.refresh()

    lists = []
    for list_title in reminders_service.lists.keys():
        lists.append({
            "guid": list_title,
            "title": list_title,
        })
    return lists


def fetch_reminders_via_legacy(api, collection_guid=None):
    """Fetch reminders using the legacy /rd/startup API."""
    reminders_service = api.reminders
    reminders_service.refresh()

    result = []
    for list_title, tasks in reminders_service.lists.items():
        if collection_guid and list_title != collection_guid:
            continue

        for i, task in enumerate(tasks):
            reminder = {
                "guid": f"{list_title}_{i}",
                "pGuid": list_title,
                "title": task.get("title", "Untitled") if isinstance(task, dict) else str(task),
                "description": (task.get("desc", "") or "") if isinstance(task, dict) else "",
                "priority": task.get("priority", 0) if isinstance(task, dict) else 0,
                "isCompleted": False,
                "dueDate": None,
                "completedDate": None,
            }

            if isinstance(task, dict):
                due = task.get("due")
                if due and hasattr(due, 'isoformat'):
                    reminder["dueDate"] = due.isoformat()

            result.append(reminder)
    return result


# ---------------------------------------------------------------------------
# Commands
# ---------------------------------------------------------------------------

def cmd_init(args):
    """Initialize iCloud session. Returns whether 2FA is needed."""
    try:
        api = get_api(args.email, args.password, get_session_dir(args.session_dir))

        if api.requires_2fa:
            print(json.dumps({
                "needs2FA": True,
                "error": None
            }))
        elif api.requires_2sa:
            print(json.dumps({
                "needs2FA": True,
                "error": None
            }))
        else:
            print(json.dumps({
                "needs2FA": False,
                "error": None
            }))
    except Exception as e:
        print(json.dumps({
            "needs2FA": False,
            "error": str(e)
        }))


def cmd_verify(args):
    """Submit 2FA verification code."""
    try:
        api = get_api(args.email, args.password, get_session_dir(args.session_dir))

        if not api.requires_2fa and not api.requires_2sa:
            print(json.dumps({
                "success": True,
                "error": None,
                "message": "Already verified"
            }))
            return

        result = api.validate_2fa_code(args.code)

        if result:
            if not api.is_trusted_session:
                api.trust_session()
            print(json.dumps({
                "success": True,
                "error": None
            }))
        else:
            print(json.dumps({
                "success": False,
                "error": "Invalid verification code"
            }))
    except Exception as e:
        print(json.dumps({
            "success": False,
            "error": str(e)
        }))


def cmd_lists(args):
    """Fetch all reminder lists. Tries CloudKit first, falls back to legacy API."""
    try:
        api = get_api(args.email, args.password, get_session_dir(args.session_dir))

        if api.requires_2fa or api.requires_2sa:
            print(json.dumps({
                "error": "2FA required. Please reconnect Apple Reminders.",
                "needs2FA": True,
                "lists": []
            }))
            return

        # Try CloudKit first (modern iOS 13+ reminders)
        lists, ck_error = fetch_lists_via_cloudkit(api)

        if lists:
            print(json.dumps({
                "error": None,
                "lists": lists,
                "source": "cloudkit"
            }))
            return

        # Fallback to legacy API
        sys.stderr.write(f"[REMINDERS] CloudKit failed ({ck_error}), trying legacy API...\n")
        legacy_lists = fetch_lists_via_legacy(api)

        print(json.dumps({
            "error": None,
            "lists": legacy_lists,
            "source": "legacy"
        }))

    except Exception as e:
        print(json.dumps({
            "error": str(e),
            "lists": []
        }))


def cmd_reminders(args):
    """Fetch reminders. Tries CloudKit first, falls back to legacy API."""
    try:
        api = get_api(args.email, args.password, get_session_dir(args.session_dir))

        if api.requires_2fa or api.requires_2sa:
            print(json.dumps({
                "error": "2FA required. Please reconnect Apple Reminders.",
                "reminders": []
            }))
            return

        # Try CloudKit first
        reminders, ck_error = fetch_reminders_via_cloudkit(api, args.collection_guid)

        if reminders or ck_error is None:
            print(json.dumps({
                "error": None,
                "reminders": reminders,
                "source": "cloudkit"
            }))
            return

        # Fallback to legacy API
        sys.stderr.write(f"[REMINDERS] CloudKit failed ({ck_error}), trying legacy API...\n")
        legacy_reminders = fetch_reminders_via_legacy(api, args.collection_guid)

        print(json.dumps({
            "error": None,
            "reminders": legacy_reminders,
            "source": "legacy"
        }))

    except Exception as e:
        print(json.dumps({
            "error": str(e),
            "reminders": []
        }))


def main():
    parser = argparse.ArgumentParser(description='iCloud Reminders bridge')
    parser.add_argument('--session-dir', type=str, default=None)
    subparsers = parser.add_subparsers(dest='command')

    # init
    p_init = subparsers.add_parser('init')
    p_init.add_argument('email')
    p_init.add_argument('password')

    # verify
    p_verify = subparsers.add_parser('verify')
    p_verify.add_argument('email')
    p_verify.add_argument('password')
    p_verify.add_argument('code')

    # lists
    p_lists = subparsers.add_parser('lists')
    p_lists.add_argument('email')
    p_lists.add_argument('password')

    # reminders
    p_rem = subparsers.add_parser('reminders')
    p_rem.add_argument('email')
    p_rem.add_argument('password')
    p_rem.add_argument('--collection-guid', type=str, default=None)

    args = parser.parse_args()

    if args.command == 'init':
        cmd_init(args)
    elif args.command == 'verify':
        cmd_verify(args)
    elif args.command == 'lists':
        cmd_lists(args)
    elif args.command == 'reminders':
        cmd_reminders(args)
    else:
        parser.print_help()
        sys.exit(1)


if __name__ == '__main__':
    main()