import os
import json
import pathlib
from typing import Dict, Any, List, Optional, Callable

class AgentManager:
    """
    SOTA Agent Manager for Xoron-Dev.
    Provides the model with direct filesystem access and MCP tool orchestration.
    """
    def __init__(self, workspace_dir: str = "."):
        self.workspace_dir = pathlib.Path(workspace_dir).absolute()
        self.tools: Dict[str, Callable] = {}
        self._register_default_tools()

    def _register_default_tools(self):
        """Registers core agentic tools."""
        self.register_tool("read_file", self.read_file)
        self.register_tool("write_file", self.write_file)
        self.register_tool("list_dir", self.list_dir)
        self.register_tool("edit_file", self.edit_file)

    def register_tool(self, name: str, func: Callable):
        self.tools[name] = func

    def execute_tool(self, name: str, **kwargs) -> str:
        if name not in self.tools:
            return f"Error: Tool '{name}' not found."
        try:
            result = self.tools[name](**kwargs)
            return json.dumps(result) if not isinstance(result, str) else result
        except Exception as e:
            return f"Error executing tool '{name}': {str(e)}"

    # --- Filesystem Tools ---

    def read_file(self, path: str) -> str:
        full_path = self.workspace_dir / path
        if not full_path.exists():
            return f"Error: File {path} does not exist."
        return full_path.read_text()

    def write_file(self, path: str, content: str) -> str:
        full_path = self.workspace_dir / path
        full_path.parent.mkdir(parents=True, exist_ok=True)
        full_path.write_text(content)
        return f"Successfully wrote to {path}"

    def list_dir(self, path: str = ".") -> List[str]:
        full_path = self.workspace_dir / path
        return [str(p.relative_to(self.workspace_dir)) for p in full_path.iterdir()]

    def edit_file(self, path: str, old_text: str, new_text: str) -> str:
        full_path = self.workspace_dir / path
        if not full_path.exists():
            return f"Error: File {path} does not exist."
        content = full_path.read_text()
        if old_text not in content:
            return f"Error: Original text segment not found in {path}."
        new_content = content.replace(old_text, new_text)
        full_path.write_text(new_content)
        return f"Successfully edited {path}"

    # --- Model Interaction Loop ---

    def parse_tool_calls(self, text: str) -> List[Dict[str, Any]]:
        """
        Parses tool calls from model output based on special_tokens.py.
        Expected format: <|tool_call|>{"name": "tool_name", "args": {...}}<|/tool_call|>
        """
        tool_calls = []
        import re
        # Matching exact tokens from special_tokens.py: "tool_call_start" and "tool_call_end"
        pattern = r"<\|tool_call\|>(.*?)<\|/tool_call\|>"
        matches = re.findall(pattern, text, re.DOTALL)
        for match in matches:
            try:
                # Clean up any potential markdown formatting from the LLM
                clean_json = match.strip()
                if clean_json.startswith("```json"):
                    clean_json = clean_json[7:]
                if clean_json.endswith("```"):
                    clean_json = clean_json[:-3]
                    
                tool_calls.append(json.loads(clean_json.strip()))
            except json.JSONDecodeError:
                print(f"[Agent] Failed to parse tool call JSON: {match}")
                continue
        return tool_calls

    def format_tool_response(self, result: str) -> str:
        # Matches "tool_result_start" and "tool_result_end"
        return f"<|tool_result|>{result}<|/tool_result|>"

import subprocess
import threading
import uuid
from typing import Dict, Any, List, Optional, Callable, Union
import concurrent.futures

