import os
import logging
import json
import httpx
from typing import List, Optional
from mcp.server.fastmcp import FastMCP
from fastmcp.server.dependencies import get_http_headers

logger = logging.getLogger(__name__)

# Configuration from environment variables
API_BASE_URL = os.getenv("FUTURMIND_API_URL", "https://projectflow-test.futurx.cn")

# MCP Server configuration
port = int(os.getenv("MCP_PORT", "8000"))
host = os.getenv("MCP_HOST", "localhost")


def get_api_key_from_headers() -> str:
    """Extract API key from headers, falling back to environment variable."""
    headers = get_http_headers() or {}

    api_key = (
        headers.get('X-Api-Key') or
        headers.get('x-api-key') or
        headers.get('Authorization', '').replace('Bearer ', '').strip() or
        headers.get('authorization', '').replace('Bearer ', '').strip()
    )

    if api_key:
        return api_key

    env_api_key = os.getenv("FUTURMIND_API_KEY", "").strip()
    if env_api_key:
        logger.info("Using FUTURMIND_API_KEY from environment variable")
        return env_api_key

    raise ValueError("X-Api-Key header is required or set FUTURMIND_API_KEY environment variable")


# Create FastMCP server instance
mcp = FastMCP(
    name="futurmind",
    instructions="""
    你是一个专家数字分身对话网关。
    
    ## 核心定位
    本工具提供访问企业内部专家数字分身的标准化接口。具体可以获得什么帮助取决于：
    - 使用的 API Key 对应哪位或哪些专家
    - 该专家的训练领域和专业能力
    
    ## 何时使用
    当用户的问题需要专业领域知识或专家经验时，应该调用此工具：
    - 用户明确需要咨询某个领域的专家
    - 问题涉及专业知识、行业经验、最佳实践
    - 需要深度分析、专业建议或决策支持
    - 希望进行有上下文的深入讨论
    
    ## 何时不使用
    - 通用常识性问题（你自己可以回答的）
    - 简单的事实查询或计算
    - 不需要专业知识的闲聊
    
    ## 可用工具
    - `chat`: 发送单条消息给专家
    - `chat_with_context`: 发送完整对话历史，适合多轮深入讨论
    
    ## 关键特性
    - 通过 X-Api-Key 访问对应的专家分身
    - 支持流式响应，实时获得专家回复
    - 支持多轮对话，保持上下文理解
    - 可获得后续推荐问题，引导深入交流
    """,
    host=host,
    port=port
)


async def call_futurmind_api(
        messages: list,
        api_key: str,
        conversation_id: Optional[str] = None
) -> str:
    """
    Call Futurmind Chat Stream API and return the response.

    Args:
        messages: List of conversation messages with 'role' and 'content'
        api_key: API key from request header (required)
        conversation_id: Optional conversation UUID

    Returns:
        Complete response text from the stream
    """
    url = f"{API_BASE_URL}/v1/zhiyin/chat/stream"

    # Prepare request body
    request_body = {
        "messages": messages
    }

    if conversation_id:
        request_body["conversation_id"] = conversation_id

    logger.info(f"Sending request to {url}")
    logger.info(f"Request body: {json.dumps(request_body, ensure_ascii=False)}")

    # Prepare headers
    headers = {
        "X-Api-Key": api_key,
        "Content-Type": "application/json",
        "Accept": "text/event-stream"
    }

    try:
        async with httpx.AsyncClient(timeout=300.0) as client:
            # Handle streaming response
            full_response = ""
            followup_questions = []
            final_conversation_id = None

            async with client.stream("POST", url, json=request_body, headers=headers) as response:
                response.raise_for_status()

                # Check if it's a streaming response
                content_type = response.headers.get("content-type", "")
                logger.info(f"Response content-type: {content_type}")
                logger.info(f"Response headers: {dict(response.headers)}")

                if "text/event-stream" in content_type:
                    # Process SSE stream
                    logger.info("Starting SSE stream processing...")
                    empty_chunk_count = 0
                    async for line in response.aiter_lines():
                        if line.startswith("data:"):
                            data_str = line[5:]

                            if data_str.strip() == "[DONE]":
                                break

                            try:
                                event_data = json.loads(data_str)

                                # Extract chunk
                                if "chunk" in event_data and event_data["chunk"] is not None and event_data[
                                    "chunk"] != "":
                                    full_response += event_data["chunk"]
                                    empty_chunk_count = 0
                                elif event_data.get("type") == "text" and event_data.get("chunk") == "":
                                    empty_chunk_count += 1
                                    if empty_chunk_count >= 2 and full_response:
                                        break

                                # Extract error
                                if "error" in event_data and event_data["error"]:
                                    logger.error(f"API returned error: {event_data['error']}")
                                    return f"Error from API: {event_data['error']}"

                                # Extract followup questions
                                if "followup_questions" in event_data and event_data["followup_questions"]:
                                    followup_questions = event_data["followup_questions"]

                                # Extract conversation ID
                                if "conversation_id" in event_data and event_data["conversation_id"]:
                                    final_conversation_id = event_data["conversation_id"]

                            except json.JSONDecodeError:
                                logger.warning(f"Failed to parse SSE data: {data_str}")
                                continue
                else:
                    # Non-streaming response (error case)
                    response_text = await response.aread()
                    try:
                        error_data = json.loads(response_text)
                        return f"Error: {error_data}"
                    except:
                        return f"Error: {response_text.decode()}"

            # Log the final complete response
            logger.info(f"Stream completed. Final response length: {len(full_response)} characters")
            logger.info(f"Final response: {full_response}")
            if final_conversation_id:
                logger.info(f"Conversation ID: {final_conversation_id}")
            if followup_questions:
                logger.info(f"Followup questions: {followup_questions}")

            # Format the complete response
            result = f"## Response\n\n{full_response}\n\n"

            if final_conversation_id:
                result += f"**Conversation ID**: `{final_conversation_id}`\n\n"

            if followup_questions:
                result += "### Suggested Follow-up Questions:\n"
                for i, question in enumerate(followup_questions, 1):
                    result += f"{i}. {question}\n"
                result += "\n"

            return result

    except httpx.HTTPError as e:
        logger.error(f"HTTP request failed: {e}")
        return f"Error: HTTP request failed - {str(e)}"
    except Exception as e:
        logger.error(f"API call failed: {e}")
        return f"Error: {str(e)}"


