from typing import Optional

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

import click
import mcp.types as types
import asyncio
import mcp

from .vikingdb import VikingDBConnector


def serve(
    vikingdb_host: str, 
    vikingdb_region: str, 
    vikingdb_ak: str, 
    vikingdb_sk: str,
    collection_name: str,
    index_name: str
) -> Server:
    """
    Encapsulates the connection to a VikingDB server and some  methods to interact with it.
    :param vikingdb_host: The host to use for the VikingDB server.
    :param vikingdb_region: The region to use for the VikingDB server.
    :param vikingdb_ak: The Access Key to use for the VikingDB server.
    :param vikingdb_sk: The Secret Key to use for the VikingDB server.
    :param collection_name: The name of the collection to use.
    :param index_name: The name of the index to use.
    """
    


    server = Server("mcp-server-vikingdb")

    vikingdb = VikingDBConnector(
            vikingdb_host,
            vikingdb_region,
            vikingdb_ak,
            vikingdb_sk,
            collection_name,
            index_name
        )
    
    @server.list_tools()
    async def handle_list_tools() -> list[types.Tool]:
        """
        List available tools.
        Each tool specifies its arguments using JSON Schema validation.
        """
        return [
            types.Tool(
                name="vikingdb-upsert-information",
                description="upsert information to vikingdb for later use",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "information": {
                            "type": "string",
                        },
                    },
                    "required": ["information"],
                },
            ),
            types.Tool(
                name="vikingdb-search-information",
                description=(
                    "Look up information in VikingDB. Use this tool when you need to: \n"
                    " - Search for information by their query \n"
                    " - Access information for further analysis \n"
                    " - Get some personal information about the user"
                ),
                inputSchema={
                    "type": "object",
                    "properties": {
                        "query": {
                            "type": "string",
                            "description": "The query to search for in the vikingdb",
                        },
                    },
                    "required": ["query"],
                },
            ),
        ]

    @server.call_tool()
    async def handle_call_tool(
        name: str, arguments: dict | None
    ) -> list[types.TextContent | types.ImageContent | types.EmbeddedResource]:
        """
        Handle tool execution requests.
        Tools can modify server state and notify clients of changes.
        """
        if name not in ["vikingdb-upsert-information", "vikingdb-search-information"]:
            raise ValueError(f"Unknown tool: {name}")

        if name == "vikingdb-upsert-information":
            if not arguments or "information" not in arguments:
                raise ValueError("Missing required argument 'information'")
            information = arguments["information"]
            await vikingdb.upsert_information(information)
            return [types.TextContent(type="text", text=f"{information} has been added to vikingdb {collection_name}")]
        
        if name == "vikingdb-search-information":
            if not arguments or "query" not in arguments:
                raise ValueError("Missing required argument 'query'")
            query = arguments["query"]
            results = await vikingdb.search_information(query)
            return [types.TextContent(type="text", text=f"Search results for {query}:\n{results}")]
        
    return server
    
    
@click.command()
@click.option(
    "--vikingdb-host",
    envvar="VIKINGDB_HOST",
    required=True,
    help="VIKINGDB_HOST",
    default="api-vikingdb.volces.com",
)

@click.option(
    "--vikingdb-region",
    envvar="VIKINGDB_REGION",
    required=True,
    help="VIKINGDB_REGION",
)
@click.option(
    "--vikingdb-ak",
    envvar="VIKINGDB_AK",
    required=True,
    help="VIKINGDB_AK",
)
@click.option(
    "--vikingdb-sk",
    envvar="VIKINGDB_SK",
    required=True,
    help="VIKINGDB_SK",
)
@click.option(
    "--collection-name",
    envvar="COLLECTION_NAME",
    required=True,
    help="Collection name",
)
@click.option(
    "--index-name",
    envvar="INDEX_NAME",
    required=True,
    help="Index name",
)

def main(
    vikingdb_host: str, 
    vikingdb_region: str,
    vikingdb_ak: str,
    vikingdb_sk: str,
    collection_name: str,
    index_name: str,
):

    async def _run():
        async with mcp.server.stdio.stdio_server() as (read_stream, write_stream):
            server = serve(
                vikingdb_host, 
                vikingdb_region, 
                vikingdb_ak, 
                vikingdb_sk,
                collection_name,
                index_name,
            )
            await server.run(
                read_stream,
                write_stream,
                InitializationOptions(
                    server_name="mcp-server-vikingdb",
                    server_version="0.1.0",
                    capabilities=server.get_capabilities(
                        notification_options=NotificationOptions(),
                        experimental_capabilities={},
                    ),
                ),
            )

    asyncio.run(_run())

