#!/usr/bin/env python3
"""
YUI CPU FIX v2.0 — apply_patch.py
Fixes:
  1. Xóa duplicate _startup_bg_tasks thread (nguyên nhân test suite chạy 2x)
  2. Global flag _TEST_SUITE_DONE — chỉ chạy test 1 lần toàn bộ process
  3. LOG-SCANNER false positive — filter bỏ test suite output khỏi error buffer
  4. Cohere AI reply test — skip API call thật, chỉ check key existence
"""
import sys, re

APP = "/app/app.py"
if len(sys.argv) > 1:
    APP = sys.argv[1]

with open(APP, "r", encoding="utf-8", errors="replace") as f:
    src = f.read()

original_len = len(src)
changes = []

# ════════════════════════════════════════════════════════
# FIX 1: Xóa duplicate _startup_bg_tasks (cái đầu tiên)
# ════════════════════════════════════════════════════════
OLD_DUP_THREAD = """# Background: test suite + warm cache
def _startup_bg_tasks():
    import time as _sbt
    _sbt.sleep(3)
    try: _run_full_test_suite()
    except Exception as _te: print(f\"[TEST] {_te}\")
    _sbt.sleep(2)
    try:
        _kw_lists = [v for k,v in globals().items() if isinstance(v,list) and k.endswith(\"_signals\") or k.endswith(\"_keywords\") or k.endswith(\"_triggers\")]
        warm_cache(_kw_lists[:10])
    except Exception: pass
import threading as _sbt_thr2
_sbt_thr2.Thread(target=_startup_bg_tasks, daemon=True, name=\"StartupBg\").start()"""

NEW_DUP_THREAD = """# [CPU-FIX] Duplicate _startup_bg_tasks đã bị xóa — chỉ dùng cái ở dưới"""

if OLD_DUP_THREAD in src:
    src = src.replace(OLD_DUP_THREAD, NEW_DUP_THREAD, 1)
    changes.append("FIX 1 OK: Xóa duplicate _startup_bg_tasks thread")
else:
    changes.append("FIX 1 SKIP: Không tìm thấy duplicate thread (có thể đã fix rồi)")

# ════════════════════════════════════════════════════════
# FIX 2: Global flag _TEST_SUITE_DONE vào _run_full_test_suite
# ════════════════════════════════════════════════════════
OLD_TEST_SUITE_DEF = """def _run_full_test_suite():
    \"\"\"Chay 47 test kiem tra toan bo he thong, in bang ket qua chi tiet.\"\"\"
    import time as _t
    results = []
    p = [0]
    f = [0]"""

NEW_TEST_SUITE_DEF = """# [CPU-FIX] Global flag: chỉ chạy test suite 1 lần toàn bộ process
_TEST_SUITE_DONE = [False]
_TEST_SUITE_LOCK = __import__('threading').Lock()

def _run_full_test_suite():
    \"\"\"Chay 47 test kiem tra toan bo he thong, in bang ket qua chi tiet.\"\"\"
    # [CPU-FIX] Chỉ chạy 1 lần — tránh multi-worker và duplicate thread chạy lại
    with _TEST_SUITE_LOCK:
        if _TEST_SUITE_DONE[0]:
            print(\"[TEST] Đã chạy rồi — skip để tiết kiệm CPU\")
            return None, None
        _TEST_SUITE_DONE[0] = True
    import time as _t
    results = []
    p = [0]
    f = [0]"""

if OLD_TEST_SUITE_DEF in src:
    src = src.replace(OLD_TEST_SUITE_DEF, NEW_TEST_SUITE_DEF, 1)
    changes.append("FIX 2 OK: Global flag _TEST_SUITE_DONE added")
else:
    changes.append("FIX 2 SKIP: Không tìm thấy _run_full_test_suite def")