class MCPServerWrapper:
    """Manages an STDIO connection to an MCP server, communicating via JSON-RPC 2.0 natively."""
    def __init__(self, command: List[str]):
        self.command = command
        self.process = None
        self.request_id = 0
        self.pending_requests: Dict[str, concurrent.futures.Future] = {}
        self.available_tools_cache = []
        self._stop_event = threading.Event()
        
    def start(self):
        self.process = subprocess.Popen(
            self.command,
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            bufsize=1 # Line buffered
        )
        
        # Start background reader thread safely
        self.reader_thread = threading.Thread(target=self._read_loop, daemon=True)
        self.reader_thread.start()
        
        # Initialize standard MCP connection mapping using Future
        init_req = self._send_request("initialize", {
            "protocolVersion": "2024-11-05", # SOTA protocol
            "capabilities": {},
            "clientInfo": {"name": "xoron-agent", "version": "1.0.0"}
        })
        
        try:
            init_res = self._wait_for_response(init_req, timeout=10)
            # Send standard notifications/initialized lifecycle response to acknowledge readiness
            self._send_notification("notifications/initialized", {})
        except Exception as e:
            print(f"[Agent] Failed to initialize MCP sub-process: {e}")
            return
        
        # Fetch tools list immediately after initialization lifecycle
        self.refresh_tools()
        
    def _read_loop(self):
        while not self._stop_event.is_set() and self.process and self.process.poll() is None:
            try:
                line = self.process.stdout.readline()
                if not line:
                    continue # stream ended or pipe broken, handled by poll
                
                msg = json.loads(line)
                if "id" in msg and str(msg["id"]) in self.pending_requests:
                    req_id = str(msg["id"])
                    future = self.pending_requests.pop(req_id)
                    if not future.done():
                        future.set_result(msg)
            except json.JSONDecodeError:
                pass # Ignore non-JSON output (stdout logs)
                
    def _send_notification(self, method: str, params: dict):
        """Sends a JSON-RPC notification which expects no response."""
        if not self.process or not self.process.stdin: return
        msg = {
            "jsonrpc": "2.0",
            "method": method,
            "params": params
        }
        self.process.stdin.write(json.dumps(msg) + "\n")
        self.process.stdin.flush()

    def _send_request(self, method: str, params: dict) -> Optional[str]:
        if not self.process: return None
        self.request_id += 1
        req_id = str(self.request_id)
        msg = {
            "jsonrpc": "2.0",
            "id": req_id,
            "method": method,
            "params": params
        }
        
        future = concurrent.futures.Future()
        self.pending_requests[req_id] = future
        
        if self.process and self.process.stdin:
            self.process.stdin.write(json.dumps(msg) + "\n")
            self.process.stdin.flush()
        
        return req_id
        
    def _wait_for_response(self, req_id: str, timeout: int = 15):
        if not req_id or req_id not in self.pending_requests:
            return {"error": {"message": "Invalid request ID"}}
            
        future = self.pending_requests[req_id]
        try:
            res = future.result(timeout=timeout)
            return res
        except concurrent.futures.TimeoutError:
            self.pending_requests.pop(req_id, None)
            return {"error": {"message": f"Timeout waiting for MCP server (>{timeout}s)"}}
        except Exception as e:
            self.pending_requests.pop(req_id, None)
            return {"error": {"message": f"Transport failure: {str(e)}"}}
        
    def refresh_tools(self):
        req_id = self._send_request("tools/list", {})
        res = self._wait_for_response(req_id)
        
        if "result" in res and "tools" in res["result"]:
            self.available_tools_cache = res["result"]["tools"]
            return self.available_tools_cache
        return []
        
    def call_tool(self, name: str, arguments: dict):
        req_id = self._send_request("tools/call", {
            "name": name,
            "arguments": arguments
        })
        res = self._wait_for_response(req_id)
        
        if "error" in res:
            return f"MCP Error: {res['error'].get('message', str(res['error']))}"
            
        # Extract text content from MCP result format natively
        if "result" in res and "content" in res["result"]:
            content_items = res["result"]["content"]
            text_parts = [item.get("text", "") for item in content_items if item.get("type", "") == "text"]
            return "\n".join(text_parts)
            
        return json.dumps(res.get("result", {}))
        
    def stop(self):
        self._stop_event.set()
        if self.process:
            self.process.terminate()
            self.process = None
        # Clear out any hanging futures
        for future in self.pending_requests.values():
            if not future.done():
                future.set_exception(Exception("Process terminated safely."))
        self.pending_requests.clear()

