#!/usr/bin/env python3
"""
Lease Recovery Service for Maya3 Pipeline.

Runs periodically to reclaim expired video leases from crashed workers.
This is a CRITICAL service for production - should run on a reliable node.

Usage:
    python lease_recovery.py                    # One-time run
    python lease_recovery.py --daemon           # Run as daemon
    python lease_recovery.py --interval 300    # Custom interval (default: 5 min)
"""

import os
import sys
import time
import argparse
import logging
from datetime import datetime, timezone

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))

from dotenv import load_dotenv
load_dotenv()

from src.supabase_client import SupabaseClient, STATUS_PENDING, STATUS_CLAIMED

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s | %(levelname)-8s | %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)
logger = logging.getLogger("LeaseRecovery")


def reclaim_expired_leases(client: SupabaseClient) -> int:
    """
    Reset videos with expired leases back to PENDING.

    Returns:
        Number of videos reclaimed
    """
    now = datetime.now(timezone.utc).isoformat()

    # Find and update expired claims
    count = client.reclaim_expired_leases()

    if count > 0:
        logger.info(f"Reclaimed {count} expired leases")

    return count


def check_stale_workers(client: SupabaseClient, threshold_sec: int = 600) -> list:
    """
    Find workers that haven't reported heartbeat in threshold seconds.

    Returns:
        List of stale worker info dicts
    """
    from supabase import create_client

    sb = create_client(os.environ['URL'], os.environ['SUPABASE_ADMIN'])
    threshold = datetime.now(timezone.utc) - __import__('datetime').timedelta(seconds=threshold_sec)

    result = sb.table('worker_heartbeats').select(
        'worker_id, machine_id, gpu_id, status, current_video_id, last_heartbeat'
    ).neq('status', 'offline').lt('last_heartbeat', threshold.isoformat()).execute()

    stale = result.data or []

    if stale:
        logger.warning(f"Found {len(stale)} stale workers:")
        for w in stale:
            logger.warning(f"  {w['worker_id']} - last seen: {w['last_heartbeat']}")

    return stale


def run_daemon(interval_sec: int = 300):
    """Run lease recovery as a daemon."""
    logger.info(f"Starting lease recovery daemon (interval: {interval_sec}s)")

    client = SupabaseClient()

    while True:
        try:
            # Reclaim expired leases
            reclaimed = reclaim_expired_leases(client)

            # Check for stale workers (warning only)
            stale = check_stale_workers(client)

            # Log summary
            logger.info(f"Cycle complete - Reclaimed: {reclaimed}, Stale workers: {len(stale)}")

        except Exception as e:
            logger.error(f"Error in recovery cycle: {e}")

        time.sleep(interval_sec)


def main():
    parser = argparse.ArgumentParser(description="Maya3 Lease Recovery Service")
    parser.add_argument('--daemon', '-d', action='store_true', help='Run as daemon')
    parser.add_argument('--interval', '-i', type=int, default=300, help='Check interval (seconds)')
    args = parser.parse_args()

    client = SupabaseClient()

    if args.daemon:
        run_daemon(args.interval)
    else:
        # One-time run
        reclaimed = reclaim_expired_leases(client)
        stale = check_stale_workers(client)
        print(f"Reclaimed: {reclaimed} leases, Stale workers: {len(stale)}")


if __name__ == "__main__":
    main()