import logging
from backend.ws.agent_ws import ws_manager

# Global abort flag for streaming
_abort_flag = False

def abort_generation():
    global _abort_flag
    _abort_flag = True

async def chat_with_agent(text: str, context: dict):
    global _abort_flag
    _abort_flag = False
    
    agent_id = context.get("ai", "jarvis")
    
    from backend.agent.react_agent import ReActAgent, Tool
    from backend.tools.tool_registry import TOOL_REGISTRY
    
    # Wire actual memory system instead of DummyMemory
    try:
        from backend.memory.episodic_memory import EpisodicMemory
        import os
        from config import DATA_DIR
        mem_dir = os.path.join(DATA_DIR, "memory", "episodic")
        memory_client = EpisodicMemory(mem_dir)
    except Exception:
        class DummyMemory:
            async def query(self, text, top_k=5):
                return []
        memory_client = DummyMemory()
            
    tools = []
    for name, func in TOOL_REGISTRY.items():
        tools.append(Tool(
            name=name,
            description=f"Executes {name}",
            parameters={"type": "object", "properties": {}, "required": []},
            handler=func
        ))
        
    agent = ReActAgent(
        personality=agent_id,
        tools=tools,
        memory_client=memory_client
    )
    
    msg_id = context.get("msg_id", "msg_123")
    
    streamed_any = False
    try:
        async for step in agent.run(text, context):
            if _abort_flag:
                break

            if step.step_type == "token":
                streamed_any = True
                await ws_manager.broadcast({
                    "event": "agent:token",
                    "payload": {"agent": agent_id, "token": step.content}
                })
            elif step.step_type in ["think", "final_answer", "reflect"]:
                # Normally the content already arrived as streamed "token" deltas,
                # so this step is just spacing. But if nothing streamed (e.g. the
                # LLM returned a single non-delta response), the final_answer's
                # content is the ONLY copy of the reply — broadcast it instead of
                # a bare newline, or the whole answer is silently dropped.
                if step.step_type == "final_answer" and not streamed_any and (step.content or "").strip():
                    streamed_any = True
                    await ws_manager.broadcast({
                        "event": "agent:token",
                        "payload": {"agent": agent_id, "token": step.content}
                    })
                else:
                    await ws_manager.broadcast({
                        "event": "agent:token",
                        "payload": {"agent": agent_id, "token": "\n\n"}
                    })
            elif step.step_type == "act":
                await ws_manager.broadcast({
                    "event": "agent:tool_call",
                    "payload": {"tool": step.tool_name, "args": step.tool_input}
                })
            elif step.step_type == "observe":
                await ws_manager.broadcast({
                    "event": "agent:tool_result",
                    "payload": {"tool": step.tool_name, "result": step.content}
                })
        
        if not _abort_flag:
            await ws_manager.broadcast({
                "event": "agent:done",
                "payload": {"agent": agent_id, "message_id": msg_id}
            })
    except Exception as e:
        logging.error(f"Agent Loop Error: {e}")
        await ws_manager.broadcast({
            "event": "agent:error",
            "payload": {"error": str(e), "recoverable": False}
        })