"""
Official Linear MCP client session manager.

This module intentionally keeps the session lifecycle small:
1) connect
2) call_tool
3) close

The manager owns one background asyncio loop and one official MCP session.
`call_tool` retries once after transport/session failures and preserves semantic
tool errors returned by the official server.
"""

from __future__ import annotations

import asyncio
import concurrent.futures
import json
import logging
import os
import re
import shlex
import threading
import time
from datetime import timedelta
from typing import Any
from urllib.parse import urlsplit, urlunsplit

from mcp.client.session import ClientSession
from mcp.client.stdio import StdioServerParameters, stdio_client
from mcp.client.streamable_http import streamablehttp_client

logger = logging.getLogger(__name__)

DEFAULT_OFFICIAL_MCP_URL = "https://mcp.linear.app/mcp"
DEFAULT_TRANSPORT = "stdio"
DEFAULT_STDIO_COMMAND = "npx"
DEFAULT_STDIO_ARGS_PREFIX = ["-y", "mcp-remote"]
DEFAULT_AUTH_RETRY_COOLDOWN_SECONDS = 300.0
DEFAULT_CLOSE_JOIN_TIMEOUT_SECONDS = 5.0


class OfficialToolError(RuntimeError):
    """Raised when the official MCP call fails."""

    def __init__(self, code: str, message: str):
        super().__init__(message)
        self.code = code
        self.message = message


