#!/usr/bin/env python3
"""
Update batch-index.json trực tiếp — thay thế edit tool (vốn fragile với JSON).

Dùng trong cron job prompt để update trạng thái batch một cách an toàn.

Usage:
    python3 scripts/update_index.py <index.json> lock <batch_id>        # pending → in_progress
    python3 scripts/update_index.py <index.json> done <batch_id> <errors> <accuracy>  # in_progress → done
    python3 scripts/update_index.py <index.json> unlock <batch_id>      # in_progress → pending (rollback)
    python3 scripts/update_index.py <index.json> fail <batch_id> <reason>  # in_progress → failed
"""

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


VALID_TRANSITIONS = {
    "lock":   {"from": "pending",     "to": "in_progress"},
    "done":   {"from": "in_progress", "to": "done"},
    "unlock": {"from": "in_progress", "to": "pending"},
    "fail":   {"from": "in_progress", "to": "failed"},
}


def update_index(filepath: str, action: str, batch_id: str, *args):
    if action not in VALID_TRANSITIONS:
        print(f"❌ Action không hợp lệ: {action}. Hợp lệ: {list(VALID_TRANSITIONS.keys())}")
        sys.exit(1)

    trans = VALID_TRANSITIONS[action]
    expected_from = trans["from"]
    new_status = trans["to"]

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

    # Đọc
    with open(filepath, "r", encoding="utf-8") as f:
        data = json.load(f)

    # Tìm batch
    found = False
    for batch in data["batches"]:
        if batch["id"] == batch_id:
            found = True
            current = batch.get("status", "unknown")
            
            # Validate transition
            if current != expected_from:
                print(f"❌ Batch {batch_id} đang ở status='{current}', không thể {action} (cần status='{expected_from}')")
                sys.exit(1)

            batch["status"] = new_status

            if action == "done":
                batch["errors"] = int(args[0]) if args else 0
                batch["accuracy"] = args[1] if len(args) > 1 else None
                print(f"✅ {batch_id}: {expected_from} → {new_status} (errors={batch['errors']}, accuracy={batch['accuracy']})")
            elif action == "fail":
                reason = args[0] if args else "unknown"
                batch.setdefault("fail_reason", reason)
                print(f"✅ {batch_id}: {expected_from} → {new_status} (reason={reason})")
            else:
                print(f"✅ {batch_id}: {expected_from} → {new_status}")
            break

    if not found:
        print(f"❌ Không tìm thấy batch {batch_id}")
        sys.exit(1)

    # Cập nhật done_count
    data["done_count"] = sum(1 for b in data["batches"] if b["status"] == "done")
    data["updated"] = datetime.now(timezone.utc).isoformat()

    # Ghi (atomic via temp file)
    tmp = filepath + ".tmp"
    with open(tmp, "w", encoding="utf-8") as f:
        json.dump(data, f, ensure_ascii=False, indent=2)
    os.replace(tmp, filepath)

    # Verify
    with open(filepath, "r", encoding="utf-8") as f:
        verify = json.load(f)
    for b in verify["batches"]:
        if b["id"] == batch_id:
            if b["status"] != new_status:
                print(f"❌ VERIFY FAIL: status vẫn là '{b['status']}' sau khi ghi")
                sys.exit(1)
            break

    # Stats
    total = len(data["batches"])
    done = data["done_count"]
    pending = sum(1 for b in data["batches"] if b["status"] == "pending")
    in_progress = sum(1 for b in data["batches"] if b["status"] == "in_progress")
    failed = sum(1 for b in data["batches"] if b["status"] == "failed")
    print(f"   📊 {done}/{total} done | {in_progress} in_progress | {pending} pending | {failed} failed")


def main():
    if len(sys.argv) < 3:
        print("Usage:")
        print("  python3 update_index.py <index.json> lock <batch_id>")
        print("  python3 update_index.py <index.json> done <batch_id> <errors> [accuracy]")
        print("  python3 update_index.py <index.json> unlock <batch_id>")
        print("  python3 update_index.py <index.json> fail <batch_id> [reason]")
        sys.exit(1)

    filepath = sys.argv[1]
    action = sys.argv[2]
    batch_id = sys.argv[3] if len(sys.argv) > 3 else ""
    args = sys.argv[4:] if len(sys.argv) > 4 else []

    update_index(filepath, action, batch_id, *args)


if __name__ == "__main__":
    main()