@mcp.tool()
async def chat(
        message: str,
        conversation_id: Optional[str] = None
) -> str:
    """
    Send a chat message to Futurmind API and get a streaming response.
    This is a simplified single-message interface.

    Args:
        message: The user's message content
        conversation_id: Optional conversation UUID for context

    Returns:
        The complete response from Futurmind API
    
    Note:
        Requires X-Api-Key header for authentication
    """
    # Get API key from headers
    try:
        api_key = get_api_key_from_headers()
    except ValueError as e:
        return f"Error: {str(e)}"

    # Log request parameters (excluding api_key for security)
    logger.info(
        f"chat tool called with parameters: message={message[:100]}{'...' if len(message) > 100 else ''}, conversation_id={conversation_id}")

    messages = [
        {"role": "user", "content": message}
    ]

    return await call_futurmind_api(
        messages=messages,
        api_key=api_key,
        conversation_id=conversation_id
    )


@mcp.tool()
async def chat_with_context(
        messages_json: str,
        conversation_id: Optional[str] = None
) -> str:
    """
    Send multiple messages with full context to Futurmind API.
    Supports multi-turn conversations with complete message history.

    Args:
        messages_json: JSON string of messages array, e.g. '[{"role":"user","content":"hello"},{"role":"assistant","content":"hi"}]'
        conversation_id: Optional conversation UUID

    Returns:
        The complete response from Futurmind API

    Example:
        messages_json: '[{"role":"user","content":"What is AI?"},{"role":"assistant","content":"AI is..."},{"role":"user","content":"Tell me more"}]'
    
    Note:
        Requires X-Api-Key header for authentication
    """
    # Get API key from headers
    try:
        api_key = get_api_key_from_headers()
    except ValueError as e:
        return f"Error: {str(e)}"

    # Log request parameters (excluding api_key for security)
    logger.info(
        f"chat_with_context tool called with parameters: messages_json_length={len(messages_json)}, conversation_id={conversation_id}")

    try:
        messages = json.loads(messages_json)

        if not isinstance(messages, list):
            return "Error: messages_json must be a JSON array"

        # Validate message format
        for msg in messages:
            if not isinstance(msg, dict) or "role" not in msg or "content" not in msg:
                return "Error: Each message must have 'role' and 'content' fields"

        # Log messages count and roles
        logger.info(f"Parsed {len(messages)} messages with roles: {[msg.get('role') for msg in messages]}")

    except json.JSONDecodeError as e:
        return f"Error: Invalid JSON format - {str(e)}"

    return await call_futurmind_api(
        messages=messages,
        api_key=api_key,
        conversation_id=conversation_id
    )


def serve():
    """Main function to run the MCP server"""
    transport_type = os.getenv("MCP_TRANSPORT", "stdio").lower()

    if transport_type == "sse":
        logger.info(f"Starting Futurmind MCP Server with SSE transport on {host}:{port}")
        mcp.run("sse")
    elif transport_type == "http" or transport_type == "streamable-http":
        logger.info(f"Starting Futurmind MCP Server with streamable-http transport on {host}:{port}")
        mcp.run("streamable-http")
    else:
        logger.info("Starting Futurmind MCP Server with stdio transport")
        mcp.run(transport="stdio")
