import os
import sys

lab = r"C:\Users\intel\AppData\Local\OptimusRuntime\Labs\AR74\12_ar88_authority_operational_closure"
candidate_ws = os.path.join(lab, "candidate_workspace")

target_file = os.path.join(candidate_ws, "artifacts", "brain_host", "brain_host_tls.py")

with open(target_file, "r", encoding="utf-8") as f:
    content = f.read()

target_str = """                                    logger.info("[BrainHost] Dispatched live shadow hook for task %s, result calling dispatch...", tid)
                                    # Forward direct in-process reference to OPERATIONAL_TASKS so hook updates live dictionary directly
                                    _snap["operational_tasks"] = OPERATIONAL_TASKS
                                    dispatch_live_shadow_hook(_snap, async_dispatch=True)"""

# Safe updater callable registered for tid
new_str = """                                    logger.info("[BrainHost] Dispatched live shadow hook for task %s, result calling dispatch...", tid)
                                    def _canonical_task_bridge_updater(update_payload):
                                        if not isinstance(update_payload, dict):
                                            return False
                                        target_tid = str(update_payload.get("task_id") or update_payload.get("mission_id") or tid).strip()
                                        with OPERATIONAL_TASKS_LOCK:
                                            if target_tid in OPERATIONAL_TASKS:
                                                t_entry = OPERATIONAL_TASKS[target_tid]
                                                new_st = update_payload.get("task_status") or update_payload.get("status")
                                                if new_st:
                                                    t_entry["task_status"] = new_st
                                                    t_entry["status"] = new_st
                                                    t_entry["state"] = update_payload.get("state") or new_st
                                                    t_entry["canonical_terminal_status"] = new_st
                                                act_ev = update_payload.get("activity_event")
                                                if act_ev:
                                                    t_entry[act_ev.lower()] = True
                                                    activities = t_entry.setdefault("activity", [])
                                                    activities.append({
                                                        "timestamp": time.time(),
                                                        "event": act_ev,
                                                        "status": update_payload.get("status") or t_entry.get("status"),
                                                        "message": update_payload.get("message") or act_ev,
                                                        "incident_id": update_payload.get("incident_id"),
                                                        "shadow_run_id": update_payload.get("shadow_run_id"),
                                                    })
                                                return True
                                        return False

                                    try:
                                        from app.remediation.self_repair.live_shadow_hook import register_canonical_task_state_updater
                                        register_canonical_task_state_updater(_canonical_task_bridge_updater, mission_id=tid, task_id=tid)
                                    except Exception as _ex_reg:
                                        logger.debug("[BrainHost] Could not register canonical task state updater: %s", _ex_reg)

                                    _snap["operational_tasks"] = OPERATIONAL_TASKS
                                    dispatch_live_shadow_hook(_snap, async_dispatch=True, task_state_updater_callable=_canonical_task_bridge_updater)"""

if target_str in content:
    content = content.replace(target_str, new_str)
    with open(target_file, "w", encoding="utf-8") as f:
        f.write(content)
    print("SUCCESS: brain_host_tls.py patched with canonical task bridge updater!")
else:
    print("ERROR: target_str not found in brain_host_tls.py")