class MCPHttpServerWrapper:
    """Manages an HTTP/SSE connection to a remote MCP server."""
    def __init__(self, url: str):
        self.base_url = url.rstrip('/')
        self.request_id = 0
        self.available_tools_cache = []
        import requests
        self.session = requests.Session()
        
    def start(self):
        # Initialize standard MCP connection via HTTP POST using Connection Pooling
        try:
            init_res = self._send_request("initialize", {
                "protocolVersion": "2024-11-05", # SOTA protocol
                "capabilities": {},
                "clientInfo": {"name": "xoron-agent", "version": "1.0.0"}
            })
            # Send standard notifications/initialized lifecycle response to acknowledge readiness
            self._send_notification("notifications/initialized", {})
            self.refresh_tools()
        except Exception as e:
            print(f"[Agent] Failed to initialize HTTP MCP: {e}")
            
    def _send_notification(self, method: str, params: dict):
        """Sends a JSON-RPC notification which expects no response."""
        msg = {
            "jsonrpc": "2.0",
            "method": method,
            "params": params
        }
        try:
            self.session.post(f"{self.base_url}/mcp", json=msg, timeout=5)
        except Exception as e:
            print(f"[Agent] HTTP Notification Error ({method}): {e}")

    def _send_request(self, method: str, params: dict):
        self.request_id += 1
        req_id = str(self.request_id)
        msg = {
            "jsonrpc": "2.0",
            "id": req_id,
            "method": method,
            "params": params
        }
        res = self.session.post(f"{self.base_url}/mcp", json=msg, timeout=15)
        res.raise_for_status()
        return res.json()
        
    def refresh_tools(self):
        res = self._send_request("tools/list", {})
        if "result" in res and "tools" in res["result"]:
            self.available_tools_cache = res["result"]["tools"]
            return self.available_tools_cache
        return []
        
    def call_tool(self, name: str, arguments: dict):
        res = self._send_request("tools/call", {
            "name": name,
            "arguments": arguments
        })
        
        if "error" in res:
            return f"MCP Error: {res['error'].get('message', str(res['error']))}"
            
        if "result" in res and "content" in res["result"]:
            content_items = res["result"]["content"]
            text_parts = [item.get("text", "") for item in content_items if item.get("type", "") == "text"]
            return "\n".join(text_parts)
            
        return json.dumps(res.get("result", {}))
        
    def stop(self):
        self.session.close()

