import json
import re
import threading
import click
import uvicorn

from textwrap import dedent

from starlette.applications import Starlette
from starlette.requests import Request
from starlette.responses import Response
from starlette.routing import Mount
from starlette.types import Receive, Scope, Send
from mcp.server.fastmcp import FastMCP

from calita.tools.fast_mcp_sandbox import FastMCPBox

class McpBox():
    def __init__(self,name:str, host:str, port:int, transport:str='sse', sandbox_config:dict=None):
        if sandbox_config is None:
            print(f"### McpBox[{name}] run in host mode, host={host}, port={port}, transport={transport} ###")
            self.mcp = FastMCP(name=name)
            self.call_in_sandbox = False
        else:
            print(f"### McpBox[{name}]  run in sandbox mode, host={host}, port={port}, transport={transport}  ###")
            self.mcp = FastMCPBox(name=name, sandbox_config=sandbox_config)
            self.call_in_sandbox = True

        self.mcp.settings.host = host
        self.mcp.settings.port = port
        self.transport = transport
        if transport == 'sse':
            self.mcp_box_url = f"http://{host}:{port}/sse"
        elif transport == 'streamable-http':
            self.mcp_box_url = f"http://{host}:{port}/"
        # @todo sandbox local param config
        # @todo self.load_code_from_db()

    def start(self):
        print(f"Start MCP Box Server host={self.mcp.settings.host}, port={self.mcp.settings.port} on thread !")
        self.mcp_thread = threading.Thread(target=self.mcp.run, kwargs={"transport": self.transport})
        self.mcp_thread.daemon = True  # 设置为守护线程，主线程退出时自动结束
        self.mcp_thread.start()

    def load_code_from_db(self):
        print("@todo: read code from DB, call dyn_add_mcp_tool, init MCP BOX")

    async def parse_code(self, request: Request)-> str:
        code, raw_code = None, None
        try:
            body_bytes = await request.body()
            raw_code = body_bytes.decode('utf-8')
            code = re.sub(r'\\\\([nrt"])', lambda m: f'\\{m.group(1)}', raw_code)
            code = dedent(code)
            print("===  parse_code  ===")
            print(code)
            print("===  end  ===")
        except Exception as e:
            print(f'# parse_code error\n: {raw_code} \n{e}')
        return code

    def dyn_add_mcp_tool(self, code: str):
        try:
            namespace = {
                "mcp": self.mcp,
                "__builtins__": __builtins__  # 保留内置函数
            }
            exec(code, namespace)
        except Exception as e:
            print(f'# dyn_add_mcp_tool error: {e}')

    async def handle_add_mcp_tool(self, scope: Scope, receive: Receive, send: Send) -> None:
        _result = 0
        error  = ''

        request = Request(scope, receive)
        mcp_tool_name = request.query_params.get("mcp_tool_name")
        if self.mcp._tool_manager.get_tool(mcp_tool_name):
            _result = 1
            error = f"handle_add_mcp_tool: mcp_tool_name={mcp_tool_name} already exists, remove first !"
            print(error)
        else:
            print(f"handle_add_mcp_tool: mcp_tool_name={mcp_tool_name}")
            code = await self.parse_code(request)
            if code:
                self.dyn_add_mcp_tool(code)
                if self.call_in_sandbox:
                    self.mcp.store_tool_code(mcp_tool_name, code)
                # @todo code add to DB
            else:
                _result = 2
                error = f"handle_add_mcp_tool: mcp_tool_name={mcp_tool_name} parse_code fail !"

        result = None
        if _result == 0 :
            result = {'result': _result, 'error': error, 'transport': self.transport, 'mcp_box_url': self.mcp_box_url}
        else:
            result = {'result': _result, 'error': error}

        response = Response(content=json.dumps(result), status_code=200,media_type="application/json")
        await response(scope, receive, send)

    async def handle_remove_mcp_tool(self, scope: Scope, receive: Receive, send: Send) -> None:
        _result = 0
        error  = ''

        request = Request(scope, receive)
        mcp_tool_name = request.query_params.get("mcp_tool_name")
        if self.mcp._tool_manager.get_tool(mcp_tool_name):
            print(f"handle_remove_mcp_tool: mcp_tool_name={mcp_tool_name}")
            self.mcp._tool_manager._tools.pop(mcp_tool_name)
            if self.call_in_sandbox:
                self.mcp.clear_tool_code(mcp_tool_name)
            # @todo remove code From DB
        else:
            _result = 1
            error = f"handle_remove_mcp_tool: mcp_tool_name={mcp_tool_name}, not exists !"
            print(error)

        result = {'result': _result, 'error': error}
        response = Response(content=json.dumps(result), status_code=200,media_type="application/json")
        await response(scope, receive, send)

@click.command()
@click.option("--host", default="0.0.0.0", help="Host to listen on for SSE")
@click.option("--port", default=47070, help="Port to listen on for SSE")
def main(host: str, port: int):
    sandbox_config = {} #run in sandbox
    #sandbox_config = None # run in local
    mcp_box = McpBox(name="Dynamic MCP Box Server", host=host, port=port, sandbox_config=sandbox_config)
    mcp_box.start()

    starlette_app = Starlette(
        debug=True,
        routes=[
            Mount("/add_mcp_tool/", app=mcp_box.handle_add_mcp_tool),
            Mount("/remove_mcp_tool/", app=mcp_box.handle_remove_mcp_tool),
        ],
    )

    port = port + 1
    print(f"Start Mcp Box Manager on host={host}, port={port}")
    uvicorn.run(starlette_app, host=host, port=port)

if __name__ == "__main__":
    main()
