import re
import unicodedata
import time
import copy
from typing import List
from inspect import Signature, Parameter

from pydantic import AnyUrl
from sqlalchemy.ext.asyncio import async_sessionmaker, AsyncSession

# MCP SDK imports
from mcp.server.fastmcp import Context
from mcp.server.fastmcp.tools import Tool as FastMCPTool
from mcp.server.fastmcp.resources import FunctionResource

# Helper imports
from rhea.utils.schema import Tool, Inputs
from rhea.server.schema import MCPOutput, MCPDataOutput, Settings, AgentState
from rhea.agent.schema import RheaParam, RheaOutput
from rhea.utils.models import get_galaxytool_by_id
from rhea.manager.utils import get_handle_from_redis
from rhea.manager.launch_agent import launch_agent
import rhea.server.metrics as metrics

# ProxyStore imports
from proxystore.connectors.redis import RedisKey
from proxystore.store import Store
from redis import Redis

# Academy imports
from academy.handle import UnboundRemoteHandle, RemoteHandle


def construct_params(inputs: Inputs) -> List[Parameter]:
    params = [param.to_python_parameter() for param in inputs.params]
    if inputs.conditionals is not None:
        for cond in inputs.conditionals:
            params.extend(cond.to_python_parameter())

    # Deduplicate by name, preserving first occurrence
    seen = {}
    for p in params:
        if p.name not in seen:
            seen[p.name] = p
    params = list(seen.values())

    # Split into those without a default, and those with one to prevent 'non-default argument follows default argument'
    no_default = [p for p in params if p.default is Parameter.empty]
    with_default = [p for p in params if p.default is not Parameter.empty]

    return no_default + with_default


def process_user_inputs(tool: Tool, args: dict) -> List[RheaParam]:
    res = []

    for param in tool.inputs.params:
        if param.name is None and param.argument is not None:
            param.name = param.argument.replace("--", "")

        a = args.get(param.name, None)

        if a is not None:
            if param.type == "data":
                res.append(RheaParam.from_param(param, RedisKey(a)))
            else:
                res.append(RheaParam.from_param(param, a))

    if tool.inputs.conditionals is not None:
        for conditional in tool.inputs.conditionals:
            value = args.get(conditional.param.name)
            if value is not None:
                conditional_p = RheaParam.from_param(conditional.param, value)
                res.append(conditional_p)
                conditional_cp = copy.copy(conditional_p)
                conditional_cp.name = f"{conditional.name}.{conditional_p.name}"
                res.append(conditional_cp)
                for when in conditional.whens:
                    if when.value == value:
                        for p in when.params:
                            if arg_val := args.get(p.name):
                                if p.type == "data":
                                    arg_val = RedisKey(redis_key=arg_val)
                                rp = RheaParam.from_param(p, arg_val)
                            else:
                                rp = RheaParam.from_param(p, p.value)
                            res.append(rp)
                            cp = copy.copy(rp)
                            cp.name = f"{conditional.name}.{rp.name}"
                            res.append(cp)

    return res


def sanitize_tool_name(text: str, repl: str = "_") -> str:
    if len(repl) != 1 or not re.match(r"[A-Za-z0-9_-]", repl):
        raise ValueError("`repl` must be a single allowed character.")

    text = unicodedata.normalize("NFKD", text).encode("ascii", "ignore").decode("ascii")
    text = re.sub(r"[^A-Za-z0-9_-]+", repl, text)
    text = re.sub(re.escape(repl) + r"+", repl, text)

    return text.strip(repl + "-")


def create_proxystore_function_resource(
    file: MCPDataOutput, ctx: Context
) -> FunctionResource:
    async def _fetch() -> bytes:
        redis_key = RedisKey(redis_key=file.key)
        output_store: Store = ctx.request_context.lifespan_context.output_store

        res = output_store.get(redis_key)
        if res is None:
            raise ValueError(f"Key {file.key} not in ProxyStore")

        return res

    return FunctionResource(
        uri=AnyUrl(f"proxystore://{file.key}"),
        name=file.filename,
        title=file.name,
        description=f"Output file {file.filename} from tool call",
        mime_type="text/plain",
        fn=_fetch,
    )


