"""
Elasticsearch 连接与 API 封装（elasticsearch-py）。
"""
from __future__ import annotations

import contextlib
import base64
import json
import urllib.request
from typing import Any, Generator

from elasticsearch import Elasticsearch

from .config import ConnectionConfig, SafetyConfig


def _unwrap(resp: Any) -> Any:
    """
    elasticsearch-py 8.x API may return TransportApiResponse/ObjectApiResponse wrappers.
    MCP tool schema expects plain JSON-serializable objects.
    """
    if resp is None:
        return None
    body = getattr(resp, "body", None)
    return body if body is not None else resp


def get_connection_config(
    profile: str | None,
    default: ConnectionConfig,
    connection_map: dict[str, ConnectionConfig],
) -> ConnectionConfig:
    if not (profile or "").strip():
        return default
    key = (profile or "").strip()
    if key not in connection_map:
        raise ValueError(f"未知的 connection profile: {profile}，可用: {', '.join(sorted(connection_map))}")
    return connection_map[key]


@contextlib.contextmanager
def connect(
    config: ConnectionConfig,
    safety: SafetyConfig,
) -> Generator[Elasticsearch, None, None]:
    scheme = "https" if config.use_ssl else "http"
    url = f"{scheme}://{config.host}:{config.port}"
    auth = (config.user, config.password) if config.user else None
    # ES 7.x / OpenSearch / 部分代理环境对 Content-Type/Accept 更敏感，统一使用 application/json
    headers = {"Content-Type": "application/json", "Accept": "application/json"}
    client = Elasticsearch(
        [url],
        basic_auth=auth,
        request_timeout=safety.request_timeout,
        headers=headers,
    )
    # Attach config for raw-http execution path
    client._mcp_connection_config = config
    if config.skip_product_check:
        client._verified_elasticsearch = True  # 兼容 OpenSearch/未带 X-Elastic-Product 的服务
    try:
        yield client
    finally:
        client.close()


def list_indices(client: Elasticsearch, pattern: str = "*") -> list[dict[str, Any]]:
    """列出索引，支持 pattern。返回 [{"name": "...", "docs_count": ...}, ...]"""
    cat = _unwrap(client.cat.indices(index=pattern, format="json"))
    if not cat:
        return []
    return [{"name": x.get("index", ""), "docs_count": x.get("docs.count", "0"), "store.size": x.get("store.size", "")} for x in cat]


def get_mapping(client: Elasticsearch, index: str) -> dict[str, Any]:
    """获取索引 mapping。"""
    return _unwrap(client.indices.get_mapping(index=index)) or {}


def execute_request(
    client: Elasticsearch,
    method: str,
    path: str,
    body: dict[str, Any] | None = None,
    request_timeout: int = 30,
) -> dict[str, Any]:
    """执行单条 REST 请求。path 如 /index/_search，不含 host。"""
    # NOTE: For some managed ES gateways (e.g. certain Alibaba Cloud setups),
    # elasticsearch-py (via elastic-transport) may fail to attach Basic auth headers,
    # while raw HTTP Basic auth works. We execute via raw HTTP to be predictable.
    cfg = getattr(client, "_mcp_connection_config", None)
    if not isinstance(cfg, ConnectionConfig):
        # Fallback to elasticsearch-py execution when config is unavailable.
        if body is not None:
            client = client.options(headers={"Content-Type": "application/json", "Accept": "application/json"})
        return _unwrap(client.transport.perform_request(method, path, body=body))

    scheme = "https" if cfg.use_ssl else "http"
    url = f"{scheme}://{cfg.host}:{cfg.port}{path}"
    data = None if body is None else json.dumps(body).encode("utf-8")
    req = urllib.request.Request(url=url, method=method.upper(), data=data)
    req.add_header("Accept", "application/json")
    if body is not None:
        req.add_header("Content-Type", "application/json")
    if cfg.user:
        token = base64.b64encode(f"{cfg.user}:{cfg.password}".encode("utf-8")).decode("ascii")
        req.add_header("Authorization", f"Basic {token}")
    try:
        with urllib.request.urlopen(req, timeout=request_timeout) as resp:
            raw = resp.read()
            if not raw:
                return {}
            try:
                return json.loads(raw.decode("utf-8"))
            except Exception:
                return {"text": raw.decode("utf-8", errors="replace")}
    except urllib.error.HTTPError as e:
        raw = e.read()
        try:
            payload = json.loads(raw.decode("utf-8")) if raw else {}
        except Exception:
            payload = {"text": raw.decode("utf-8", errors="replace")}
        return {"status": e.code, **payload}