# ════════════════════════════════════════════════════════
# FIX 3: LOG-SCANNER false positive — filter test suite output
# ════════════════════════════════════════════════════════
OLD_PATCHED_PRINT = """def _patched_print(*args, **kwargs):
    _orig_print(*args, **kwargs)
    msg = \" \".join(str(a) for a in args)
    # Chỉ capture dòng có dấu hiệu lỗi
    if any(kw in msg for kw in [\"❌\",\"⚠️\",\"Error\",\"error\",\"Traceback\",\"lỗi\",\"loi\",\"failed\",\"Failed\",\"exception\"]):
        _log_error_to_buffer(msg)
_builtins.print = _patched_print"""

NEW_PATCHED_PRINT = """# [CPU-FIX] Test suite output không được log vào error buffer
_TEST_SUITE_SKIP_PATTERNS = [
    \"[v] \", \"[x] \", \"[!] \",           # Test result lines
    \"URAL AI — FULL SYSTEM\",            # Test suite header
    \"KET QUA TONG:\",                    # Test suite footer
    \"HANG MUC\",                         # Test suite column header
    \"SAN SANG PRODUCTION\",             # Test suite pass message
    \"PASS (\", \"FAIL: 0\",               # Test counts
    \"[TEST]\",                           # Test runner
    \"[LOG-SCANNER]\",                    # Scanner itself (avoid recursion)
]

def _patched_print(*args, **kwargs):
    _orig_print(*args, **kwargs)
    msg = \" \".join(str(a) for a in args)
    # [CPU-FIX] Bỏ qua test suite output — không nhầm là lỗi
    if any(skip in msg for skip in _TEST_SUITE_SKIP_PATTERNS):
        return
    # Chỉ capture dòng có dấu hiệu lỗi
    if any(kw in msg for kw in [\"❌\",\"⚠️\",\"Error\",\"error\",\"Traceback\",\"lỗi\",\"loi\",\"failed\",\"Failed\",\"exception\"]):
        _log_error_to_buffer(msg)
_builtins.print = _patched_print"""

if OLD_PATCHED_PRINT in src:
    src = src.replace(OLD_PATCHED_PRINT, NEW_PATCHED_PRINT, 1)
    changes.append("FIX 3 OK: LOG-SCANNER false positive filter added")
else:
    changes.append("FIX 3 SKIP: Không tìm thấy _patched_print (có thể đã fix)")

# ════════════════════════════════════════════════════════
# FIX 4: Cohere AI reply test — skip API call thật
# ════════════════════════════════════════════════════════
OLD_AI_TEST = """    # 9. QUICK AI REPLY
    def _ai_test():
        try:
            if _COHERE_GUEST:
                import cohere as _co2
                co2 = _co2.ClientV2(api_key=_COHERE_GUEST[0])
                r   = co2.chat(
                    model=_cohere_hunter.get_chat() or \"command-r-08-2024\",
                    messages=[{\"role\":\"user\",\"content\":\"Reply: OK\"}],
                    max_tokens=5
                )
                return bool(r.message.content[0].text)
        except Exception:
            return False
        return False
    test(\"Cohere AI reply\",   _ai_test, critical=False)"""

NEW_AI_TEST = """    # 9. QUICK AI REPLY — [CPU-FIX] Chỉ check key, không gọi API thật (tốn 12-15s mỗi lần)
    def _ai_test():
        # Chỉ verify key tồn tại — không gọi API để tiết kiệm CPU/quota
        return bool(_COHERE_GUEST) or bool(_COHERE_ADMIN)
    test(\"Cohere key check\",   _ai_test, critical=False)"""

if OLD_AI_TEST in src:
    src = src.replace(OLD_AI_TEST, NEW_AI_TEST, 1)
    changes.append("FIX 4 OK: Cohere AI reply test → lightweight key check")
else:
    changes.append("FIX 4 SKIP: Không tìm thấy _ai_test block")

