"""Lightweight MCP streamable-HTTP client using raw httpx + JSON-RPC 2.0.

Replaces the mcp SDK's ClientSession / streamable_http_client to eliminate
heavy transitive imports (starlette, uvicorn, pydantic-settings) that add
~300ms to every CLI invocation.  Only httpx is needed.

Supports both JSON and SSE response formats from the server.

Session caching: sessions are cached locally so subsequent CLI calls reuse
an existing server-side session and skip the initialize handshake (~400ms
saved per call on prod).  On 404 "Session not found", the client
transparently re-handshakes and caches the new session ID.
"""

import asyncio
import hashlib
import json
import logging
from contextlib import asynccontextmanager
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

import httpx

from mcpbundles.config import AuthError, CONFIG_DIR

logger = logging.getLogger(__name__)

_PROTOCOL_VERSION = "2025-03-26"
_CLIENT_INFO = {"name": "mcpbundles-cli", "version": "1.4.0"}


# ---------------------------------------------------------------------------
# Lightweight types (replace mcp.types)
# ---------------------------------------------------------------------------

@dataclass
class TextContent:
    type: str
    text: str


@dataclass
class ImageContent:
    type: str
    data: str
    mimeType: str


@dataclass
class CallToolResult:
    content: list[TextContent | ImageContent]
    isError: bool


@dataclass
class ToolInfo:
    name: str
    description: str | None = None
    inputSchema: dict[str, Any] | None = None


@dataclass
class ResourceInfo:
    uri: str
    name: str | None = None
    description: str | None = None
    mimeType: str | None = None


@dataclass
class ResourceContent:
    uri: str
    mimeType: str | None = None
    text: str | None = None


@dataclass
class PromptArgument:
    name: str
    description: str | None = None
    required: bool = False


@dataclass
class PromptInfo:
    name: str
    description: str | None = None
    arguments: list[PromptArgument] | None = None


@dataclass
class PromptMessage:
    role: str
    content: Any = None


# ---------------------------------------------------------------------------
# Result containers
# ---------------------------------------------------------------------------

@dataclass
class ListToolsResult:
    tools: list[ToolInfo] = field(default_factory=list)


@dataclass
class ListResourcesResult:
    resources: list[ResourceInfo] = field(default_factory=list)


@dataclass
class ListPromptsResult:
    prompts: list[PromptInfo] = field(default_factory=list)


@dataclass
class ReadResourceResult:
    contents: list[ResourceContent] = field(default_factory=list)


@dataclass
class GetPromptResult:
    description: str | None = None
    messages: list[PromptMessage] = field(default_factory=list)


# ---------------------------------------------------------------------------
# Exceptions
# ---------------------------------------------------------------------------

class ConnectionError(Exception):
    pass


class ToolCallError(Exception):
    def __init__(self, message: str, content: list[Any] | None = None):
        super().__init__(message)
        self.content = content


# ---------------------------------------------------------------------------
# JSON-RPC + SSE helpers
# ---------------------------------------------------------------------------

def _build_headers(auth: str) -> dict[str, str]:
    headers: dict[str, str] = {"Accept": "application/json, text/event-stream"}
    if auth.startswith("mb_"):
        headers["X-API-Key"] = auth
    else:
        headers["Authorization"] = f"Bearer {auth}"
    return headers


def _extract_result(resp: httpx.Response, req_id: int) -> Any:
    """Extract the JSON-RPC result from either a JSON or SSE response."""
    content_type = resp.headers.get("content-type", "")
    if "text/event-stream" in content_type:
        return _parse_sse_result(resp.text, req_id)
    data = resp.json()
    if "error" in data:
        msg = data["error"].get("message", "JSON-RPC error")
        raise ConnectionError(msg)
    return data.get("result")


def _parse_sse_result(text: str, req_id: int) -> Any:
    """Parse an SSE stream to find the JSON-RPC response matching req_id."""
    for block in text.split("\n\n"):
        data_parts: list[str] = []
        for line in block.split("\n"):
            if line.startswith("data:"):
                data_parts.append(line[5:].lstrip(" "))
        if not data_parts:
            continue
        try:
            payload = json.loads("".join(data_parts))
        except json.JSONDecodeError:
            continue
        if payload.get("id") == req_id:
            if "error" in payload:
                raise ConnectionError(payload["error"].get("message", "JSON-RPC error"))
            return payload.get("result")
    raise ConnectionError("No matching response in SSE stream")


