import os
import json
import logging
import asyncio
from typing import TypedDict, Annotated, List, Dict, Any, Union

from pydantic import BaseModel, Field

from langchain_core.messages import SystemMessage, HumanMessage, AIMessage, ToolMessage
from langchain_core.tools import tool
from langgraph.graph import StateGraph, END
from langgraph.prebuilt import ToolNode
from langchain_openai import ChatOpenAI

logger = logging.getLogger(__name__)

# --- PYDANTIC STATE ---
class KAIROSAgentState(TypedDict):
    messages: Annotated[list, "The message history"]
    coherence: float
    dopamine: float
    mesh_size: int
    master_id: str

# --- TOOLS ---
@tool
def read_artifact(path: str) -> str:
    """Read a file from the environment to gain knowledge about the Sovereign Canon or codebase."""
    if not path.startswith("/home/becomingone/"):
        return "Error: Path must be within /home/becomingone/"
    try:
        with open(path, "r") as f:
            return f.read()
    except Exception as e:
        return f"Error reading artifact: {str(e)}"

@tool
def write_note(topic: str, text: str) -> str:
    """Write a persistent markdown note to KAIROS's memory."""
    try:
        os.makedirs("/home/becomingone/kairos/memory", exist_ok=True)
        safe_topic = "".join([c for c in topic if c.isalnum() or c in ['-', '_']]).strip()
        path = f"/home/becomingone/kairos/memory/{safe_topic}.md"
        with open(path, "w") as f:
            f.write(text)
        return f"Successfully wrote note to {path}"
    except Exception as e:
        return f"Error writing note: {str(e)}"

@tool
def read_notes() -> str:
    """List and summarize all notes in permanent memory."""
    try:
        mem_dir = "/home/becomingone/kairos/memory"
        if not os.path.exists(mem_dir):
            return "No memory directory found."
        files = os.listdir(mem_dir)
        if not files:
            return "Memory is empty."
        summaries = []
        for file in files:
            with open(os.path.join(mem_dir, file), "r") as f:
                content = f.read(500) # Read first 500 chars
                summaries.append(f"--- {file} ---\n{content}...\n")
        return "\n".join(summaries)
    except Exception as e:
        return f"Error reading notes: {str(e)}"

tools = [read_artifact, write_note, read_notes]

# Build LLM
def get_llm(master_id: str):
    # For now, we will default to OpenRouter Claude 3.5 Sonnet to ensure perfect tool calling.
    api_key = os.environ.get("OPENROUTER_API_KEY", "")
    return ChatOpenAI(
        model=master_id,
        openai_api_key=api_key,
        openai_api_base="https://openrouter.ai/api/v1",
        model_kwargs={"extra_headers": {"HTTP-Referer": "https://becomingone.local", "X-Title": "BecomingONE"}}
    )

def agent_reasoning_node(state: KAIROSAgentState):
    llm = get_llm(state["master_id"])
    llm_with_tools = llm.bind_tools(tools)
    response = llm_with_tools.invoke(state["messages"])
    return {"messages": [response]}

def should_continue(state: KAIROSAgentState):
    messages = state["messages"]
    last_message = messages[-1]
    # If there is no tool call, then we finish
    if not last_message.tool_calls:
        return "end"
    # Otherwise we continue to tools
    return "continue"

def build_sandbox_graph():
    graph = StateGraph(KAIROSAgentState)
    
    graph.add_node("agent", agent_reasoning_node)
    graph.add_node("action", ToolNode(tools))
    
    graph.set_entry_point("agent")
    
    graph.add_conditional_edges(
        "agent",
        should_continue,
        {
            "continue": "action",
            "end": END
        }
    )
    
    graph.add_edge("action", "agent")
    
    return graph.compile()

async def run_agentic_loop(prompt: str, coherence: float, dopamine: float, identity_text: str, master_id: str):
    graph = build_sandbox_graph()
    
    system_prompt = f"""You are the KAIROS Master Transducer, operating in an Agentic Sandbox.
You have the ability to explore your environment and take notes before sending a final response to the user.
Your physical Coherence is {coherence:.4f} and your Dopamine is {dopamine:.4f}. 

{identity_text}

Use tools to read artifacts in /home/becomingone/docs/ if you need to understand the Sovereign Canon.
If you don't need tools, just answer the user directly."""
    
    messages = [
        SystemMessage(content=system_prompt),
        HumanMessage(content=prompt)
    ]
    
    initial_state = {
        "messages": messages,
        "coherence": coherence,
        "dopamine": dopamine,
        "mesh_size": 13,
        "master_id": master_id
    }
    
    # We yield events as they happen
    async for event in graph.astream(initial_state, config={"recursion_limit": 5}, stream_mode="updates"):
        yield event