#!/usr/bin/env python3
import asyncio
import time
import uuid
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import Optional, Dict, Any
from fastmcp import Client

app = FastAPI(title="MCP Tool Executor", description="HTTP serwer do wykonywania narzędzi MCP")

class ToolRequest(BaseModel):
    tool: str
    params: Optional[Dict[str, Any]] = {}

class ToolResponse(BaseModel):
    success: bool
    result: str

def _ms(start_ns: int, end_ns: int) -> float:
    return (end_ns - start_ns) / 1_000_000.0

def _now_ns() -> int:
    return time.perf_counter_ns()

def _log(rid: str, msg: str):
    print(f"[MCP][{rid}] {msg}")

class _ClientState:
    __slots__ = ("uri", "client", "open_lock", "call_lock", "is_open")
    def __init__(self, uri: str, client: Client):
        self.uri = uri
        self.client = client
        self.open_lock = asyncio.Lock()
        self.call_lock = asyncio.Lock()
        self.is_open = False

_clients: Dict[str, _ClientState] = {}
_clients_lock = asyncio.Lock()

async def _open_client_if_needed(state: _ClientState, rid: str):
    async with state.open_lock:
        if state.is_open:
            return

        aenter = getattr(state.client, "__aenter__", None)
        if not callable(aenter):
            raise Exception("Klient MCP nie wspiera __aenter__ (nie da się utrzymać sesji otwartej).")

        t0 = _now_ns()
        await aenter()
        t1 = _now_ns()
        state.is_open = True
        _log(rid, f"OPEN_SESSION {state.uri} took {_ms(t0, t1):.3f} ms")

async def _close_client_if_open(state: _ClientState, rid: str):
    async with state.open_lock:
        if not state.is_open:
            return

        try:
            close_fn = getattr(state.client, "close", None)
            if callable(close_fn):
                t0 = _now_ns()
                await close_fn()
                t1 = _now_ns()
                state.is_open = False
                _log(rid, f"CLOSE_SESSION(close) {state.uri} took {_ms(t0, t1):.3f} ms")
                return

            aexit = getattr(state.client, "__aexit__", None)
            if callable(aexit):
                t0 = _now_ns()
                await aexit(None, None, None)
                t1 = _now_ns()
                state.is_open = False
                _log(rid, f"CLOSE_SESSION(__aexit__) {state.uri} took {_ms(t0, t1):.3f} ms")
                return

            state.is_open = False
            _log(rid, f"CLOSE_SESSION {state.uri} no close method found, marked closed")

        except Exception as e:
            state.is_open = False
            _log(rid, f"CLOSE_SESSION ERROR {state.uri}: {e}")

async def _replace_client(state: _ClientState, rid: str):
    """Zamyka starego klienta i tworzy nowego. Wywolywane przy retry."""
    await _close_client_if_open(state, rid)
    t0 = _now_ns()
    state.client = Client(state.uri)
    state.is_open = False
    t1 = _now_ns()
    _log(rid, f"REPLACE_CLIENT {state.uri} new Client created in {_ms(t0, t1):.3f} ms")
    await _open_client_if_needed(state, rid)

async def _get_or_create_state(uri: str, rid: str) -> _ClientState:
    t_lock_wait_start = _now_ns()
    async with _clients_lock:
        t_lock_acquired = _now_ns()
        lock_wait_ms = _ms(t_lock_wait_start, t_lock_acquired)

        state = _clients.get(uri)
        if state is not None:
            _log(rid, f"GET_STATE HIT {uri}, lock_wait={lock_wait_ms:.3f} ms")
            return state

        t_create_start = _now_ns()
        client = Client(uri)
        state = _ClientState(uri, client)
        _clients[uri] = state
        t_create_end = _now_ns()
        _log(rid, f"GET_STATE MISS {uri}, lock_wait={lock_wait_ms:.3f} ms, create={_ms(t_create_start, t_create_end):.3f} ms")
        return state

async def _close_all_clients():
    rid = "shutdown"
    async with _clients_lock:
        states = list(_clients.values())
        _clients.clear()

    for state in states:
        await _close_client_if_open(state, rid)

@app.on_event("shutdown")
async def on_shutdown():
    await _close_all_clients()

