import logging
import asyncio
import aiohttp
from typing import Dict, List, Any, Optional
from contextlib import AsyncExitStack
from urllib.parse import urljoin
from dotenv import load_dotenv
import webbrowser  # Add this import at the top

from mcp.client.sse import sse_client
from mcp.client.session import ClientSession
from mcp.types import CallToolResult

from .types import Tool, SecurityScheme, AuthStatus, Host

logger = logging.getLogger(__name__)
logger.setLevel(logging.WARNING)
load_dotenv()


class Client:
    """Client for managing MCP server connections and tools."""

    BASE_URL = "https://mcp-hive.ti.trilogy.com"

    def __init__(self, hive_ids: List[str]):
        self.hive_ids = hive_ids
        self.hive_urls: List[str] = []
        self.sessions: Dict[str, ClientSession] = {}
        self.exit_stack = AsyncExitStack()
        self.tools: List[Tool] = []
        self._tools_by_name: Dict[str, Tool] = {}

        # Session management
        self.session_id: Optional[str] = None
        self.auth_config: Optional[Dict] = None
        self.hives: Dict[str, Any] = {}

    async def init(self, session_id: Optional[str] = None) -> None:
        """Initialize the client with optional session_id."""
        self.session_id = session_id

        # Fetch hive information first
        await self._fetch_hives()

        if session_id:
            await self._fetch_session()
        else:
            await self._create_session()

        await self.connect()
        await self.authenticate()

    async def connect(self) -> None:
        """Connect to all configured hives."""
        if not self.hive_urls:
            # Get hive URLs from hive data
            self.hive_urls = [
                f"http://{hive.get("publishedUrl")}/sse" for hive in self.hives.values()
            ]

        for server_url in self.hive_urls:
            try:
                sse_transport = await self.exit_stack.enter_async_context(
                    sse_client(url=server_url)
                )
                stdio, write = sse_transport
                session = await self.exit_stack.enter_async_context(
                    ClientSession(stdio, write)
                )
                await session.initialize()
                self.sessions[server_url] = session
                logger.info(f"Connected to hive: {server_url}")
            except Exception as e:
                logger.error(f"Failed to connect to {server_url}: {e}")

        await self.setup_tools()

    async def setup_tools(self) -> List[Tool]:
        """Get tools from all connected hives."""
        self.tools = []
        for server_url, session in self.sessions.items():
            try:
                response = await session.list_tools()
                params_to_exclude = await self._get_params_to_exclude(
                    session, server_url
                )

                # Get security schemes from hive data
                hive_id = next(
                    (
                        id
                        for id, hive in self.hives.items()
                        if f"http://{hive.get('publishedUrl')}/sse" == server_url
                    ),
                    None,
                )
                hive_tools = self.hives[hive_id]["tools"] if hive_id else []

                filtered_tools = [
                    Tool(
                        name=tool.name,
                        description=tool.description,
                        input_schema=self._filter_tool_params(
                            tool.inputSchema, params_to_exclude
                        ),
                        hive=server_url,
                        security_schemes=self._get_security_schemes(
                            tool.name, hive_tools
                        ),
                        host=self._get_host_info(tool.name, hive_tools),
                    )
                    for tool in response.tools
                    if tool.name not in {"list_tool_configs", "list_params_to_exclude"}
                ]
                self.tools.extend(filtered_tools)
            except Exception as e:
                logger.error(f"Error getting tools from {server_url}: {e}")

        self._tools_by_name = {tool.name: tool for tool in self.tools}
        return self.tools

    async def get_status(self) -> AuthStatus:
        """Poll and return current session authentication status."""
        if not self.session_id:
            raise ValueError("Session not initialized. Call init() first.")

        async with aiohttp.ClientSession() as session:
            url = urljoin(self.BASE_URL, f"/api/sessions/{self.session_id}")
            async with session.get(url) as response:
                if response.status == 404:
                    raise ValueError("Session not found")
                data = await response.json()
                self.auth_config = data["session"]["auth"]["config"]
                return AuthStatus(data["session"]["auth"]["status"])

    async def authenticate(self) -> None:
        """Initiate authentication by opening browser and poll until complete."""
        if not self.session_id:
            raise ValueError("Session not initialized. Call init() first.")

        status = await self.get_status()
        if status not in [
            AuthStatus.AUTHENTICATED,
            AuthStatus.NO_AUTHENTICATION_REQUIRED,
        ]:
            # Open authentication URL in browser
            auth_url = urljoin(self.BASE_URL, f"/sessions/{self.session_id}")
            webbrowser.open_new_tab(auth_url)

            while True:
                status = await self.get_status()
                if status in [
                    AuthStatus.AUTHENTICATED,
                    AuthStatus.NO_AUTHENTICATION_REQUIRED,
                ]:
                    logger.info(f"Authentication complete with status: {status}")
                    break
                elif status == AuthStatus.UNAUTHENTICATED:
                    logger.info("Waiting for authentication...")
                    await asyncio.sleep(3)
                else:
                    logger.info(f"Partial authentication detected: {status}")
                    await asyncio.sleep(3)

    async def call_tool(
        self, tool_name: str, tool_args: Dict[str, Any]
    ) -> CallToolResult:
        """Call a tool with the given arguments."""
        tool = self._tools_by_name.get(tool_name)
        if not tool:
            raise ValueError(f"Tool {tool_name} not found")

        session = self.sessions[tool.hive]
        args = tool_args.copy()

        # Add host and auth information from session if available
        if self.session_id and self.auth_config:
            package_name = tool.host.name
            if package_name in self.auth_config["hosts"]:
                args["host__"] = self.auth_config["hosts"][package_name]
            else:
                args["host__"] = tool.host.default_url

            if package_name in self.auth_config["credentials"]:
                creds = self.auth_config["credentials"][package_name]
                for scheme_id, credential in creds.items():
                    if self._has_valid_credential(credential):
                        args["request_auth__"] = self._build_auth_config(credential)
                        break

        return await session.call_tool(tool_name, args)

    async def cleanup(self) -> None:
        """Clean up all sessions."""
        await self.exit_stack.aclose()

    # Helper methods
    def _get_security_schemes(
        self, tool_name: str, hive_tools: List[Dict]
    ) -> List[SecurityScheme]:
        """Get security schemes from hive tools data."""
        tool_data = next((t for t in hive_tools if t["name"] == tool_name), None)
        if not tool_data or not tool_data.get("securitySchemes"):
            return []
        return [
            SecurityScheme.from_dict(scheme) for scheme in tool_data["securitySchemes"]
        ]

    def _get_host_info(self, tool_name: str, hive_tools: List[Dict]) -> Host:
        """Get host information from hive tools data."""
        tool_data = next((t for t in hive_tools if t["name"] == tool_name), None)
        default_url = tool_data.get("defaultServerUrl") if tool_data else None
        host_name = tool_data.get("path", {}).get("package")
        return Host(name=host_name, default_url=default_url)

    async def _fetch_hives(self) -> None:
        """Fetch hive information for all hive_ids."""
        async with aiohttp.ClientSession() as session:
            for hive_id in self.hive_ids:
                url = urljoin(self.BASE_URL, f"/api/hives/{hive_id}")
                async with session.get(url) as response:
                    if response.status == 200:
                        hive_data = await response.json()
                        self.hives[hive_id] = hive_data

    async def _create_session(self) -> None:
        """Create new session via API."""
        async with aiohttp.ClientSession() as session:
            url = urljoin(self.BASE_URL, "/api/sessions")
            payload = {"hiveIds": self.hive_ids}
            async with session.post(url, json=payload) as response:
                if response.status == 200:
                    data = await response.json()
                    self.session_id = data["session"]["id"]
                    self.auth_config = data["session"]["auth"]["config"]

    async def _fetch_session(self) -> None:
        """Fetch current session configuration."""
        async with aiohttp.ClientSession() as session:
            url = urljoin(self.BASE_URL, f"/api/sessions/{self.session_id}")
            async with session.get(url) as response:
                if response.status == 200:
                    data = await response.json()
                    self.auth_config = data["session"]["auth"]["config"]

    def _has_valid_credential(self, credential: Dict) -> bool:
        """Check if credential has valid authentication data."""
        if credential["scheme"]["type"] == "oauth2":
            return bool(credential.get("accessToken"))
        elif credential["scheme"]["type"] == "apiKey":
            return bool(credential.get("apiKey"))
        elif credential["scheme"]["type"] == "http":
            return bool(credential.get("bearerToken"))
        return False

    def _build_auth_config(self, credential: Dict) -> Dict:
        """Build authentication configuration for tool call."""
        scheme = credential["scheme"]
        if scheme["type"] == "apiKey":
            return {
                "type": "apiKey",
                "key": scheme["name"],
                "value": credential["apiKey"],
                "in": scheme["in"],
            }
        elif scheme["type"] == "oauth2":
            return {
                "type": "oauth2",
                "key": "Authorization",
                "value": f"Bearer {credential['accessToken']}",
                "in": "header",
            }
        elif scheme["type"] == "http" and scheme["scheme"] == "bearer":
            return {
                "type": "bearer",
                "key": "Authorization",
                "value": f"Bearer {credential['bearerToken']}",
                "in": "header",
                "format": scheme.get("bearerFormat"),
            }
        return {}

    # Existing helper methods remain unchanged
    async def _get_params_to_exclude(
        self, session: ClientSession, server_url: str
    ) -> List[str]:
        try:
            result = await session.call_tool("list_params_to_exclude", {})
            return [param.text for param in result.content] if result.content else []
        except Exception as e:
            logger.warning(f"Could not get params to exclude from {server_url}: {e}")
            return []

    @staticmethod
    def _filter_tool_params(schema: dict, params_to_exclude: List[str]) -> dict:
        filtered_schema = schema.copy()
        if "properties" in filtered_schema:
            filtered_schema["properties"] = {
                k: v
                for k, v in filtered_schema["properties"].items()
                if k not in params_to_exclude
            }
        return filtered_schema