def _parse_content_blocks(blocks: list[dict]) -> list[TextContent | ImageContent]:
    result: list[TextContent | ImageContent] = []
    for b in blocks:
        btype = b.get("type", "text")
        if btype == "text":
            result.append(TextContent(type="text", text=b.get("text", "")))
        elif btype == "image":
            result.append(ImageContent(
                type="image", data=b.get("data", ""), mimeType=b.get("mimeType", "image/png"),
            ))
    return result


def _parse_prompt_content(content: Any) -> Any:
    if isinstance(content, dict) and content.get("type") == "text":
        return TextContent(type="text", text=content.get("text", ""))
    return content


# ---------------------------------------------------------------------------
# Session cache
# ---------------------------------------------------------------------------

_SESSION_CACHE_FILE = CONFIG_DIR / "mcp_session_cache.json"


def _cache_key(url: str, auth: str) -> str:
    auth_hash = hashlib.sha256(auth.encode()).hexdigest()[:12]
    return f"{url}|{auth_hash}"


def _load_session_cache() -> dict[str, str]:
    try:
        return json.loads(_SESSION_CACHE_FILE.read_text())
    except (FileNotFoundError, json.JSONDecodeError, OSError):
        return {}


def _get_cached_session(url: str, auth: str) -> str | None:
    return _load_session_cache().get(_cache_key(url, auth))


def _cache_session(url: str, auth: str, session_id: str) -> None:
    cache = _load_session_cache()
    cache[_cache_key(url, auth)] = session_id
    try:
        CONFIG_DIR.mkdir(parents=True, exist_ok=True)
        _SESSION_CACHE_FILE.write_text(json.dumps(cache))
    except OSError:
        logger.debug("Failed to write session cache", exc_info=True)


def _evict_cached_session(url: str, auth: str) -> None:
    cache = _load_session_cache()
    if cache.pop(_cache_key(url, auth), None) is not None:
        try:
            _SESSION_CACHE_FILE.write_text(json.dumps(cache))
        except OSError:
            pass


# ---------------------------------------------------------------------------
# Session
# ---------------------------------------------------------------------------

class MCPSession:
    """Lightweight MCP session over streamable HTTP.

    Handles transparent re-handshake on session expiry (404).
    """

    def __init__(self, client: httpx.AsyncClient, url: str, session_id: str, auth: str) -> None:
        self._client = client
        self._url = url
        self._session_id = session_id
        self._auth = auth
        self._next_id = 2
        self._pending_tasks: list[asyncio.Task[Any]] = []

    async def _handshake(self) -> str:
        """Perform MCP initialize handshake, return new session ID."""
        init_resp = await self._client.post(self._url, json={
            "jsonrpc": "2.0",
            "method": "initialize",
            "params": {
                "protocolVersion": _PROTOCOL_VERSION,
                "capabilities": {},
                "clientInfo": _CLIENT_INFO,
            },
            "id": 1,
        })
        init_resp.raise_for_status()
        sid = init_resp.headers.get("mcp-session-id", "")
        _extract_result(init_resp, 1)

        task = asyncio.create_task(self._client.post(
            self._url,
            json={"jsonrpc": "2.0", "method": "notifications/initialized"},
            headers={"Mcp-Session-Id": sid},
        ))
        self._pending_tasks.append(task)
        return sid

    async def _request(self, method: str, params: dict[str, Any] | None = None) -> Any:
        req_id = self._next_id
        self._next_id += 1

        body: dict[str, Any] = {"jsonrpc": "2.0", "method": method, "id": req_id}
        if params is not None:
            body["params"] = params

        resp = await self._client.post(
            self._url, json=body,
            headers={"Mcp-Session-Id": self._session_id},
        )

        if resp.status_code == 404:
            _evict_cached_session(self._url, self._auth)
            self._session_id = await self._handshake()
            _cache_session(self._url, self._auth, self._session_id)
            logger.info("Session re-established after expiry: %s", self._url)
            resp = await self._client.post(
                self._url, json=body,
                headers={"Mcp-Session-Id": self._session_id},
            )

        resp.raise_for_status()
        return _extract_result(resp, req_id)

    async def cleanup(self) -> None:
        for task in self._pending_tasks:
            try:
                await task
            except Exception:
                pass

    async def list_tools(self) -> ListToolsResult:
        data = await self._request("tools/list")
        tools = [
            ToolInfo(
                name=t["name"],
                description=t.get("description"),
                inputSchema=t.get("inputSchema"),
            )
            for t in (data or {}).get("tools", [])
        ]
        return ListToolsResult(tools=tools)

    async def list_resources(self) -> ListResourcesResult:
        data = await self._request("resources/list")
        resources = [
            ResourceInfo(
                uri=str(r.get("uri", "")),
                name=r.get("name"),
                description=r.get("description"),
                mimeType=r.get("mimeType"),
            )
            for r in (data or {}).get("resources", [])
        ]
        return ListResourcesResult(resources=resources)

    async def list_prompts(self) -> ListPromptsResult:
        data = await self._request("prompts/list")
        prompts = [
            PromptInfo(
                name=p["name"],
                description=p.get("description"),
                arguments=[
                    PromptArgument(
                        name=a["name"],
                        description=a.get("description"),
                        required=a.get("required", False),
                    )
                    for a in p.get("arguments", [])
                ] if p.get("arguments") else None,
            )
            for p in (data or {}).get("prompts", [])
        ]
        return ListPromptsResult(prompts=prompts)

    async def read_resource(self, uri: str) -> ReadResourceResult:
        data = await self._request("resources/read", {"uri": uri})
        contents = [
            ResourceContent(
                uri=str(c.get("uri", uri)),
                mimeType=c.get("mimeType"),
                text=c.get("text"),
            )
            for c in (data or {}).get("contents", [])
        ]
        return ReadResourceResult(contents=contents)

    async def get_prompt(self, name: str, arguments: dict[str, Any] | None = None) -> GetPromptResult:
        params: dict[str, Any] = {"name": name}
        if arguments:
            params["arguments"] = arguments
        data = await self._request("prompts/get", params)
        data = data or {}
        messages = [
            PromptMessage(
                role=m.get("role", "assistant"),
                content=_parse_prompt_content(m.get("content")),
            )
            for m in data.get("messages", [])
        ]
        return GetPromptResult(description=data.get("description"), messages=messages)

    async def call_tool(self, name: str, arguments: dict[str, Any] | None = None) -> CallToolResult:
        data = await self._request("tools/call", {"name": name, "arguments": arguments or {}})
        data = data or {}
        content = _parse_content_blocks(data.get("content", []))
        return CallToolResult(content=content, isError=data.get("isError", False))


