from __future__ import annotations

import base64
from typing import Any

from mcp.server.fastmcp import FastMCP
from Crypto.PublicKey import RSA
from Crypto.Cipher import PKCS1_v1_5

from ..backend_client import BackendRequestContext, BackendResponse
from ..config import (
    DEFAULT_PROFILE,
    get_default_base_url,
    get_default_qf_version,
    normalize_base_url,
)
from ..errors import QingflowApiError, raise_tool_error
from ..session_store import SessionStore
from .base import ToolBase


class AuthTools(ToolBase):
    def __init__(self, sessions: SessionStore, backend) -> None:
        super().__init__(sessions, backend)

    def register(self, mcp: FastMCP) -> None:
        @mcp.tool()
        def auth_login(
            profile: str = DEFAULT_PROFILE,
            base_url: str | None = None,
            qf_version: str | None = None,
            email: str = "",
            password: str = "",
            persist: bool = True,
        ) -> dict[str, Any]:
            return self.auth_login(profile=profile, base_url=base_url, qf_version=qf_version, email=email, password=password, persist=persist)

        @mcp.tool()
        def auth_use_token(
            profile: str = DEFAULT_PROFILE,
            base_url: str | None = None,
            qf_version: str | None = None,
            token: str = "",
            ws_id: int | None = None,
            persist: bool = False,
        ) -> dict[str, Any]:
            return self.auth_use_token(profile=profile, base_url=base_url, qf_version=qf_version, token=token, ws_id=ws_id, persist=persist)

        @mcp.tool()
        def auth_whoami(profile: str = DEFAULT_PROFILE) -> dict[str, Any]:
            return self.auth_whoami(profile=profile)

        @mcp.tool()
        def auth_logout(profile: str = DEFAULT_PROFILE, forget_persisted: bool = False) -> dict[str, Any]:
            return self.auth_logout(profile=profile, forget_persisted=forget_persisted)

    def auth_login(
        self,
        *,
        profile: str = DEFAULT_PROFILE,
        base_url: str | None = None,
        qf_version: str | None = None,
        email: str,
        password: str,
        persist: bool,
    ) -> dict[str, Any]:
        normalized_base_url = self._normalize_base_url(base_url)
        normalized_qf_version, qf_version_source = self._resolve_qf_version_input(qf_version)
        if not email or not password:
            raise_tool_error(QingflowApiError.config_error("email and password are required"))

        # Try to fetch public key and encrypt password
        public_key_str = self._fetch_public_key(normalized_base_url, qf_version=normalized_qf_version)
        encrypted_password = self._encrypt_password(password, public_key_str) if public_key_str else password

        try:
            # Try 'email' first (aPaas/Public Cloud style)
            login_response = self.backend.public_request_with_meta(
                "POST",
                normalized_base_url,
                "/user/login",
                json_body={"email": email, "password": encrypted_password},
                qf_version=normalized_qf_version,
            )
        except QingflowApiError as error:
            # If failed, try 'account' (QMC/Private Cloud style)
            try:
                login_response = self.backend.public_request_with_meta(
                    "POST",
                    normalized_base_url,
                    "/user/login",
                    json_body={"account": email, "password": encrypted_password},
                    qf_version=normalized_qf_version,
                )
            except QingflowApiError:
                # If both failed, raise the original error
                self._handle_error(profile, error)
                raise AssertionError("unreachable")
        login_result = login_response.data
        if not isinstance(login_result, dict):
            raise_tool_error(QingflowApiError(category="auth", message="Login did not return a valid result"))

        token = login_result.get("token")
        login_token = login_result.get("loginToken")
        
        if not token and login_token:
            raise_tool_error(
                QingflowApiError.not_supported(
                    "Current environment requires an additional login challenge. Qingflow MCP v1 only supports direct token login."
                )
            )
        
        if not token:
            raise_tool_error(QingflowApiError(category="auth", message="Login did not return a valid Qingflow token"))
        detected_qf_version = login_response.qf_response_version
        resolved_qf_version, resolved_qf_version_source = self._resolve_backend_qf_version(
            detected_qf_version,
            fallback_qf_version=normalized_qf_version,
            fallback_source=qf_version_source,
        )

        user_info = login_result.get("userVO") or login_result.get("userInfo") or {}
        if not isinstance(user_info, dict):
            user_info = {}
        verified_user_info, verified_qf_version = self._try_fetch_user_info(
            normalized_base_url,
            token,
            qf_version=resolved_qf_version,
            qf_version_source=resolved_qf_version_source,
        )
        if verified_qf_version:
            resolved_qf_version, resolved_qf_version_source = self._resolve_backend_qf_version(
                verified_qf_version,
                fallback_qf_version=resolved_qf_version,
                fallback_source=resolved_qf_version_source,
            )
        if isinstance(verified_user_info, dict):
            user_info = verified_user_info
        last_ws_info = user_info.get("lastWsInfo") or {}
        session_profile = self.sessions.save_session(
            profile=profile,
            base_url=normalized_base_url,
            qf_version=resolved_qf_version,
            qf_version_source=resolved_qf_version_source,
            token=token,
            login_token=login_token,
            uid=int(user_info.get("uid")),
            email=user_info.get("email"),
            nick_name=user_info.get("nickName"),
            persist=persist,
        )
        return {
            "profile": session_profile.profile,
            "base_url": session_profile.base_url,
            "qf_version": session_profile.qf_version,
            "qf_version_source": session_profile.qf_version_source,
            "uid": session_profile.uid,
            "email": session_profile.email,
            "nick_name": session_profile.nick_name,
            "selected_ws_id": session_profile.selected_ws_id,
            "selected_ws_name": session_profile.selected_ws_name,
            "suggested_ws_id": last_ws_info.get("wsId"),
            "suggested_ws_name": last_ws_info.get("wsName") or last_ws_info.get("workspaceName"),
            "persisted": session_profile.persisted,
            "request_route": self._request_route_payload(
                BackendRequestContext(
                    base_url=session_profile.base_url,
                    token=token,
                    ws_id=session_profile.selected_ws_id,
                    qf_version=session_profile.qf_version,
                    qf_version_source=session_profile.qf_version_source,
                )
            ),
        }

    def auth_use_token(
        self,
        *,
        profile: str = DEFAULT_PROFILE,
        base_url: str | None = None,
        qf_version: str | None = None,
        token: str,
        ws_id: int | None = None,
        persist: bool = False,
    ) -> dict[str, Any]:
        normalized_base_url = self._normalize_base_url(base_url)
        normalized_qf_version, qf_version_source = self._resolve_qf_version_input(qf_version)
        if not token:
            raise_tool_error(QingflowApiError.config_error("token is required"))
        if ws_id is not None and ws_id <= 0:
            raise_tool_error(QingflowApiError.config_error("ws_id must be positive"))
        try:
            user_info, detected_qf_version = self._fetch_user_info(
                normalized_base_url,
                token,
                ws_id,
                qf_version=normalized_qf_version,
                qf_version_source=qf_version_source,
            )
            resolved_qf_version, resolved_qf_version_source = self._resolve_backend_qf_version(
                detected_qf_version,
                fallback_qf_version=normalized_qf_version,
                fallback_source=qf_version_source,
            )
            uid = user_info.get("uid")
            if uid is None:
                raise_tool_error(QingflowApiError(category="auth", message="Token validation did not return valid user info"))
            last_ws_info = user_info.get("lastWsInfo") or {}
            session_profile = self.sessions.save_session(
                profile=profile,
                base_url=normalized_base_url,
                qf_version=resolved_qf_version,
                qf_version_source=resolved_qf_version_source,
                token=token,
                login_token=None,
                uid=int(uid),
                email=user_info.get("email"),
                nick_name=user_info.get("nickName") or user_info.get("displayName") or user_info.get("name"),
                persist=persist,
            )
            selected_ws_name = None
            if ws_id is not None:
                workspace = self._fetch_workspace(
                    normalized_base_url,
                    token,
                    ws_id,
                    qf_version=resolved_qf_version,
                    qf_version_source=resolved_qf_version_source,
                )
                workspace_qf_version = self._workspace_system_version(workspace)
                if resolved_qf_version is None and workspace_qf_version is not None:
                    resolved_qf_version = workspace_qf_version
                    resolved_qf_version_source = "workspace_system_version"
                    session_profile = self.sessions.update_route(
                        profile,
                        qf_version=resolved_qf_version,
                        qf_version_source=resolved_qf_version_source,
                    )
                selected_ws_name = workspace.get("workspaceName") or workspace.get("wsName") or workspace.get("remark")
                session_profile = self.sessions.select_workspace(profile, ws_id=ws_id, ws_name=selected_ws_name)
            return {
                "profile": session_profile.profile,
                "base_url": session_profile.base_url,
                "qf_version": session_profile.qf_version,
                "qf_version_source": session_profile.qf_version_source,
                "uid": session_profile.uid,
                "email": session_profile.email,
                "nick_name": session_profile.nick_name,
                "selected_ws_id": session_profile.selected_ws_id,
                "selected_ws_name": session_profile.selected_ws_name,
                "suggested_ws_id": last_ws_info.get("wsId"),
                "suggested_ws_name": last_ws_info.get("wsName") or last_ws_info.get("workspaceName"),
                "persisted": session_profile.persisted,
                "request_route": self._request_route_payload(
                    BackendRequestContext(
                        base_url=session_profile.base_url,
                        token=token,
                        ws_id=session_profile.selected_ws_id,
                        qf_version=session_profile.qf_version,
                        qf_version_source=session_profile.qf_version_source,
                    )
                ),
            }
        except QingflowApiError as error:
            self._handle_error(profile, error)
            raise AssertionError("unreachable")

    def auth_whoami(self, *, profile: str = DEFAULT_PROFILE) -> dict[str, Any]:
        try:
            session_profile, _, context = self._require_context(profile, require_workspace=False)
        except QingflowApiError as error:
            self._handle_error(profile, error)
            raise AssertionError("unreachable")
        return {
            "profile": session_profile.profile,
            "base_url": session_profile.base_url,
            "qf_version": session_profile.qf_version,
            "qf_version_source": session_profile.qf_version_source,
            "uid": session_profile.uid,
            "email": session_profile.email,
            "nick_name": session_profile.nick_name,
            "selected_ws_id": session_profile.selected_ws_id,
            "selected_ws_name": session_profile.selected_ws_name,
            "persisted": session_profile.persisted,
            "request_route": self._request_route_payload(context),
        }

    def auth_logout(self, *, profile: str = DEFAULT_PROFILE, forget_persisted: bool = False) -> dict[str, Any]:
        if not self.sessions.has_profile(profile):
            raise_tool_error(QingflowApiError.auth_required(profile))
        self.sessions.logout(profile, forget_persisted=forget_persisted)
        return {
            "profile": profile,
            "logged_out": True,
            "forgot_persisted": forget_persisted,
        }

    def _normalize_base_url(self, base_url: str | None) -> str:
        normalized_base_url = normalize_base_url(base_url) or get_default_base_url()
        if not normalized_base_url:
            raise_tool_error(
                QingflowApiError.config_error(
                    "base_url is required or configure default_base_url / QINGFLOW_MCP_DEFAULT_BASE_URL"
                )
            )
        return normalized_base_url

    def _normalize_qf_version(self, qf_version: str | None) -> str | None:
        if qf_version is not None:
            normalized = str(qf_version).strip()
            return normalized or None
        return get_default_qf_version()

    def _resolve_qf_version_input(self, qf_version: str | None) -> tuple[str | None, str]:
        if qf_version is not None:
            normalized = self._normalize_qf_version(qf_version)
            return normalized, "explicit" if normalized else "unset"
        normalized = self._normalize_qf_version(None)
        if normalized:
            return normalized, "default_config"
        return None, "unset"

    def _resolve_backend_qf_version(
        self,
        backend_qf_version: str | None,
        *,
        fallback_qf_version: str | None,
        fallback_source: str,
    ) -> tuple[str | None, str]:
        if backend_qf_version:
            return backend_qf_version, "backend_response"
        return fallback_qf_version, fallback_source

    def _workspace_system_version(self, workspace: Any) -> str | None:
        if not isinstance(workspace, dict):
            return None
        value = workspace.get("systemVersion")
        if value is None:
            return None
        normalized = str(value).strip()
        return normalized or None

    def _fetch_user_info(
        self,
        base_url: str,
        token: str,
        ws_id: int | None,
        *,
        qf_version: str | None,
        qf_version_source: str | None,
    ) -> tuple[dict[str, Any], str | None]:
        request_context = BackendRequestContext(
            base_url=base_url,
            token=token,
            ws_id=ws_id,
            qf_version=qf_version,
            qf_version_source=qf_version_source,
        )
        try:
            user_response = self.backend.request_with_meta("GET", request_context, "/user")
            user_info = user_response.data
            if isinstance(user_info, dict):
                return user_info, user_response.qf_response_version
        except QingflowApiError as original_error:
            if ws_id is not None:
                raise original_error
            first_workspace, workspace_qf_version = self._fetch_first_workspace(
                base_url,
                token,
                qf_version=qf_version,
                qf_version_source=qf_version_source,
            )
            if not first_workspace:
                raise original_error
            first_ws_id = first_workspace.get("wsId")
            if not first_ws_id:
                raise original_error
            effective_qf_version = workspace_qf_version or qf_version
            effective_qf_version_source = "backend_response" if workspace_qf_version else qf_version_source
            fallback_context = BackendRequestContext(
                base_url=base_url,
                token=token,
                ws_id=int(first_ws_id),
                qf_version=effective_qf_version,
                qf_version_source=effective_qf_version_source,
            )
            user_response = self.backend.request_with_meta("GET", fallback_context, "/user")
            user_info = user_response.data
            if isinstance(user_info, dict):
                return user_info, user_response.qf_response_version or effective_qf_version
            raise original_error
        raise_tool_error(QingflowApiError(category="auth", message="Token validation did not return valid user info"))

    def _try_fetch_user_info(
        self,
        base_url: str,
        token: str,
        *,
        qf_version: str | None,
        qf_version_source: str | None,
    ) -> tuple[dict[str, Any] | None, str | None]:
        try:
            return self._fetch_user_info(
                base_url,
                token,
                None,
                qf_version=qf_version,
                qf_version_source=qf_version_source,
            )
        except QingflowApiError:
            return None, None

    def _fetch_first_workspace(
        self,
        base_url: str,
        token: str,
        *,
        qf_version: str | None,
        qf_version_source: str | None,
    ) -> tuple[dict[str, Any] | None, str | None]:
        page_response = self.backend.request_with_meta(
            "POST",
            BackendRequestContext(
                base_url=base_url,
                token=token,
                ws_id=None,
                qf_version=qf_version,
                qf_version_source=qf_version_source,
            ),
            "/user/workspaceList/pageQuery",
            json_body={"pageNum": 1, "pageSize": 1},
        )
        page = page_response.data
        if not isinstance(page, dict):
            return None, page_response.qf_response_version
        workspaces = page.get("list") or []
        if not workspaces:
            return None, page_response.qf_response_version
        first_workspace = workspaces[0]
        return (first_workspace if isinstance(first_workspace, dict) else None), page_response.qf_response_version

    def _fetch_workspace(
        self,
        base_url: str,
        token: str,
        ws_id: int,
        *,
        qf_version: str | None,
        qf_version_source: str | None,
    ) -> dict[str, Any]:
        workspace = self.backend.request(
            "GET",
            BackendRequestContext(
                base_url=base_url,
                token=token,
                ws_id=None,
                qf_version=qf_version,
                qf_version_source=qf_version_source,
            ),
            f"/user/workspace/{ws_id}",
        )
        if not isinstance(workspace, dict):
            raise_tool_error(QingflowApiError(category="workspace", message=f"Workspace {ws_id} is not accessible"))
        return workspace

    def _request_route_payload(self, context: BackendRequestContext) -> dict[str, Any]:
        describe_route = getattr(self.backend, "describe_route", None)
        if callable(describe_route):
            payload = describe_route(context)
            if isinstance(payload, dict):
                return payload
        return {
            "base_url": context.base_url,
            "qf_version": context.qf_version,
            "qf_version_source": context.qf_version_source or ("context" if context.qf_version else "unknown"),
        }

    def _fetch_public_key(self, base_url: str, *, qf_version: str | None) -> str | None:
        # Endpoints to try (order matters, lowercase 'pubkey' is for Public Cloud)
        endpoints = ["/user/pubkey", "/api/user/pubkey", "/user/publicKey", "/api/user/publicKey"]
        for endpoint in endpoints:
            try:
                # We use unwrap=False to handle various response formats
                result = self.backend.public_request("GET", base_url, endpoint, unwrap=False, qf_version=qf_version)
                if isinstance(result, dict):
                    # Try various common response structures
                    data = result.get("data") or result.get("result") or result
                    if isinstance(data, dict):
                        # Try case-insensitive keys
                        for key in ["pubkey", "publicKey", "pubKey"]:
                            if key in data:
                                return str(data[key])
                    if isinstance(data, str) and not data.startswith("{"):
                        return data
            except Exception:
                continue
        
        return None

    def _encrypt_password(self, password: str, public_key_str: str) -> str:
        try:
            if not public_key_str.startswith("-----BEGIN"):
                key_content = public_key_str.strip()
                formatted_key = f"-----BEGIN PUBLIC KEY-----\n{key_content}\n-----END PUBLIC KEY-----"
            else:
                formatted_key = public_key_str
            
            key = RSA.import_key(formatted_key)
            cipher = PKCS1_v1_5.new(key)
            encrypted = cipher.encrypt(password.encode("utf-8"))
            return base64.b64encode(encrypted).decode("utf-8")
        except Exception as e:
            # If encryption fails, fallback to plain text (might be a legacy or custom environment)
            return password