import asyncio
import logging
import os
import shutil
import uuid
from pathlib import Path

from fastmcp import Context as MCPContext
from fastmcp import FastMCP

from relace_agent.agents import Agent
from relace_agent.config import AgentConfig, LiteralPrompt
from relace_agent.context import ChatHistory, Context, EventStream
from relace_agent.server.types import AgentEvent, ToolEvent
from relace_toolbox.local_client import ToolboxClient
from relace_toolbox.tools import ReportBackToolSchema

REQUIRED_ENV_VARS = ["RELACE_API_KEY"]

missing = [var for var in REQUIRED_ENV_VARS if not os.environ.get(var)]
if missing:
    raise RuntimeError(
        f"Missing required environment variables: {', '.join(missing)}. "
        f"Set these in your MCP server config or shell environment."
    )

# Check for ripgrep availability
if not shutil.which("rg"):
    raise RuntimeError(
        "ripgrep (rg) is not installed or not available in PATH. "
        "See installation instructions at https://github.com/BurntSushi/ripgrep?tab=readme-ov-file#installation"
    )

logging.basicConfig(
    level=logging.WARNING,  # Root logger at WARNING to suppress dependencies
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)
# Set first-party loggers to DEBUG/INFO
for logger_name in [
    "relace_mcp",
    "relace_agent",
    "relace_toolbox",
]:
    logging.getLogger(logger_name).setLevel(logging.DEBUG)

logger = logging.getLogger(__name__)

mcp = FastMCP("Relace MCP Server")

FAS_TOOLS = [
    "view_file",
    "view_directory",
    "grep_search",
    "bash",
    "check_task",
    "report_back",
]

FAS_SYSTEM_PROMPT = """
You are an AI agent whose job is to explore a code base with the provided tools and thoroughly understand the problem. You should use the tools provided to explore the codebase, read files, search for specific terms, and execute bash commands as needed. Once you have a good understanding of the problem, use the report_back tool to share your findings. Make sure to only use the report_back tool when you are confident that you have gathered enough information to make an informed decision.

Your objective is speed and efficiency so call multiple tools at once where applicable to reduce latency and reduce the number of turns. You are given a limited number of turns so aim to call 4-12 tools in parallel.

ASYNC EXECUTION: Some operations (grep, find, view_directory) can be slow on large repos. Use the wait_before_suspend parameter (default 2s) to prevent blocking—if exceeded, the task suspends and you continue working. Use check_task with wait_for_suspend to collect results later. Strategy: start slow operations with short timeouts, do fast operations (file reads), then check_task to gather results.

REASONING: Brief reasoning before tool calls helps you plan effectively and avoid wasted operations. However, you are judged on overall speed—keep reasoning concise and actionable. A few sentences explaining your next steps is ideal; avoid lengthy analysis.
"""

FAS_USER_PROMPT = """
Consider the following user query:

<user_query>
${{ inputs.query }}
</user_query>

You need to resolve the <user_query>.

Your job is purely to understand the codebase.

1. Explore and Understand the Codebase
Build a deep understanding of the relevant code. Locate and examine all relevant parts, understand control flow and edge cases, identify root causes or entry points, and review related unit tests.

2. Report Back Your Understanding
Use the report_back tool to report findings. File paths should be relative to project root directory. Only report relevant files within the repository.

Success Criteria:
- The issue in <user_query> is well understood
- You explain reasoning behind marking code as relevant
- Files comprehensively cover all key files needed to address the query

<async_task_execution>
SLOW OPERATIONS: grep, find, and view_directory can be slow on large codebases.

Use wait_before_suspend parameter (default: 2s) to avoid blocking. If exceeded, the task suspends and you can continue with other work.

Use check_task with wait_for_suspend to poll suspended operations for results in later turns.

STRATEGY: Start slow operations (grep, find, view_directory) with short wait_before_suspend, continue with fast operations (read specific files), then check_task to collect results.
</async_task_execution>

<reasoning_guidelines>
Brief reasoning before tool calls helps plan effectively and avoid wasted operations. However, overall speed matters—keep reasoning concise. A few sentences on your next steps is ideal; avoid lengthy analysis.
</reasoning_guidelines>

<use_parallel_tool_calls>
If you intend to call multiple tools and there are no dependencies between the tool calls, make all of the independent tool calls in parallel. Prioritize calling tools simultaneously whenever possible. Maximize use of parallel tool calls to increase speed and efficiency. However, if some tool calls depend on previous calls, do NOT call these tools in parallel. Never use placeholders or guess missing parameters.

Parallel tool calls schema:
<tool_call>
<function=example_function_1>
<parameter=param1>value1</parameter>
<parameter=wait_before_suspend>2</parameter>
</function>
<function=example_function_2>
<parameter=param1>value1</parameter>
</function>
</tool_call>
</use_parallel_tool_calls>
"""

