import asyncio
import os
import stat
import threading
from concurrent.futures import Future
from contextlib import asynccontextmanager
from logging import getLogger
from typing import Any, Coroutine

import anyio
import anyio.lowlevel
import mcp.types as types
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
from mcp import ClientSession
from mcp.shared.message import SessionMessage

logger = getLogger(__name__)


@asynccontextmanager
async def file_rpc_client(
    input_file: str = "/opt/astabench/tools/conn_in",
    output_file: str = "/opt/astabench/tools/conn_out",
    encoding: str = "utf-8",
    encoding_error_handler: str = "strict",
):
    """
    Client transport for file-based RPC using actual tail -f subprocess.
    Modeled after the [mcp python sdk](https://github.com/modelcontextprotocol/python-sdk/blob/main/src/mcp/server/stdio.py)
    implementation.
    """

    read_stream: MemoryObjectReceiveStream[SessionMessage | Exception]
    read_stream_writer: MemoryObjectSendStream[SessionMessage | Exception]

    write_stream: MemoryObjectSendStream[SessionMessage]
    write_stream_reader: MemoryObjectReceiveStream[SessionMessage]

    read_stream_writer, read_stream = anyio.create_memory_object_stream(0)
    write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
    tail_process = None

    # We will read differently depending if files are fifos
    files_are_fifos = False
    if os.path.exists(input_file) and os.path.exists(output_file):
        input_file_stat = os.stat(input_file)
        output_file_stat = os.stat(output_file)

        input_is_fifo = stat.S_ISFIFO(input_file_stat.st_mode)
        output_is_fifo = stat.S_ISFIFO(output_file_stat.st_mode)

        if input_is_fifo and output_is_fifo:
            files_are_fifos = True
        elif (not input_is_fifo) and (not output_is_fifo):
            files_are_fifos = False
        else:
            raise ValueError(
                "Both input and output files must be either FIFOs or regular files."
            )

    reader_ready = anyio.Event()

    async def file_reader():
        nonlocal tail_process
        try:
            async with read_stream_writer:
                # For non-fifos, we want to skip past any existing content when we start reading
                cmd = ["tail", "-f"] + ([] if files_are_fifos else ["-n", "0"])
                tail_process = await asyncio.create_subprocess_exec(
                    *cmd,
                    input_file,
                    stdout=asyncio.subprocess.PIPE,
                    stderr=asyncio.subprocess.PIPE,
                )

                reader_ready.set()

                while True:

                    line = await tail_process.stdout.readline()
                    if not line:  # EOF
                        break

                    line_str = line.decode(
                        encoding, errors=encoding_error_handler
                    ).strip()

                    if line_str:
                        try:
                            message = types.JSONRPCMessage.model_validate_json(line_str)
                            await read_stream_writer.send(SessionMessage(message))
                        except Exception as exc:
                            await read_stream_writer.send(exc)
        except anyio.ClosedResourceError:
            if tail_process and tail_process.returncode is None:
                tail_process.terminate()
            await anyio.lowlevel.checkpoint()

    async def file_writer():
        try:
            # Make sure reader is ready so there's no race condition
            await reader_ready.wait()
            async with write_stream_reader:
                async for sessionmessage in write_stream_reader:
                    jsondata = sessionmessage.message.model_dump_json(
                        by_alias=True, exclude_none=True
                    )
                    async with await anyio.open_file(
                        output_file,
                        "a",
                        encoding=encoding,
                        errors=encoding_error_handler,
                    ) as f:
                        await f.write(jsondata + "\n")
                        await f.flush()
        except anyio.ClosedResourceError:
            await anyio.lowlevel.checkpoint()

    async with anyio.create_task_group() as tg:
        tg.start_soon(file_reader)
        tg.start_soon(file_writer)
        try:
            yield read_stream, write_stream
        finally:
            # Clean up the tail process if it's still running
            if tail_process and tail_process.returncode is None:
                tail_process.terminate()
                try:
                    await asyncio.wait_for(tail_process.wait(), timeout=1.0)
                except asyncio.TimeoutError:
                    tail_process.kill()


@asynccontextmanager
async def mcp_session(
    input_file: str = "/opt/astabench/tools/conn_in",
    output_file: str = "/opt/astabench/tools/conn_out",
):
    """
    Helper for creating a session with the file-based RPC client.
    """
    async with file_rpc_client(input_file, output_file) as (read_stream, write_stream):
        async with ClientSession(read_stream, write_stream) as session:
            try:
                # 10 seconds to allow for very slow file polling; FIFOs should be <1 second in practice
                await asyncio.wait_for(session.initialize(), timeout=10)
            except asyncio.TimeoutError:
                raise asyncio.TimeoutError(
                    "Session initialization timed out when connecting to tool server."
                )
            yield session


def run_async_thread(coro: Coroutine) -> Any:
    """Run an async function synchronously regardless of whether we are
    currently in an async context (e.g. Jupyter) or a sync context.  This runs
    the coroutine in a separate thread with a new event loop and returns the
    result."""

    future: Future = Future()

    def run_in_thread():
        try:
            future.set_result(asyncio.run(coro))
        except Exception as e:
            future.set_exception(e)

    threading.Thread(target=run_in_thread, daemon=True).start()
    return future.result()