# ---------------------------------------------------------------------------
# Context manager
# ---------------------------------------------------------------------------

@asynccontextmanager
async def hub_session(
    auth: str,
    base_url: str,
    endpoint: str = "/hub/",
    session_id: str | None = None,
):
    """Open an MCP session to the Hub or a bundle endpoint.

    Uses cached session IDs when available to skip the initialize handshake.
    On 404 (expired/unknown session), MCPSession._request transparently
    re-handshakes and caches the new session ID.
    """
    url = f"{base_url}{endpoint}"
    headers = _build_headers(auth)

    try:
        async with httpx.AsyncClient(headers=headers, timeout=httpx.Timeout(60.0)) as client:
            session = MCPSession(client, url, "", auth)

            cached_sid = _get_cached_session(url, auth)
            if cached_sid:
                session._session_id = cached_sid
                logger.debug("Using cached session for %s", url)
            else:
                session._session_id = await session._handshake()
                _cache_session(url, auth, session._session_id)
                logger.info("MCP session initialized: %s", url)

            try:
                yield session
            finally:
                await session.cleanup()

    except (ToolCallError, AuthError, ConnectionError):
        raise
    except httpx.HTTPStatusError as exc:
        if exc.response.status_code in (401, 403):
            raise AuthError(
                f"Authentication failed ({exc.response.status_code}). Run 'mcpbundles login'.",
            ) from exc
        raise ConnectionError(f"Server error: {exc.response.status_code}") from exc
    except httpx.ConnectError as exc:
        raise ConnectionError(f"Cannot connect to {url}: {exc}") from exc
    except Exception as exc:
        raise ConnectionError(f"Connection failed: {exc}") from exc


# ---------------------------------------------------------------------------
# Public helpers
# ---------------------------------------------------------------------------

async def call_tool(session: MCPSession, name: str, arguments: dict[str, Any] | None = None) -> CallToolResult:
    result = await session.call_tool(name, arguments or {})
    if result.isError:
        text_parts = [c.text for c in result.content if isinstance(c, TextContent)]
        raise ToolCallError("\n".join(text_parts) if text_parts else "Tool call failed", result.content)
    return result


def extract_text(result: CallToolResult) -> str:
    parts: list[str] = []
    for block in result.content:
        if isinstance(block, TextContent):
            parts.append(block.text)
    return "\n".join(parts)


def extract_json(result: CallToolResult) -> Any:
    text = extract_text(result)
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        return text
