#!/usr/bin/env python3
"""
Retry Manager — Quản lý batch bị failed.

- Reset 1 batch failed về pending
- Reset tất cả failed về pending
- Kiểm tra file bị hỏng

Usage:
    python3 scripts/retry_manager.py batch-index.json reset <batch_id>
    python3 scripts/retry_manager.py batch-index.json reset-all
    python3 scripts/retry_manager.py batch-index.json status
"""

import json
import sys
import os
from datetime import datetime, timezone


def load_index(filepath: str) -> dict:
    with open(filepath, "r", encoding="utf-8") as f:
        return json.load(f)


def save_index(filepath: str, data: dict):
    backup = filepath + ".bak"
    if os.path.exists(filepath):
        with open(filepath, "r") as src:
            with open(backup, "w") as dst:
                dst.write(src.read())
    with open(filepath, "w", encoding="utf-8") as f:
        json.dump(data, f, ensure_ascii=False, indent=2)
    print(f"✅ Đã lưu (backup: {backup})")


def reset_batch(filepath: str, batch_id: str):
    data = load_index(filepath)
    for batch in data["batches"]:
        if batch["id"] == batch_id:
            old_status = batch["status"]
            batch["status"] = "pending"
            batch["errors"] = 0
            batch["accuracy"] = None
            save_index(filepath, data)
            print(f"✅ Batch {batch_id}: {old_status} → pending")
            return
    print(f"❌ Không tìm thấy batch {batch_id}")
    sys.exit(1)


def reset_all_failed(filepath: str):
    data = load_index(filepath)
    count = 0
    for batch in data["batches"]:
        if batch["status"] == "failed":
            batch["status"] = "pending"
            batch["errors"] = 0
            batch["accuracy"] = None
            count += 1
    if count == 0:
        print("ℹ️  Không có batch failed nào")
    else:
        save_index(filepath, data)
        print(f"✅ Đã reset {count} batch failed → pending")


def show_status(filepath: str):
    data = load_index(filepath)
    print(f"Project: {data.get('project', 'N/A')}")
    print(f"Created: {data.get('created', 'N/A')}")
    print(f"\n{'ID':<20} {'Pages':<10} {'Status':<15} {'Errors':<8} {'Accuracy':<10}")
    print("-" * 65)
    for batch in data["batches"]:
        status_icon = {"pending": "⏳", "in_progress": "🔄", "done": "✅", "failed": "❌"}.get(batch["status"], "?")
        print(f"{batch['id']:<20} {batch['pages']:<10} {status_icon} {batch['status']:<13} {batch.get('errors', 0):<8} {batch.get('accuracy') or '':<10}")

    # Thống kê
    total = len(data["batches"])
    done = sum(1 for b in data["batches"] if b["status"] == "done")
    failed = sum(1 for b in data["batches"] if b["status"] == "failed")
    in_progress = sum(1 for b in data["batches"] if b["status"] == "in_progress")
    pending = sum(1 for b in data["batches"] if b["status"] == "pending")

    print(f"\n📊 {total} batches: {done} done | {in_progress} in_progress | {pending} pending | {failed} failed")
    print(f"   Progress: {done}/{total} ({done/total*100:.0f}%)" if total > 0 else "   Progress: 0%")


def main():
    if len(sys.argv) < 2:
        print("Usage:")
        print("  python3 retry_manager.py <batch-index.json> status")
        print("  python3 retry_manager.py <batch-index.json> reset <batch_id>")
        print("  python3 retry_manager.py <batch-index.json> reset-all")
        sys.exit(1)

    filepath = sys.argv[1]
    action = sys.argv[2] if len(sys.argv) > 2 else "status"

    if not os.path.exists(filepath):
        print(f"❌ File {filepath} không tồn tại")
        sys.exit(1)

    if action == "status":
        show_status(filepath)
    elif action == "reset":
        if len(sys.argv) < 4:
            print("❌ Thiếu batch_id: python3 retry_manager.py <index> reset <batch_id>")
            sys.exit(1)
        reset_batch(filepath, sys.argv[3])
    elif action == "reset-all":
        reset_all_failed(filepath)
    else:
        print(f"❌ Action không hợp lệ: {action}")
        sys.exit(1)


if __name__ == "__main__":
    main()