"""Named connection pool with persistent session IDs and encrypted auth.

Each named connection targets a different MCP endpoint (Hub, bundle, different
workspace/server). Auth tokens (API keys or OAuth access tokens) are encrypted
at rest using the same Fernet key as global OAuth tokens. OAuth refresh data
is stored alongside so connections can auto-renew without user interaction.
"""

import json
import logging
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path

from cryptography.fernet import Fernet, InvalidToken

from mcpbundles.config import CONFIG_DIR, get_or_create_key

logger = logging.getLogger(__name__)

SESSIONS_DIR = CONFIG_DIR / "sessions"


def _encrypt(plaintext: str) -> str:
    cipher = Fernet(get_or_create_key())
    return cipher.encrypt(plaintext.encode()).decode()


def _decrypt(ciphertext: str) -> str:
    cipher = Fernet(get_or_create_key())
    return cipher.decrypt(ciphertext.encode()).decode()


@dataclass
class ConnectionInfo:
    name: str
    url: str
    endpoint: str
    workspace_id: str | None
    session_id: str | None
    is_default: bool
    created_at: str
    auth_token_encrypted: str | None = None
    auth_refresh_encrypted: str | None = None

    @property
    def display_target(self) -> str:
        from urllib.parse import urlparse
        parsed = urlparse(self.url)
        return f"{parsed.hostname or self.url}{self.endpoint}"

    @property
    def has_auth(self) -> bool:
        return self.auth_token_encrypted is not None

    @property
    def auth_type(self) -> str:
        if self.auth_refresh_encrypted:
            return "oauth"
        if self.auth_token_encrypted:
            return "api-key"
        return "none"

    def get_auth_token(self) -> str | None:
        if not self.auth_token_encrypted:
            return None
        try:
            return _decrypt(self.auth_token_encrypted)
        except (InvalidToken, Exception):
            logger.warning("Failed to decrypt auth token for connection '%s'", self.name)
            return None

    def set_auth_token(self, token: str) -> None:
        self.auth_token_encrypted = _encrypt(token)

    def get_refresh_data(self) -> dict | None:
        if not self.auth_refresh_encrypted:
            return None
        try:
            return json.loads(_decrypt(self.auth_refresh_encrypted))
        except (InvalidToken, json.JSONDecodeError, Exception):
            logger.warning("Failed to decrypt refresh data for connection '%s'", self.name)
            return None

    def set_refresh_data(self, data: dict) -> None:
        self.auth_refresh_encrypted = _encrypt(json.dumps(data))

    def clear_auth(self) -> None:
        self.auth_token_encrypted = None
        self.auth_refresh_encrypted = None


class ConnectionManager:
    """Manage multiple named MCP connections with encrypted auth storage."""

    def __init__(self) -> None:
        SESSIONS_DIR.mkdir(parents=True, exist_ok=True)

    def _path(self, name: str) -> Path:
        return SESSIONS_DIR / f"{name}.json"

    def _load_conn(self, data: dict) -> ConnectionInfo:
        known_fields = {
            "name", "url", "endpoint", "workspace_id", "session_id",
            "is_default", "created_at", "auth_token_encrypted", "auth_refresh_encrypted",
        }
        filtered = {k: v for k, v in data.items() if k in known_fields}
        return ConnectionInfo(**filtered)

    def list_connections(self) -> list[ConnectionInfo]:
        conns: list[ConnectionInfo] = []
        if not SESSIONS_DIR.exists():
            return conns
        for f in sorted(SESSIONS_DIR.glob("*.json")):
            try:
                data = json.loads(f.read_text())
                conns.append(self._load_conn(data))
            except Exception:
                logger.warning("Corrupt session file: %s", f)
        return conns

    def get(self, name: str) -> ConnectionInfo | None:
        path = self._path(name)
        if not path.exists():
            return None
        try:
            return self._load_conn(json.loads(path.read_text()))
        except Exception:
            return None

    def get_default(self) -> ConnectionInfo | None:
        for conn in self.list_connections():
            if conn.is_default:
                return conn
        conns = self.list_connections()
        return conns[0] if conns else None

    def create(
        self,
        name: str,
        url: str,
        endpoint: str = "/hub/",
        workspace_id: str | None = None,
        is_default: bool = False,
        auth_token: str | None = None,
        refresh_data: dict | None = None,
    ) -> ConnectionInfo:
        if is_default:
            self._clear_default()
        conn = ConnectionInfo(
            name=name,
            url=url,
            endpoint=endpoint,
            workspace_id=workspace_id,
            session_id=None,
            is_default=is_default,
            created_at=datetime.now(timezone.utc).isoformat(),
        )
        if auth_token:
            conn.set_auth_token(auth_token)
        if refresh_data:
            conn.set_refresh_data(refresh_data)
        self._save(conn)
        return conn

    def update_auth(self, name: str, auth_token: str, refresh_data: dict | None = None) -> bool:
        conn = self.get(name)
        if not conn:
            return False
        conn.set_auth_token(auth_token)
        if refresh_data is not None:
            conn.set_refresh_data(refresh_data)
        self._save(conn)
        return True

    def remove(self, name: str) -> bool:
        path = self._path(name)
        if path.exists():
            path.unlink()
            return True
        return False

    def remove_all(self) -> int:
        count = 0
        for f in SESSIONS_DIR.glob("*.json"):
            f.unlink()
            count += 1
        return count

    def save_session_id(self, name: str, session_id: str) -> None:
        conn = self.get(name)
        if conn:
            conn.session_id = session_id
            self._save(conn)

    def invalidate_session(self, name: str) -> None:
        conn = self.get(name)
        if conn:
            conn.session_id = None
            self._save(conn)

    def set_default(self, name: str) -> bool:
        conn = self.get(name)
        if not conn:
            return False
        self._clear_default()
        conn.is_default = True
        self._save(conn)
        return True

    def _clear_default(self) -> None:
        for conn in self.list_connections():
            if conn.is_default:
                conn.is_default = False
                self._save(conn)

    def _save(self, conn: ConnectionInfo) -> None:
        path = self._path(conn.name)
        data = {
            "name": conn.name,
            "url": conn.url,
            "endpoint": conn.endpoint,
            "workspace_id": conn.workspace_id,
            "session_id": conn.session_id,
            "is_default": conn.is_default,
            "created_at": conn.created_at,
            "auth_token_encrypted": conn.auth_token_encrypted,
            "auth_refresh_encrypted": conn.auth_refresh_encrypted,
        }
        path.write_text(json.dumps(data, indent=2))
        path.chmod(0o600)
