import asyncio
from mcp import StdioServerParameters, stdio_client
from mcp.types import JSONRPCMessage

from mcp_bridge.config import config
from mcp_bridge.mcp_clients.session import McpClientSession
from .AbstractClient import GenericMcpClient
from loguru import logger
import shutil
import os
from anyio.streams.memory import MemoryObjectReceiveStream


class FilteredStream:
    def __init__(self, stream: MemoryObjectReceiveStream[JSONRPCMessage | Exception]):
        self._stream = stream

    async def __aenter__(self):
        await self._stream.__aenter__()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        await self._stream.__aexit__(exc_type, exc_val, exc_tb)

    def __aiter__(self):
        return self

    async def __anext__(self):
        while True:
            message = await self._stream.__anext__()
            # Handle string messages that aren't valid JSON
            if isinstance(message, str):
                stripped = message.strip()
                if not (stripped.startswith("{") or stripped.startswith("[")):
                    logger.debug(f"Filtered non-JSON output: {message}")
                    continue

            # Handle exception objects that contain validation errors for non-JSON strings
            if isinstance(message, Exception):
                error_str = str(message)
                if "Invalid JSON" in error_str and "input_value" in error_str:
                    # Extract the original message from the error
                    import re
                    match = re.search(r"input_value='([^']*)'", error_str)
                    if match:
                        original_message = match.group(1)
                        logger.debug(f"Filtered exception for non-JSON: {original_message}")
                        continue

            return message

    async def receive(self):
        return await self.__anext__()


# Keywords to identify virtual environment variables
venv_keywords = ["CONDA", "VIRTUAL", "PYTHON"]

class StdioClient(GenericMcpClient):
    config: StdioServerParameters

    def __init__(self, name: str, config: StdioServerParameters) -> None:
        super().__init__(name=name)

        # logger.debug(f"initializing settings for {name}: {config.command} {" ".join(config.args)}")

        own_config = config.model_copy(deep=True)

        env = dict(os.environ.copy())

        env = {
            key: value for key, value in env.items()
            if not any(key.startswith(keyword) for keyword in venv_keywords)
        }

        if config.env is not None:
            env.update(config.env)

        own_config.env = env

        command = shutil.which(config.command)
        if command is None:
            logger.error(f"could not find command {config.command}")
            exit(1)

        own_config.command = command

        # this changes the default to ignore
        if "encoding_error_handler" not in config.model_fields_set:
            own_config.encoding_error_handler = "ignore"

        self.config = own_config

    async def _maintain_session(self):
        logger.debug(f"starting maintain session for {self.name}")
        async with stdio_client(self.config) as client:
            logger.debug(f"entered stdio_client context manager for {self.name}")
            assert client[0] is not None, f"missing read stream for {self.name}"
            assert client[1] is not None, f"missing write stream for {self.name}"

            # Wrap the read stream with FilteredStream to filter out non-JSON messages
            filtered_read_stream = FilteredStream(client[0])

            async with McpClientSession(filtered_read_stream, client[1]) as session:
                logger.debug(f"entered client session context manager for {self.name}")
                await session.initialize()
                logger.debug(f"finished initialise session for {self.name}")
                self.session = session

                try:
                    while True:
                        await asyncio.sleep(10)
                        if config.logging.log_server_pings:
                            logger.debug(f"pinging session for {self.name}")

                        await session.send_ping()

                except Exception as exc:
                    logger.error(f"ping failed for {self.name}: {exc}")
                    self.session = None

        logger.debug(f"exiting session for {self.name}")