# ════════════════════════════════════════════════════════
# FIX 5: Thêm worker-space detection — worker skip test suite
# ════════════════════════════════════════════════════════
OLD_STARTUP_BG = """# Background startup tasks — không block port 7860
# CODE MỚI ĐÃ SỬA CHỐNG LOOP LOG (Boss copy đoạn này dán đè vào):
def _startup_bg_tasks():
    import time as _sbt
    _sbt.sleep(3)
    
    try: 
        print(\"[TEST] Khởi chạy bộ kiểm thử hệ thống...\")
        _run_full_test_suite()
    except Exception as _te: 
        print(f\"[TEST] Lỗi kiểm thử: {_te}\")
        
    _sbt.sleep(2)
    try:
        _kw = [v for k,v in globals().items() if isinstance(v,list) and \"search_triggers\" in k]
        warm_cache(_kw)
        print(\"[KW-MATCH] ✅ Warm cache done\")
    except Exception: pass

import threading as _sbt_thr
_sbt_thr.Thread(target=_startup_bg_tasks, daemon=True, name=\"StartupBg\").start()"""

NEW_STARTUP_BG = """# Background startup tasks — không block port 7860
# [CPU-FIX v2.0] Chỉ chạy test suite trên MAIN SPACE, không chạy trên worker replicas
def _startup_bg_tasks():
    import time as _sbt, os as _os
    _sbt.sleep(5)  # Đợi app fully init
    
    # Detect worker space — worker không cần chạy full test suite
    _space_id = _os.environ.get(\"SPACE_ID\", \"\").lower()
    _is_worker = any(w in _space_id for w in [\"-quok\", \"-worker\", \"-backup\", \"-strong\", \"-creative\", \"-code\"])
    
    if not _is_worker:
        try: 
            print(\"[TEST] Khởi chạy bộ kiểm thử hệ thống...\")
            _run_full_test_suite()
        except Exception as _te: 
            print(f\"[TEST] Lỗi kiểm thử: {_te}\")
    else:
        print(f\"[TEST] Worker space ({_space_id}) — skip test suite để tiết kiệm CPU\")
        
    _sbt.sleep(2)
    try:
        _kw = [v for k,v in globals().items() if isinstance(v,list) and \"search_triggers\" in k]
        warm_cache(_kw)
        print(\"[KW-MATCH] ✅ Warm cache done\")
    except Exception: pass

import threading as _sbt_thr
_sbt_thr.Thread(target=_startup_bg_tasks, daemon=True, name=\"StartupBg\").start()"""

if OLD_STARTUP_BG in src:
    src = src.replace(OLD_STARTUP_BG, NEW_STARTUP_BG, 1)
    changes.append("FIX 5 OK: Worker space detection — worker skip test suite")
else:
    changes.append("FIX 5 SKIP: Không tìm thấy _startup_bg_tasks block chính")

# ════════════════════════════════════════════════════════
# SUMMARY
# ════════════════════════════════════════════════════════
print("\n" + "="*60)
print("  YUI CPU FIX v2.0 — KẾT QUẢ")
print("="*60)
for i, c in enumerate(changes, 1):
    icon = "✅" if "OK:" in c else "⚠️"
    print(f"  {icon} {c}")

ok_count = sum(1 for c in changes if "OK:" in c)
print(f"\n  Áp dụng thành công: {ok_count}/{len(changes)} fix")

if ok_count == 0:
    print("\n  ❌ Không có fix nào áp dụng được — kiểm tra lại source!")
    sys.exit(1)

# Backup và ghi
backup = APP + ".backup_cpufix"
import shutil
shutil.copy2(APP, backup)
print(f"\n  💾 Backup: {backup}")

with open(APP, "w", encoding="utf-8") as f:
    f.write(src)

new_len = len(src)
print(f"  📝 Đã ghi {APP} ({original_len:,} → {new_len:,} chars)")
print("\n  🎉 XONG! Restart Space để áp dụng fix.")
print("="*60)