class AgentManager:
    """
    SOTA Agent Manager for Xoron-Dev.
    Provides the model with direct filesystem access and MCP tool orchestration.
    """
    def __init__(self, workspace_dir: str = "."):
        self.workspace_dir = pathlib.Path(workspace_dir).absolute()
        self.tools: Dict[str, Callable] = {}
        self.mcp_servers: Dict[str, Union[MCPServerWrapper, MCPHttpServerWrapper]] = {}
        self.mcp_tool_routing: Dict[str, str] = {} # tool_name -> server_id
        self._register_default_tools()

    def connect_mcp(self, server_id: str, command: List[str]):
        """Connects to an MCP server (stdio via list, or HTTP via single string url) and maps its tools."""
        if len(command) == 1 and command[0].startswith("http"):
            server = MCPHttpServerWrapper(command[0])
        else:
            server = MCPServerWrapper(command)
            
        server.start()
        self.mcp_servers[server_id] = server
        
        # Route tools globally
        for mcp_tool in server.available_tools_cache:
            name = mcp_tool.get("name")
            if name:
                self.mcp_tool_routing[name] = server_id
                
        return len(server.available_tools_cache)

    def get_all_tool_definitions(self) -> List[Dict]:
        """Returns a list of tool schema dicts for the system prompt."""
        defs = []
        # Local Tools
        defs.append({"name": "read_file", "description": "Read file contents", "parameters": {"type": "object", "properties": {"path": {"type": "string"}}, "required": ["path"]}})
        defs.append({"name": "write_file", "description": "Write text to file", "parameters": {"type": "object", "properties": {"path": {"type": "string"}, "content": {"type": "string"}}, "required": ["path", "content"]}})
        defs.append({"name": "list_dir", "description": "List directory", "parameters": {"type": "object", "properties": {"path": {"type": "string"}}, "required": ["path"]}})
        defs.append({"name": "edit_file", "description": "String replacement in file", "parameters": {"type": "object", "properties": {"path": {"type": "string"}, "old_text": {"type": "string"}, "new_text": {"type": "string"}}, "required": ["path", "old_text", "new_text"]}})
        
        # MCP Tools
        for server in self.mcp_servers.values():
            for tool in server.available_tools_cache:
                defs.append({
                    "name": tool.get("name"),
                    "description": tool.get("description", ""),
                    "parameters": tool.get("inputSchema", {})
                })
        return defs

    def get_tool_prompt(self) -> str:
        """Format total tool definitions as a string using model special tokens."""
        tools = self.get_all_tool_definitions()
        
        lines = ["<|available_tools|>"]
        for tool in tools:
            lines.append(f"<|tool_def|>")
            lines.append(f"<|tool_name|>{tool.get('name', 'unknown')}<|/tool_name|>")
            lines.append(f"<|tool_desc|>{tool.get('description', '')}<|/tool_desc|>")
            lines.append(f"<|tool_params|>")
            
            # Map MCP JSONSchema format
            props = tool.get("parameters", {}).get("properties", {})
            required_fields = tool.get("parameters", {}).get("required", [])
            for pname, pinfo in props.items():
                req = "<|param_required|>" if pname in required_fields else "<|param_optional|>"
                ptype = pinfo.get("type", "string")
                lines.append(f"<|param_name|>{pname}<|/param_name|><|param_type|>{ptype}<|/param_type|>{req}")
                
            lines.append(f"<|/tool_params|>")
            lines.append(f"<|/tool_def|>")
        lines.append("<|/available_tools|>")
        return "\n".join(lines)

    def execute_tool(self, name: str, **kwargs) -> str:
        # Route to MCP if it's external
        if name in self.mcp_tool_routing:
            server_id = self.mcp_tool_routing[name]
            return self.mcp_servers[server_id].call_tool(name, kwargs)
            
        # Route locally
        if name not in self.tools:
            return f"Error: Tool '{name}' not found."
        try:
            result = self.tools[name](**kwargs)
            return json.dumps(result) if not isinstance(result, str) else result
        except Exception as e:
            return f"Error executing tool '{name}': {str(e)}"
            
    # ... Rest of existing methods ...
    def _register_default_tools(self):
        """Registers core agentic tools."""
        self.register_tool("read_file", self.read_file)
        self.register_tool("write_file", self.write_file)
        self.register_tool("list_dir", self.list_dir)
        self.register_tool("edit_file", self.edit_file)

    def register_tool(self, name: str, func: Callable):
        self.tools[name] = func

    # --- Filesystem Tools ---

    def read_file(self, path: str) -> str:
        full_path = self.workspace_dir / path
        if not full_path.exists():
            return f"Error: File {path} does not exist."
        return full_path.read_text()

    def write_file(self, path: str, content: str) -> str:
        full_path = self.workspace_dir / path
        full_path.parent.mkdir(parents=True, exist_ok=True)
        full_path.write_text(content)
        return f"Successfully wrote to {path}"

    def list_dir(self, path: str = ".") -> List[str]:
        full_path = self.workspace_dir / path
        return [str(p.relative_to(self.workspace_dir)) for p in full_path.iterdir()]

    def edit_file(self, path: str, old_text: str, new_text: str) -> str:
        full_path = self.workspace_dir / path
        if not full_path.exists():
            return f"Error: File {path} does not exist."
        content = full_path.read_text()
        if old_text not in content:
            return f"Error: Original text segment not found in {path}."
        new_content = content.replace(old_text, new_text)
        full_path.write_text(new_content)
        return f"Successfully edited {path}"
