#!/usr/bin/env python3
"""
MCP Server for Ambivo API Endpoints

This MCP server provides access to Ambivo's entity/natural_query endpoint
through standardized MCP tools. It handles authentication via JWT Bearer tokens with
enhanced security, configuration management, and error handling.
"""

import asyncio
import json
import logging
import os
import time
from typing import Any, Dict, List, Optional

import httpx
from mcp import types
from mcp.server import NotificationOptions, Server
from mcp.server.models import InitializationOptions

# Import from package modules
from .config import ServerConfig, load_config
from .security import InputValidator, RateLimiter, TokenValidator

# Load configuration
try:
    config_path = os.getenv("AMBIVO_CONFIG_FILE")
    config = load_config(config_path)
    logger = config.setup_logging()
except Exception as e:
    # Fallback logging if config fails
    logging.basicConfig(level=logging.INFO)
    logger = logging.getLogger("ambivo-mcp")
    logger.error(f"Failed to load configuration: {e}")
    # Use default config
    config = ServerConfig()

# Initialize  security components
rate_limiter = RateLimiter(
    max_requests=config.rate_limit_requests, window_seconds=config.rate_limit_window
)
input_validator = InputValidator(
    max_query_length=config.max_query_length, max_payload_size=config.max_payload_size
)
token_validator = TokenValidator(cache_ttl=config.token_cache_ttl)

# Server configuration
server = Server(config.server_name)


class AmbivoAPIClient:
    """Client for interacting with Ambivo API endpoints with enhanced error handling"""

    def __init__(self, config: ServerConfig, auth_token: Optional[str] = None):
        self.config = config
        self.base_url = config.base_url.rstrip("/")
        self.auth_token = auth_token
        self.client = httpx.AsyncClient(timeout=config.timeout)
        self.logger = logging.getLogger("ambivo-mcp.client")

    def set_auth_token(self, token: str):
        """Set the authentication token with validation"""
        try:
            if self.config.token_validation_enabled:
                token_validator.validate_token_format(token)
            self.auth_token = token
            self.logger.info("Authentication token set successfully")
        except ValueError as e:
            self.logger.error(f"Invalid token format: {e}")
            raise

    def _get_headers(self) -> Dict[str, str]:
        """Get headers for API requests"""
        headers = {"Content-Type": "application/json", "Accept": "application/json"}
        if self.auth_token:
            headers["Authorization"] = f"Bearer {self.auth_token}"
        return headers

    async def _make_request_with_retry(
        self, method: str, url: str, **kwargs
    ) -> httpx.Response:
        """Make HTTP request with retry logic"""
        last_exception = None

        for attempt in range(self.config.max_retries + 1):
            try:
                response = await self.client.request(method, url, **kwargs)
                return response
            except (httpx.TimeoutException, httpx.ConnectError, httpx.ReadError) as e:
                last_exception = e
                if attempt < self.config.max_retries:
                    wait_time = 2**attempt  # Exponential backoff
                    self.logger.warning(
                        f"Request attempt {attempt + 1} failed, retrying in {wait_time}s: {e}"
                    )
                    await asyncio.sleep(wait_time)
                else:
                    self.logger.error(
                        f"All {self.config.max_retries + 1} request attempts failed"
                    )
                    raise

        # This should never be reached, but just in case
        raise last_exception

    async def natural_query(
        self, query: str, response_format: str = "both"
    ) -> Dict[str, Any]:
        """
        Execute a natural language query against entity data with validation and error handling

        Args:
            query: Natural language query string
            response_format: Response format - "table", "natural", or "both"

        Returns:
            API response dictionary
        """
        # Validate inputs
        input_validator.validate_query(query)

        if response_format not in ["table", "natural", "both"]:
            raise ValueError(
                "Invalid response_format. Must be 'table', 'natural', or 'both'"
            )

        payload = {"query": query, "response_format": response_format}

        url = f"{self.base_url}/entity/natural_query"

        try:
            self.logger.info(f"Executing natural query: {query[:100]}...")
            start_time = time.time()

            response = await self._make_request_with_retry(
                "POST", url, json=payload, headers=self._get_headers()
            )

            elapsed_time = time.time() - start_time
            self.logger.info(f"Natural query completed in {elapsed_time:.2f}s")

            response.raise_for_status()
            result = response.json()

            self.logger.debug(f"API response: {json.dumps(result, indent=2)[:500]}...")
            return result

        except httpx.TimeoutException as e:
            self.logger.error(f"Natural query timeout: {e}")
            raise Exception(f"Request timeout after {self.config.timeout}s")
        except httpx.HTTPStatusError as e:
            self.logger.error(
                f"Natural query HTTP error: {e.response.status_code} - {e.response.text}"
            )
            raise
        except Exception as e:
            self.logger.error(f"Natural query unexpected error: {e}")
            raise

    async def close(self):
        """Close the HTTP client"""
        await self.client.aclose()


# Global client instance
api_client = AmbivoAPIClient(config, auth_token=config.auth_token)