class OfficialMcpSessionManager:
    """Thread-safe synchronous wrapper around one async MCP client session."""

    def __init__(
        self,
        transport: str | None = None,
        command: str | None = None,
        args: list[str] | None = None,
        env: dict[str, str] | None = None,
        cwd: str | None = None,
        url: str | None = None,
        headers: dict[str, str] | None = None,
        timeout_seconds: float = 30.0,
        sse_read_timeout_seconds: float = 300.0,
        read_timeout_seconds: float = 30.0,
        auth_timeout_seconds: float = 180.0,
        auth_retry_cooldown_seconds: float | None = None,
    ):
        self._transport = (transport or os.getenv("LINEAR_OFFICIAL_MCP_TRANSPORT", DEFAULT_TRANSPORT)).lower()
        self._url = url or os.getenv("LINEAR_OFFICIAL_MCP_URL", DEFAULT_OFFICIAL_MCP_URL)
        self._headers = headers or self._parse_headers_from_env()
        self._command = command or os.getenv("LINEAR_OFFICIAL_MCP_COMMAND", DEFAULT_STDIO_COMMAND)
        self._args = args or self._parse_stdio_args_from_env(default_url=self._url)
        self._env = env or self._parse_stdio_env_from_env()
        self._cwd = cwd or os.getenv("LINEAR_OFFICIAL_MCP_CWD")
        self._timeout_seconds = timeout_seconds
        self._sse_read_timeout_seconds = sse_read_timeout_seconds
        self._read_timeout_seconds = read_timeout_seconds
        self._connect_timeout_seconds = auth_timeout_seconds
        self._auth_retry_cooldown_seconds = self._parse_auth_retry_cooldown_seconds(
            auth_retry_cooldown_seconds
        )

        if self._transport not in {"stdio", "http"}:
            raise ValueError("LINEAR_OFFICIAL_MCP_TRANSPORT must be one of: stdio, http")

        self._loop: asyncio.AbstractEventLoop | None = None
        self._thread: threading.Thread | None = None
        self._lock = threading.RLock()

        self._transport_cm: Any = None
        self._session_cm: Any = None
        self._session: ClientSession | None = None

        self._failure_count = 0
        self._last_error: str | None = None
        self._last_failure_at: float | None = None
        self._last_connected_at: float | None = None
        self._last_connect_attempt_at: float | None = None
        self._auth_retry_not_before: float | None = None
        self._last_auth_failure_reason: str | None = None

    @staticmethod
    def _parse_headers_from_env() -> dict[str, str] | None:
        raw = os.getenv("LINEAR_OFFICIAL_MCP_HEADERS")
        if not raw:
            return None
        try:
            parsed = json.loads(raw)
            if isinstance(parsed, dict):
                return {str(k): str(v) for k, v in parsed.items()}
        except Exception:
            logger.warning("Ignoring invalid LINEAR_OFFICIAL_MCP_HEADERS value")
        return None

    @staticmethod
    def _parse_stdio_env_from_env() -> dict[str, str] | None:
        raw = os.getenv("LINEAR_OFFICIAL_MCP_ENV")
        if not raw:
            return None
        try:
            parsed = json.loads(raw)
            if isinstance(parsed, dict):
                return {str(k): str(v) for k, v in parsed.items()}
        except Exception:
            logger.warning("Ignoring invalid LINEAR_OFFICIAL_MCP_ENV value")
        return None

    @staticmethod
    def _parse_stdio_args_from_env(default_url: str) -> list[str]:
        raw = os.getenv("LINEAR_OFFICIAL_MCP_ARGS")
        if not raw:
            return [*DEFAULT_STDIO_ARGS_PREFIX, default_url]

        # Prefer JSON array for exact argument boundaries.
        try:
            parsed = json.loads(raw)
            if isinstance(parsed, list):
                return [str(item) for item in parsed]
        except Exception:
            pass

        # Fallback to shell-like parsing.
        try:
            return shlex.split(raw)
        except ValueError:
            logger.warning("Ignoring invalid LINEAR_OFFICIAL_MCP_ARGS value; using default args")
            return [*DEFAULT_STDIO_ARGS_PREFIX, default_url]

    @staticmethod
    def _parse_auth_retry_cooldown_seconds(override: float | None) -> float:
        if override is not None:
            return max(0.0, float(override))

        raw = os.getenv("LINEAR_OFFICIAL_AUTH_RETRY_COOLDOWN_SECONDS")
        if raw is None or raw == "":
            return DEFAULT_AUTH_RETRY_COOLDOWN_SECONDS
        try:
            return max(0.0, float(raw))
        except ValueError:
            logger.warning(
                "Ignoring invalid LINEAR_OFFICIAL_AUTH_RETRY_COOLDOWN_SECONDS value %r; using default %.0fs",
                raw,
                DEFAULT_AUTH_RETRY_COOLDOWN_SECONDS,
            )
            return DEFAULT_AUTH_RETRY_COOLDOWN_SECONDS

    @staticmethod
    def _redact_text(value: str | None) -> str | None:
        if value is None:
            return None
        redacted = value
        redacted = re.sub(r"(?i)(authorization\s*[:=]\s*)([^\s,;]+)", r"\1***", redacted)
        redacted = re.sub(r"(?i)(bearer\s+)[^\s,;]+", r"\1***", redacted)
        redacted = re.sub(
            r"(?i)((?:api[_-]?key|token|secret|password)\s*[:=]\s*)([^\s,;]+)",
            r"\1***",
            redacted,
        )
        return redacted

    @classmethod
    def _redact_args(cls, args: list[str]) -> list[str]:
        redacted: list[str] = []
        redact_next_n = 0
        for arg in args:
            lower = arg.lower()
            if redact_next_n > 0:
                redacted.append("***")
                redact_next_n -= 1
                continue
            if lower == "--header":
                redacted.append(arg)
                # Some shells split "Authorization: Bearer <token>" into multiple args.
                redact_next_n = 2
                continue
            if lower in {"--auth-token", "--token", "--api-key"}:
                redacted.append(arg)
                redact_next_n = 1
                continue
            if any(
                marker in lower
                for marker in (
                    "authorization",
                    "bearer ",
                    "token=",
                    "api_key=",
                    "apikey=",
                    "secret=",
                    "password=",
                )
            ):
                redacted.append("***")
            else:
                redacted.append(arg)
        return redacted

    def _ensure_loop(self) -> None:
        if self._loop and self._thread and self._thread.is_alive():
            return

        loop = asyncio.new_event_loop()

        def _run_loop() -> None:
            asyncio.set_event_loop(loop)
            loop.run_forever()

        thread = threading.Thread(target=_run_loop, daemon=True, name="linear-official-mcp")
        thread.start()

        self._loop = loop
        self._thread = thread

    def _submit(self, coro: Any, timeout: float | None = None) -> Any:
        if not self._loop:
            raise RuntimeError("event loop not initialized")
        future = asyncio.run_coroutine_threadsafe(coro, self._loop)
        try:
            return future.result(timeout=timeout or (self._read_timeout_seconds + 10))
        except concurrent.futures.TimeoutError:
            future.cancel()
            try:
                future.result(timeout=1)
            except Exception:
                pass
            raise

    async def _connect_async(self) -> None:
        if self._session is not None:
            return

        if self._transport == "stdio":
            params = StdioServerParameters(
                command=self._command,
                args=self._args,
                env=self._env,
                cwd=self._cwd,
            )
            self._transport_cm = stdio_client(params)
        else:
            self._transport_cm = streamablehttp_client(
                self._url,
                headers=self._headers,
                timeout=self._timeout_seconds,
                sse_read_timeout=self._sse_read_timeout_seconds,
                terminate_on_close=False,
            )

        transport_streams = await self._transport_cm.__aenter__()
        if len(transport_streams) == 3:
            read_stream, write_stream, _ = transport_streams
        elif len(transport_streams) == 2:
            read_stream, write_stream = transport_streams
        else:
            raise RuntimeError("official MCP transport returned unexpected stream tuple")

        self._session_cm = ClientSession(
            read_stream,
            write_stream,
            read_timeout_seconds=timedelta(seconds=self._read_timeout_seconds),
        )
        self._session = await self._session_cm.__aenter__()
        await self._session.initialize()
        self._last_connected_at = time.time()

    async def _disconnect_async(self) -> None:
        if self._session_cm is not None:
            try:
                await self._session_cm.__aexit__(None, None, None)
            except Exception as exc:
                self._log_cleanup_exception("Official MCP session cleanup failed", exc)
        if self._transport_cm is not None:
            try:
                await self._transport_cm.__aexit__(None, None, None)
            except Exception as exc:
                self._log_cleanup_exception("Official MCP transport cleanup failed", exc)

        self._session = None
        self._session_cm = None
        self._transport_cm = None

    def _safe_disconnect(self, message: str) -> None:
        try:
            self._submit(self._disconnect_async())
        except Exception as exc:
            self._log_cleanup_exception(message, exc)

    @staticmethod
    def _log_cleanup_exception(prefix: str, exc: Exception) -> None:
        message = str(exc)
        if "Attempted to exit cancel scope in a different task" in message:
            logger.debug("%s: %s", prefix, exc)
            return
        logger.warning("%s: %s", prefix, exc)

    def _connect_locked(self) -> None:
        self._ensure_loop()
        self._last_connect_attempt_at = time.time()
        self._submit(self._connect_async(), timeout=self._connect_timeout_seconds + 10)

    def connect(self) -> None:
        """Eagerly connect to official MCP."""
        with self._lock:
            try:
                cooldown_err = self._auth_retry_cooldown_error()
                if cooldown_err is not None:
                    raise cooldown_err
                self._connect_locked()
                self._record_success()
            except OfficialToolError:
                raise
            except Exception as exc:
                self._record_failure(exc)
                self._safe_disconnect("Official MCP disconnect after connect failure failed")
                if self._is_likely_auth_connect_error(exc):
                    self._mark_auth_retry_cooldown(exc)
                    raise OfficialToolError(
                        "official_auth_required",
                        f"official MCP authentication required: {exc}",
                    ) from exc
                raise OfficialToolError(
                    "official_unavailable",
                    f"official MCP connect failed: {exc}",
                ) from exc

    def _normalize_result(self, result: Any) -> Any:
        if getattr(result, "isError", False):
            text = self._extract_text(result)
            raise OfficialToolError(
                "official_tool_error",
                text or "official MCP returned an error",
            )

        structured = getattr(result, "structuredContent", None)
        if structured is not None:
            return structured

        text = self._extract_text(result)
        if text:
            try:
                return json.loads(text)
            except Exception:
                return {"text": text}

        if hasattr(result, "model_dump"):
            return result.model_dump()
        return result

    @staticmethod
    def _extract_text(result: Any) -> str:
        content = getattr(result, "content", None) or []
        texts: list[str] = []
        for block in content:
            if getattr(block, "type", None) == "text":
                text = getattr(block, "text", "")
                if text:
                    texts.append(text)
        return "\n".join(texts).strip()

    def _record_failure(self, exc: Exception) -> None:
        self._failure_count += 1
        self._last_failure_at = time.time()
        self._last_error = f"{exc.__class__.__name__}: {exc}"
        safe_error = self._redact_text(self._last_error) or exc.__class__.__name__
        logger.warning("Official MCP call failed (%s): %s", exc.__class__.__name__, safe_error)

    def _record_success(self) -> None:
        self._failure_count = 0
        self._last_error = None
        self._auth_retry_not_before = None
        self._last_auth_failure_reason = None

    @staticmethod
    def _is_likely_auth_connect_error(exc: Exception) -> bool:
        if isinstance(exc, (TimeoutError, concurrent.futures.TimeoutError)):
            return True
        text = f"{exc.__class__.__name__}: {exc}".lower()
        auth_signals = (
            "oauth",
            "unauthorized",
            "forbidden",
            "401",
            "403",
            "consent",
            "login",
            "browser",
            "authorize",
            "expired",
        )
        return any(signal in text for signal in auth_signals)

    def _mark_auth_retry_cooldown(self, exc: Exception) -> None:
        if self._auth_retry_cooldown_seconds <= 0:
            return
        self._auth_retry_not_before = time.time() + self._auth_retry_cooldown_seconds
        self._last_auth_failure_reason = f"{exc.__class__.__name__}: {exc}"
        safe_reason = self._redact_text(self._last_auth_failure_reason) or exc.__class__.__name__
        logger.warning(
            "Suppressing official MCP reconnect attempts for %.0fs after auth-related failure: %s",
            self._auth_retry_cooldown_seconds,
            safe_reason,
        )

    def _auth_retry_cooldown_error(self) -> OfficialToolError | None:
        if self._auth_retry_not_before is None:
            return None
        now = time.time()
        if now >= self._auth_retry_not_before:
            self._auth_retry_not_before = None
            self._last_auth_failure_reason = None
            return None
        remaining = int(self._auth_retry_not_before - now)
        reason = self._redact_text(self._last_auth_failure_reason) or "recent authentication failure"
        return OfficialToolError(
            "official_auth_cooldown",
            f"official MCP auth retry is temporarily suppressed for {remaining}s ({reason})",
        )

    def _auth_retry_remaining_seconds(self) -> int:
        if self._auth_retry_not_before is None:
            return 0
        remaining = int(self._auth_retry_not_before - time.time())
        if remaining <= 0:
            self._auth_retry_not_before = None
            self._last_auth_failure_reason = None
            return 0
        return remaining

    def call_tool(self, name: str, arguments: dict[str, Any] | None = None) -> Any:
        """Invoke an official tool with one reconnect retry on transport failures."""
        args = arguments or {}
        with self._lock:
            for attempt in range(2):
                started_without_session = self._session is None
                try:
                    if self._session is None:
                        cooldown_err = self._auth_retry_cooldown_error()
                        if cooldown_err is not None:
                            raise cooldown_err
                        self._connect_locked()

                    if self._session is None:
                        raise RuntimeError("official MCP session unavailable")

                    result = self._submit(self._session.call_tool(name, arguments=args))
                    normalized = self._normalize_result(result)
                    self._record_success()
                    return normalized
                except OfficialToolError as exc:
                    if exc.code in {"official_tool_error", "official_auth_cooldown", "official_auth_required"}:
                        # Official tool reported a semantic error; do not mask or retry.
                        raise

                    self._record_failure(exc)
                    self._safe_disconnect("Official MCP disconnect after tool error failed")
                    if attempt == 1:
                        raise OfficialToolError(
                            "official_unavailable",
                            f"official MCP call failed for tool '{name}': {exc.message}",
                        ) from exc
                except Exception as exc:
                    self._record_failure(exc)
                    self._safe_disconnect("Official MCP disconnect after call failure failed")
                    if started_without_session and self._is_likely_auth_connect_error(exc):
                        self._mark_auth_retry_cooldown(exc)
                        raise OfficialToolError(
                            "official_auth_required",
                            f"official MCP authentication required: {exc}",
                        ) from exc
                    if attempt == 1:
                        raise OfficialToolError(
                            "official_unavailable",
                            f"official MCP call failed for tool '{name}': {exc}",
                        ) from exc

        raise OfficialToolError("official_unavailable", "official MCP unavailable")

    def is_connected(self) -> bool:
        with self._lock:
            return self._session is not None

    def get_health(self) -> dict[str, Any]:
        with self._lock:
            health_url = self._url
            try:
                parsed = urlsplit(self._url)
                host = parsed.hostname or ""
                if parsed.port:
                    host = f"{host}:{parsed.port}"
                health_url = urlunsplit((parsed.scheme, host, parsed.path, "", ""))
            except Exception:
                # Keep original URL when parsing fails; error fields stay redacted.
                pass

            health: dict[str, Any] = {
                "transport": self._transport,
                "url": health_url,
                "connected": self._session is not None,
                "failureCount": self._failure_count,
                "lastError": self._redact_text(self._last_error),
                "lastFailureAt": self._last_failure_at,
                "lastConnectedAt": self._last_connected_at,
                "lastConnectAttemptAt": self._last_connect_attempt_at,
                "authRetryCooldownSeconds": self._auth_retry_cooldown_seconds,
                "authRetryRemainingSeconds": self._auth_retry_remaining_seconds(),
                "lastAuthFailureReason": self._redact_text(self._last_auth_failure_reason),
            }
            if self._transport == "stdio":
                health["command"] = self._command
                health["args"] = self._redact_args(self._args)
            else:
                health["hasHeaders"] = self._headers is not None
            return health

    def close(self) -> None:
        with self._lock:
            loop = self._loop
            thread = self._thread

            if loop:
                self._safe_disconnect("Official MCP disconnect during close failed")
                try:
                    loop.call_soon_threadsafe(loop.stop)
                except Exception as exc:
                    self._log_cleanup_exception("Official MCP loop stop failed", exc)

            if thread and thread.is_alive():
                thread.join(timeout=DEFAULT_CLOSE_JOIN_TIMEOUT_SECONDS)

            if loop and (thread is None or not thread.is_alive()):
                try:
                    loop.close()
                except Exception as exc:
                    self._log_cleanup_exception("Official MCP loop close failed", exc)
                self._loop = None
                self._thread = None
            elif thread and thread.is_alive():
                logger.warning("Official MCP loop thread did not stop within timeout")
                # Keep references so callers can try `close` again.
                self._loop = loop
                self._thread = thread

            self._session = None
            self._session_cm = None
            self._transport_cm = None