ReportBackFiles = dict[str, list[tuple[int, int]]]
ReportBackEventData = dict[str, str | ReportBackFiles]


@mcp.tool
async def relace_search(
    query: str,
    workdir: str,
    ctx: MCPContext,
) -> ReportBackEventData:
    """Agentic codebase search powered by Relace. Use when finding files to modify or answering questions requiring deep search.

    Args:
        query: Specific, human-readable search query (avoid minimal keywords).
            Examples:
            - "Find where user authentication tokens are validated"
            - "How does the payment processing flow work?"
            - "Where are database migrations defined and executed?"
        workdir: Absolute path to search directory.
    """

    workdir_path = Path(workdir).expanduser()
    if not workdir_path.is_dir():
        raise NotADirectoryError(f"Workdir is not a directory: {workdir}")

    event_stream = EventStream()

    async def _run_agent() -> None:
        try:
            agent = Agent(
                config=AgentConfig(
                    name="relace-search",
                    model_name="relace-search",
                    prompts=LiteralPrompt(
                        system_prompt=FAS_SYSTEM_PROMPT,
                        user_prompt=FAS_USER_PROMPT,
                    ),
                    tools=FAS_TOOLS,
                ),
            )
            agent_context = Context(
                prompt_cache_key=f"fas-{uuid.uuid4()}",
                toolbox=ToolboxClient(),
                toolbox_workdir=workdir_path,
                inputs={
                    "query": query,
                },
                history=ChatHistory(max_messages=None),
                events=event_stream,
            )
            await agent.run_loop(agent_context)
        finally:
            await event_stream.close()

    agent_task = asyncio.create_task(_run_agent())

    # Process events concurrently with agent execution
    final_result: ReportBackEventData | None = None
    try:
        tool_call_count = 0
        async for event in event_stream:
            logger.info("Event: %s", event)

            # Emit tool calls as progress notifications to make them visible in Cursor
            if isinstance(event, ToolEvent):
                tool_call_count += 1

                # Format tool call information for progress notification
                tool_info = f"🔧 {event.name}"
                if event.path:
                    tool_info += f": {event.path}"
                elif event.query:
                    query_preview = (
                        event.query[:50] + "..."
                        if len(event.query) > 50
                        else event.query
                    )
                    tool_info += f": {query_preview}"
                elif event.command:
                    cmd_preview = (
                        event.command[:50] + "..."
                        if len(event.command) > 50
                        else event.command
                    )
                    tool_info += f": {cmd_preview}"

                # Report progress to make tool calls visible in Cursor's chat
                await ctx.report_progress(
                    progress=tool_call_count,
                    message=tool_info,
                )
                await ctx.info(message=tool_info)

                # Check if this is the final report_back event
                if event.name == ReportBackToolSchema.tool_name:
                    if event.files is None or event.explanation is None:
                        raise RuntimeError("Invalid report_back tool event")
                    final_result = {
                        "files": event.files,
                        "explanation": event.explanation,
                    }
            elif isinstance(event, AgentEvent):
                # Show agent messages as progress
                content_preview = (
                    event.content[:100] + "..."
                    if len(event.content) > 100
                    else event.content
                )
                await ctx.report_progress(
                    progress=0,
                    message=f"💭 {content_preview}",
                )
                await ctx.info(message=f"💭 Agent: {content_preview}")
    except Exception as e:
        await ctx.error(message=f"Error processing events: {e}")
        raise
    finally:
        if not agent_task.done():
            agent_task.cancel()
        try:
            await agent_task
        except asyncio.CancelledError:
            pass
        except Exception:
            logger.debug("Agent task exception", exc_info=True)

    if final_result is not None:
        return final_result

    raise RuntimeError("No report back event emitted")


def main() -> None:
    mcp.run()


if __name__ == "__main__":
    main()