@server.list_tools()
async def handle_list_tools() -> List[types.Tool]:
    """
    List available tools.
    """
    return [
        types.Tool(
            name="natural_query",
            description="Execute natural language queries against Ambivo entity data. "
            "This tool processes natural language queries and returns structured data "
            "about leads, contacts, opportunities, and other entities.",
            inputSchema={
                "type": "object",
                "properties": {
                    "query": {
                        "type": "string",
                        "description": "Natural language query describing what data you want to retrieve. "
                        "Examples: 'Show me leads created this week', 'Find contacts with gmail addresses', "
                        "'List opportunities worth more than $10,000'",
                    },
                    "response_format": {
                        "type": "string",
                        "enum": ["table", "natural", "both"],
                        "default": "both",
                        "description": "Format of the response: 'table' for structured data, "
                        "'natural' for natural language description, 'both' for both formats",
                    },
                },
                "required": ["query"],
            },
        ),
        types.Tool(
            name="set_auth_token",
            description="Set the authentication token for API requests. "
            "This must be called before using other tools to authenticate with the Ambivo API.",
            inputSchema={
                "type": "object",
                "properties": {
                    "token": {
                        "type": "string",
                        "description": "JWT Bearer token for authentication with Ambivo API",
                    }
                },
                "required": ["token"],
            },
        ),
    ]


@server.call_tool()
async def handle_call_tool(
    name: str, arguments: Dict[str, Any] | None
) -> List[types.TextContent]:
    """
    Handle tool calls with security and rate limiting.
    """
    if arguments is None:
        arguments = {}

    start_time = time.time()
    logger.info(f"Tool call started: {name}")

    try:
        # Rate limiting (except for auth token setting)
        if name != "set_auth_token" and api_client.auth_token:
            client_id = token_validator.get_client_id_from_token(api_client.auth_token)
            if not rate_limiter.is_allowed(client_id):
                stats = rate_limiter.get_client_stats(client_id)
                return [
                    types.TextContent(
                        type="text",
                        text=f"Rate limit exceeded. Requests: {stats['requests']}/{config.rate_limit_requests}. "
                        f"Reset in {stats['reset_time'] - time.time():.0f}s",
                    )
                ]

        if name == "set_auth_token":
            token = arguments.get("token")
            if not token:
                return [
                    types.TextContent(
                        type="text", text="Error: Authentication token is required"
                    )
                ]

            # Validate and set token
            api_client.set_auth_token(token)

            # Cache the token if validation is enabled
            if config.token_validation_enabled:
                token_validator.cache_token(token)

            return [
                types.TextContent(
                    type="text",
                    text="Authentication token set successfully. You can now use other tools to query the Ambivo API.",
                )
            ]

        elif name == "natural_query":
            if not api_client.auth_token:
                return [
                    types.TextContent(
                        type="text",
                        text="Error: Authentication required. Please use the 'set_auth_token' tool first.",
                    )
                ]

            query = arguments.get("query")
            if not query:
                return [
                    types.TextContent(
                        type="text", text="Error: Query parameter is required"
                    )
                ]

            response_format = arguments.get("response_format", "both")

            try:
                result = await api_client.natural_query(query, response_format)
                return [
                    types.TextContent(
                        type="text",
                        text=f"Natural Query Results:\n\n{json.dumps(result, indent=2)}",
                    )
                ]
            except httpx.HTTPStatusError as e:
                error_msg = f"HTTP {e.response.status_code}: {e.response.text}"
                return [types.TextContent(type="text", text=f"API Error: {error_msg}")]
            except Exception as e:
                return [
                    types.TextContent(
                        type="text", text=f"Error executing natural query: {str(e)}"
                    )
                ]

        else:
            return [types.TextContent(type="text", text=f"Unknown tool: {name}")]

    except ValueError as e:
        # Input validation errors
        logger.warning(f"Validation error in tool {name}: {e}")
        return [types.TextContent(type="text", text=f"Validation Error: {str(e)}")]

    except httpx.HTTPStatusError as e:
        # HTTP errors from API
        logger.error(f"HTTP error in tool {name}: {e.response.status_code}")
        error_msg = f"API Error (HTTP {e.response.status_code})"
        try:
            error_detail = e.response.json()
            if "error_code" in error_detail:
                error_msg += f": {error_detail['error_code']}"
        except:
            error_msg += f": {e.response.text[:200]}"

        return [types.TextContent(type="text", text=error_msg)]

    except Exception as e:
        # Unexpected errors
        logger.exception(f"Unexpected error in tool {name}")
        return [types.TextContent(type="text", text=f"Unexpected error: {str(e)}")]

    finally:
        elapsed_time = time.time() - start_time
        logger.info(f"Tool call completed: {name} in {elapsed_time:.2f}s")


async def main():
    """Main entry point for the MCP server with enhanced initialization"""
    logger.info(f"Starting {config.server_name} v{config.server_version}")
    logger.info(f"Configuration: Base URL: {config.base_url}")
    logger.info(
        f"Security: Rate limit: {config.rate_limit_requests}/{config.rate_limit_window}s"
    )

    # Import here to avoid issues with event loops
    import mcp.server.stdio

    try:
        async with mcp.server.stdio.stdio_server() as (read_stream, write_stream):
            await server.run(
                read_stream,
                write_stream,
                InitializationOptions(
                    server_name=config.server_name,
                    server_version=config.server_version,
                    capabilities=server.get_capabilities(
                        notification_options=NotificationOptions(),
                        experimental_capabilities={},
                    ),
                ),
            )
    except KeyboardInterrupt:
        logger.info("Server shutdown requested")
    except Exception as e:
        logger.error(f"Server error: {e}")
        raise
    finally:
        # Cleanup
        await api_client.close()
        logger.info("Server shutdown complete")


def run_server():
    """Synchronous wrapper for the async main function"""
    import asyncio

    asyncio.run(main())


if __name__ == "__main__":
    run_server()