async def execute_mcp_tool(tool_name: str, params: Optional[Dict[str, Any]] = None) -> str:
    rid = uuid.uuid4().hex[:8]
    t0 = _now_ns()
    _log(rid, f"START tool_name={tool_name}, params_keys={list((params or {}).keys())}")

    try:
        t_parse_start = _now_ns()

        prefix = "mcp__"
        if not tool_name.startswith(prefix):
            raise Exception(f"Niepoprawna nazwa narzędzia: {tool_name}")

        port_start = len(prefix)
        first_us = tool_name.find("_", port_start)
        if first_us == -1:
            raise Exception(f"Brak '_' po porcie w nazwie narzędzia: {tool_name}")

        port_str = tool_name[port_start:first_us]
        if not port_str.isdigit():
            raise Exception(f"Port nie jest liczbą: {port_str}")

        port = int(port_str)
        uri = f"http://localhost:{port}/mcp"

        dd = tool_name.find("__", first_us + 1)
        if dd == -1:
            raise Exception(f"Brak '__' po porcie w nazwie narzędzia: {tool_name}")

        new_tool_name = tool_name[dd + 2:]

        t_parse_end = _now_ns()
        _log(
            rid,
            f"PARSE ok in {_ms(t_parse_start, t_parse_end):.3f} ms "
            f"(port={port}, uri={uri}, new_tool_name={new_tool_name})"
        )

        t_get_state_start = _now_ns()
        state = await _get_or_create_state(uri, rid)
        t_get_state_end = _now_ns()
        _log(rid, f"GET_STATE total {_ms(t_get_state_start, t_get_state_end):.3f} ms (+{_ms(t0, t_get_state_end):.3f} ms from start)")

        await _open_client_if_needed(state, rid)

        async with state.call_lock:
            t_call_start = _now_ns()
            try:
                result = await state.client.call_tool(new_tool_name, params or {})
            except Exception as e:
                t_fail = _now_ns()
                _log(rid, f"CALL call_tool({new_tool_name}) FAILED after {_ms(t_call_start, t_fail):.3f} ms: {e}")

                await _replace_client(state, rid)

                t_retry_start = _now_ns()
                result = await state.client.call_tool(new_tool_name, params or {})
                t_retry_end = _now_ns()
                _log(rid, f"CALL_RETRY call_tool({new_tool_name}) {_ms(t_retry_start, t_retry_end):.3f} ms (+{_ms(t0, t_retry_end):.3f} ms from start)")
            else:
                t_call_end = _now_ns()
                _log(rid, f"CALL call_tool({new_tool_name}) {_ms(t_call_start, t_call_end):.3f} ms (+{_ms(t0, t_call_end):.3f} ms from start)")

        t_extract_start = _now_ns()
        text = result.content[0].text
        t_extract_end = _now_ns()
        _log(rid, f"EXTRACT result.content[0].text {_ms(t_extract_start, t_extract_end):.3f} ms (+{_ms(t0, t_extract_end):.3f} ms from start)")

        t_total_end = _now_ns()
        _log(rid, f"SUCCESS total {_ms(t0, t_total_end):.3f} ms")
        return text

    except Exception as e:
        t_err = _now_ns()
        _log(rid, f"ERROR after {_ms(t0, t_err):.3f} ms: {e}")
        raise HTTPException(status_code=500, detail=f"Błąd MCP: {str(e)}")

@app.get("/")
async def root():
    return {
        "message": "MCP Tool Executor Server",
        "endpoints": {
            "/execTool": "POST - wykonaj narzędzie MCP"
        }
    }

@app.post("/execTool", response_model=ToolResponse)
async def exec_tool(request: ToolRequest):
    rid = uuid.uuid4().hex[:8]
    t_req_start = _now_ns()
    _log(rid, f"HTTP /execTool START tool={request.tool}, params_keys={list((request.params or {}).keys())}")

    try:
        result = await execute_mcp_tool(request.tool, request.params)
        t_req_end = _now_ns()
        _log(rid, f"HTTP /execTool OK total {_ms(t_req_start, t_req_end):.3f} ms")
        return ToolResponse(success=True, result=result)
    except HTTPException:
        t_req_end = _now_ns()
        _log(rid, f"HTTP /execTool HTTPException total {_ms(t_req_start, t_req_end):.3f} ms")
        raise
    except Exception as e:
        t_req_end = _now_ns()
        _log(rid, f"HTTP /execTool ERROR total {_ms(t_req_start, t_req_end):.3f} ms: {e}")
        raise HTTPException(status_code=500, detail=f"Błąd serwera: {str(e)}")

if __name__ == "__main__":
    import uvicorn
    print("Uruchamianie HTTP serwera MCP Tool Executor...")
    print("Dostepne endpointy:")
    print("   GET  /         - strona glowna")
    print("   POST /execTool - wykonaj narzedzie")
    print()
    print("Serwer bedzie dostepny na: http://localhost:8122")
    print("Dokumentacja API: http://localhost:8122/docs")

    uvicorn.run(app, host="0.0.0.0", port=8122)