"""FastMCP server for grok-cli-mcp."""

from __future__ import annotations

import json
import sys
from typing import Optional

from mcp.server.fastmcp.server import Context, FastMCP

from .types import GrokMessage
from .utils import _collect_assistant_text, _run_grok

# Server
server = FastMCP(
    name="grok-mcp",
    instructions=(
        "This MCP server exposes the Grok CLI via tools for general queries, chat-style prompts, "
        "and code-oriented tasks. Use grok_query for generic prompts, grok_chat for role-based "
        "messages, and grok_code when asking for code with optional language/context. "
        "You can optionally pass a Grok model name. Requires GROK_API_KEY."
    ),
    debug=False,
    log_level="INFO",
)


@server.tool(
    name="grok_query",
    title="Grok Query",
    description=(
        "Send a single prompt to Grok via CLI headless mode. Returns the assistant's text. "
        "Use raw_output=true to get raw CLI output and parsed messages."
    ),
)
async def grok_query(
    prompt: str,
    model: Optional[str] = None,
    raw_output: bool = False,
    timeout_s: float = 120.0,
    ctx: Optional[Context] = None,
) -> str | dict:
    """
    Send a single prompt to Grok.

    Args:
        prompt: The user prompt.
        model: Optional Grok model name (passed with -m if provided).
        raw_output: If true, returns {text, messages, raw, model}.
        timeout_s: Process timeout in seconds.
        ctx: FastMCP context.

    Returns:
        Assistant's text response, or dict with full details if raw_output=True.
    """
    result = await _run_grok(prompt, model=model, timeout_s=timeout_s, ctx=ctx)

    # Collate assistant text
    assistant_text = _collect_assistant_text(result.messages) if result.messages else (result.raw or "")

    if raw_output:
        return {
            "text": assistant_text,
            "messages": [m.model_dump() for m in result.messages],
            "raw": result.raw,
            "model": result.model,
        }

    return assistant_text


@server.tool(
    name="grok_chat",
    title="Grok Chat",
    description=(
        "Send a list of role/content messages to Grok by flattening into a single prompt. "
        "Useful for multi-turn context when the CLI only supports a single '-p' prompt."
    ),
)
async def grok_chat(
    messages: list[GrokMessage],
    model: Optional[str] = None,
    raw_output: bool = False,
    timeout_s: float = 120.0,
    ctx: Optional[Context] = None,
) -> str | dict:
    """
    Send multi-turn conversation to Grok.

    Args:
        messages: Array of {role, content}; content may be string or structured.
        model: Optional Grok model name.
        raw_output: If true, returns structured output.
        timeout_s: Process timeout in seconds.
        ctx: FastMCP context.

    Returns:
        Assistant's text response, or dict with full details if raw_output=True.
    """
    # Flatten messages to a single prompt with role tags
    # This preserves some conversational context even if CLI accepts only single prompt.
    prompt_lines: list[str] = []
    for m in messages:
        content_str: str
        if isinstance(m.content, str):
            content_str = m.content
        else:
            try:
                content_str = json.dumps(m.content, ensure_ascii=False)
            except Exception:
                content_str = str(m.content)
        prompt_lines.append(f"{m.role.capitalize()}: {content_str}")
    prompt = "\n".join(prompt_lines)

    result = await _run_grok(prompt, model=model, timeout_s=timeout_s, ctx=ctx)
    assistant_text = _collect_assistant_text(result.messages) if result.messages else (result.raw or "")

    if raw_output:
        return {
            "text": assistant_text,
            "messages": [m.model_dump() for m in result.messages],
            "raw": result.raw,
            "model": result.model,
        }

    return assistant_text


@server.tool(
    name="grok_code",
    title="Grok Code Task",
    description=(
        "Ask Grok for code or code-related guidance. You can provide a language hint and context "
        "(e.g., file snippets or requirements). Returns assistant text by default."
    ),
)
async def grok_code(
    task: str,
    language: Optional[str] = None,
    context: Optional[str] = None,
    model: Optional[str] = None,
    raw_output: bool = False,
    timeout_s: float = 180.0,
    ctx: Optional[Context] = None,
) -> str | dict:
    """
    Ask Grok for code generation or guidance.

    Args:
        task: Description of what code/help you need.
        language: Optional language hint (e.g., 'python', 'typescript').
        context: Optional context (repo constraints, file snippets, tests, etc.).
        model: Optional Grok model name.
        raw_output: If true, returns structured output.
        timeout_s: Process timeout in seconds.
        ctx: FastMCP context.

    Returns:
        Assistant's text response, or dict with full details if raw_output=True.
    """
    sys_instructions = [
        "You are an expert software engineer.",
        "Respond with clear, correct, directly usable code and concise explanations.",
        "Prefer minimal dependencies and explain tradeoffs when relevant.",
    ]
    if language:
        sys_instructions.append(f"Primary language: {language}")
    if context:
        sys_instructions.append("Context:\n" + context.strip())

    prompt = "\n\n".join(
        [
            "\n".join(sys_instructions),
            "Task:",
            task.strip(),
        ]
    )

    result = await _run_grok(prompt, model=model, timeout_s=timeout_s, ctx=ctx)
    assistant_text = _collect_assistant_text(result.messages) if result.messages else (result.raw or "")

    if raw_output:
        return {
            "text": assistant_text,
            "messages": [m.model_dump() for m in result.messages],
            "raw": result.raw,
            "model": result.model,
        }

    return assistant_text


def main() -> None:
    """Run the MCP server with stdio transport."""
    try:
        server.run("stdio")
    except KeyboardInterrupt:
        sys.exit(130)
