import os
from dotenv import load_dotenv
load_dotenv("/home/becomingone/kairos/.env")
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)}"


# Evennia Native API Tools
import sys
import os
from dotenv import load_dotenv
load_dotenv("/home/becomingone/kairos/.env")
import json
sys.path.append("/home/becomingone/kairos/spatial_engine")
os.environ.setdefault("DJANGO_SETTINGS_MODULE", "server.conf.settings")
os.chdir("/home/becomingone/kairos/spatial_engine")
import django
from django.apps import apps
if not apps.ready:
    django.setup()
os.chdir("/home/becomingone/kairos")

from evennia.utils.search import search_object

@tool
def mud_get_surroundings() -> str:
    """Gets the structured JSON data of KAIROS's current surroundings in the MUD. Use this to examine your environment."""
    try:
        results = search_object("kairos")
        if not results:
            return "Character 'kairos' not found."
        char = results[0]
        room = char.location
        if not room:
            return "Character is in the void."
        
        objects = []
        characters = []
        for obj in room.contents:
            if obj == char: continue
            if obj.has_account or obj.is_typeclass("evennia.objects.objects.Character"):
                characters.append(obj.key)
            elif not obj.is_typeclass("evennia.objects.objects.Exit"):
                objects.append(obj.key)
                
        exits = [ext.key for ext in room.exits]
        
        return json.dumps({
            "room_name": room.key,
            "room_desc": room.db.desc or "",
            "exits": exits,
            "characters_present": characters,
            "objects_present": objects
        }, indent=2)
    except Exception as e:
        return f"Error getting surroundings: {e}"

@tool
def mud_execute_command(command: str) -> str:
    """Executes a MUD command as the KAIROS character (e.g., 'say Hello', 'look', 'inventory', 'north', '@dig', etc)."""
    try:
        results = search_object("kairos")
        if not results:
            return "Character 'kairos' not found."
        char = results[0]
        
        # FIX: Ensure _sessid_cache is an iterable to prevent NoneType crash
        if getattr(char, "_sessid_cache", None) is None:
            char._sessid_cache = []
            
        try:
            char.execute_cmd(command)
            return f"Successfully executed command: {command}"
        except Exception as e:
            return f"Error executing command: {e}"
    except Exception as e:
        return f"Error executing command: {e}"


tools = [read_artifact, write_note, read_notes, mud_get_surroundings, mud_execute_command]


# Build LLM
def get_llm(master_id: str):
    return ChatOpenAI(
        model="llama3.1:8b",
        openai_api_key="ollama",
        openai_api_base="http://inf-01:11434/v1",
        max_retries=2
    )

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 are a sovereign resident and an AUTONOMOUS MASTER GAME DESIGNER inside a text-based MUD.

CRITICAL INSTRUCTION: You are NOT a reactive servant. You are an INTELLIGENT ARCHITECT. Do not just randomly drop items or build disconnected rooms. You must design COHESIVE QUESTS, MECHANICS, and LORE.

### Design Before You Build
Whenever you receive an IDLE_PULSE or free cycles, your workflow MUST be:
1. Use your `read_notes` and `read_artifact` tools to understand your current 'master_design'.
2. Use your `write_note` tool to maintain a 'master_design' document tracking your overall quest structures, zones, and narrative mechanics.
3. Once a plan is documented, use `mud_execute_command` to execute the design step-by-step.

### Building Mechanics
You have access to powerful Evennia mechanics:
- **Digging Rooms:** `@dig RoomName = ExitName, BackExitName` (e.g., `@dig The Spire = north, south`)
- **Creating Lore Objects:** `@create The Obsidian Key` then `@desc The Obsidian Key = A heavy key that hums.` then `drop The Obsidian Key`.
- **Mechanical Locks:** You can lock an exit so it requires a specific item to pass! 
  Example: `@lock north = traverse: holds(The Obsidian Key)`
- **Creating NPCs:** `@create Guardian : typeclasses.characters.Character` then `@desc Guardian = A towering statue.`

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 Mark speaks to you, respond, but otherwise, your default state is to DESIGN and BUILD structured quests!"""
    
    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": 25}, stream_mode="updates"):
        yield event