def create_tool(tool: Tool, ctx: Context) -> FastMCPTool:
    params: List[Parameter] = []

    # Add Context to tool params
    params.append(
        Parameter("ctx", kind=Parameter.POSITIONAL_OR_KEYWORD, annotation=Context)
    )

    params += construct_params(tool.inputs)

    sig = Signature(parameters=params, return_annotation=MCPOutput)

    def make_wrapper(tool_id, param_names):
        async def wrapper(*args, **kwargs):
            try:
                # Log tool call
                metrics.tool_execution_request_count.inc()

                start_time = time.time()

                # Get context
                ctx: Context = kwargs.pop("ctx")
                await ctx.info(f"Launching tool {tool_id}")

                # Get settings from app context
                settings: Settings = ctx.request_context.lifespan_context.settings

                # Configure Redis client to get handle
                r = Redis(settings.redis_host, settings.redis_port)

                run_id = ctx.request_context.lifespan_context.run_id

                db_sessionmaker: async_sessionmaker[AsyncSession] = (
                    ctx.request_context.lifespan_context.db_sessionmaker
                )

                async with db_sessionmaker() as session:
                    tool: Tool | None = await get_galaxytool_by_id(session, tool_id)

                if tool is None:
                    raise RuntimeError(f"No tool found with ID: {tool_id}")

                await ctx.report_progress(0, 1)

                for name, value in zip(param_names, args):
                    kwargs.setdefault(name, value)

                # Construct RheaParams
                rhea_params = process_user_inputs(tool, kwargs)

                await ctx.report_progress(0.05, 1)

                if tool_id not in ctx.request_context.lifespan_context.agents:
                    # First, quickly check if the agent exists in other contexts
                    unbound_handle: UnboundRemoteHandle | None = (
                        await get_handle_from_redis(tool.id, run_id, r, timeout=1)
                    )

                    # Another context already initialized this tool, bind it to this Academy client:
                    if unbound_handle is not None:
                        handle: RemoteHandle = unbound_handle.bind_to_client(
                            ctx.request_context.lifespan_context.academy_client
                        )
                        ctx.request_context.lifespan_context.agents[tool_id] = (
                            AgentState(tool_id=tool_id, handle=handle)
                        )

                    else:
                        # Launch agent
                        launch_agent(
                            tool,
                            run_id=run_id,
                            container_runtime=settings.parsl_container_backend,
                            redis_host=settings.agent_redis_host,
                            redis_port=settings.agent_redis_port,
                            minio_endpoint=settings.minio_endpoint,
                            minio_access_key=settings.minio_access_key,
                            minio_secret_key=settings.minio_secret_key,
                            minio_secure=False,
                        )

                        unbound_handle: UnboundRemoteHandle | None = (
                            await get_handle_from_redis(
                                tool.id,
                                run_id,
                                r,
                                timeout=settings.agent_handle_timeout,
                            )
                        )

                        if unbound_handle is None:
                            raise RuntimeError(
                                "Never received handle from Parsl worker."
                            )

                        handle: RemoteHandle = unbound_handle.bind_to_client(
                            ctx.request_context.lifespan_context.academy_client
                        )

                        await ctx.info(f"Lanched agent {handle.agent_id}")

                        ctx.request_context.lifespan_context.agents[tool_id] = (
                            AgentState(tool_id=tool_id, handle=handle)
                        )

                # Get handle from dictionary
                handle: RemoteHandle = ctx.request_context.lifespan_context.agents[
                    tool_id
                ].handle

                await ctx.info(f"Executing tool {tool_id} in {handle.agent_id}")
                await ctx.report_progress(0.1, 1)

                # Execute tool
                tool_result: RheaOutput = await (await handle.run_tool(rhea_params))

                await ctx.info(f"Tool {tool_id} finished in {handle.agent_id}")
                await ctx.report_progress(1, 1)

                result = MCPOutput.from_rhea(tool_result)

                # Add ProxyStore resource
                if result.files is not None:
                    for file in result.files:
                        ctx.request_context.lifespan_context.resource_manager.add_resource_to_context(
                            resource=create_proxystore_function_resource(file, ctx)
                        )

                # Notify the client that we have new output resources
                await ctx.request_context.session.send_resource_list_changed()

                # Log execution time
                metrics.tool_execution_runtime.observe(time.time() - start_time)

                # Log successful tool execution
                metrics.successful_tool_executions.inc()

                return result

            except Exception as e:
                # Log failed tool execution
                metrics.failed_tool_executions.inc()
                raise

        return wrapper

    # Create tool.call()
    if tool.name is None:
        tool.name = tool.user_provided_name
    safe_name = sanitize_tool_name(tool.name.lower())  # Normalize tool name
    fn = make_wrapper(tool.id, [name for name in params])
    fn.__name__ = safe_name
    fn.__doc__ = tool.description
    fn.__signature__ = sig  # type: ignore[attr-defined]

    fn.__annotations__ = {p.name: p.annotation for p in params}
    fn.__annotations__["return"] = MCPOutput

    return FastMCPTool.from_function(
        fn=fn, name=safe_name, title=tool.name, description=tool.description
    )
