OpenBB/cli/openbb_cli/dispatchers/openapi_schema.py

"""Build CLI menu / parser structures from an OpenAPI 3.x schema."""

from __future__ import annotations

import argparse
import json
import re
from typing import Any

import httpx

PROVIDER_TAG_RE = re.compile(r"\s*\(provider:\s*([^)]+)\)\s*$")
PROVIDER_SECTION_SPLIT_RE = re.compile(r";\s*\n\s*")


def parse_provider_sections(text: str) -> tuple[set[str], bool]:
    """Split an OpenBB-merged help string into ``(tagged_providers, has_untagged)``."""
    tagged: set[str] = set()
    has_untagged = False
    for raw_section in PROVIDER_SECTION_SPLIT_RE.split(text):
        section = raw_section.strip()
        if not section:
            continue
        match = PROVIDER_TAG_RE.search(section)
        if match is None:
            has_untagged = True
            continue
        for p in match.group(1).split(","):
            name = p.strip().lower()
            if name:
                tagged.add(name)
    return tagged, has_untagged


def param_provider_membership(
    schema: dict[str, Any],
    description: str | None,
    providers_set: set[str],
) -> list[str]:
    """Return the providers a parameter belongs to, or ``[]`` if shared by all."""
    providers_lower = {p.lower(): p for p in providers_set}
    if description:
        tagged, has_untagged = parse_provider_sections(description)
        if has_untagged:
            return []
        if tagged:
            return [providers_lower[t] for t in tagged if t in providers_lower]
    keys = [
        k
        for k, v in schema.items()
        if k in providers_set
        and k not in _OPENAPI_RESERVED_SCHEMA_KEYS
        and isinstance(v, dict)
    ]
    if keys:
        return keys
    title = schema.get("title")
    if isinstance(title, str) and title.lower() in providers_lower:
        return [providers_lower[title.lower()]]
    return []


_TYPE_MAP: dict[str, type] = {
    "string": str,
    "integer": int,
    "number": float,
    "boolean": bool,
}

_OPENAPI_RESERVED_SCHEMA_KEYS = frozenset(
    {
        "type",
        "format",
        "title",
        "description",
        "default",
        "enum",
        "const",
        "anyOf",
        "oneOf",
        "allOf",
        "items",
        "properties",
        "additionalProperties",
        "required",
        "nullable",
        "minimum",
        "maximum",
        "minLength",
        "maxLength",
        "pattern",
    }
)


def resolve_ref(spec: dict[str, Any], ref: str) -> dict[str, Any]:
    """Resolve a local OpenAPI ``$ref`` JSON pointer against ``spec``."""
    if not ref.startswith("#/"):
        return {}
    node: Any = spec
    for raw_part in ref[2:].split("/"):
        part = raw_part.replace("~1", "/").replace("~0", "~")
        if not isinstance(node, dict) or part not in node:
            return {}
        node = node[part]
    return node if isinstance(node, dict) else {}


def deref_parameter(spec: dict[str, Any], param: dict[str, Any]) -> dict[str, Any]:
    """Resolve a parameter that may be a ``$ref``."""
    seen: set[str] = set()
    while "$ref" in param:
        ref = param["$ref"]
        if ref in seen:
            return {}
        seen.add(ref)
        param = resolve_ref(spec, ref) or {}
    schema = param.get("schema")
    if isinstance(schema, dict) and "$ref" in schema:
        param = {**param, "schema": resolve_ref(spec, schema["$ref"]) or schema}
    return param


def deref_schema(
    spec: dict[str, Any],
    node: Any,
    seen: frozenset[str] | None = None,
    max_depth: int = 32,
) -> Any:
    """Recursively expand every ``$ref`` inside an OpenAPI schema fragment."""
    if max_depth <= 0:
        return node
    if seen is None:
        seen = frozenset()
    if isinstance(node, dict):
        ref = node.get("$ref")
        if isinstance(ref, str):
            if ref in seen:
                return {"$ref": ref}
            target = resolve_ref(spec, ref)
            if not target:
                return node
            return deref_schema(spec, target, seen | {ref}, max_depth - 1)
        return {k: deref_schema(spec, v, seen, max_depth - 1) for k, v in node.items()}
    if isinstance(node, list):
        return [deref_schema(spec, v, seen, max_depth - 1) for v in node]
    return node


_SUCCESS_PRIORITY = ("200", "2XX", "201", "default")
_JSON_CONTENT_TYPES = ("application/json", "application/vnd.api+json")


def _deref_response(spec: dict[str, Any], response: dict[str, Any]) -> dict[str, Any]:
    if "$ref" in response:
        return resolve_ref(spec, response["$ref"]) or {}
    return response


def extract_response_schema(
    spec: dict[str, Any], operation: dict[str, Any]
) -> dict[str, Any] | None:
    """Pull the primary success-response JSON schema for an operation."""
    responses = operation.get("responses") or {}
    response: dict[str, Any] | None = None
    for key in _SUCCESS_PRIORITY:
        candidate = responses.get(key)
        if isinstance(candidate, dict):
            response = _deref_response(spec, candidate)
            break
    if response is None:
        for value in responses.values():
            if isinstance(value, dict):
                response = _deref_response(spec, value)
                break
    if response is None:
        return None
    content = response.get("content") or {}
    media: dict[str, Any] | None = None
    for ct in _JSON_CONTENT_TYPES:
        candidate = content.get(ct)
        if isinstance(candidate, dict):
            media = candidate
            break
    if media is None:
        for value in content.values():
            if isinstance(value, dict) and "schema" in value:
                media = value
                break
    if media is None:
        return None
    schema = media.get("schema")
    if not isinstance(schema, dict):
        return None
    return deref_schema(spec, schema)


def extract_request_body_schema(
    spec: dict[str, Any], operation: dict[str, Any]
) -> dict[str, Any] | None:
    """Pull the primary request-body JSON schema for an operation."""
    rb = operation.get("requestBody")
    if not isinstance(rb, dict):
        return None
    if "$ref" in rb:
        rb = resolve_ref(spec, rb["$ref"]) or {}
    content = rb.get("content") or {}
    media: dict[str, Any] | None = None
    for ct in _JSON_CONTENT_TYPES:
        candidate = content.get(ct)
        if isinstance(candidate, dict):
            media = candidate
            break
    if media is None:
        for value in content.values():
            if isinstance(value, dict) and "schema" in value:
                media = value
                break
    if media is None:
        return None
    schema = media.get("schema")
    if not isinstance(schema, dict):
        return None
    return deref_schema(spec, schema)


def extract_response_schemas(
    spec: dict[str, Any], operation: dict[str, Any]
) -> dict[str, dict[str, dict[str, Any]]]:
    """Return every response schema as ``{status_code: {content_type: schema}}``."""
    out: dict[str, dict[str, dict[str, Any]]] = {}
    responses = operation.get("responses") or {}
    for status, response in responses.items():
        if not isinstance(response, dict):
            continue
        deref_resp = _deref_response(spec, response)
        content = deref_resp.get("content") or {}
        per_content: dict[str, dict[str, Any]] = {}
        for content_type, media in content.items():
            if not isinstance(media, dict):
                continue
            schema = media.get("schema")
            if isinstance(schema, dict):
                per_content[content_type] = deref_schema(spec, schema)
        if per_content:
            out[status] = per_content
    return out


def _resolve_schema(schema: dict[str, Any]) -> tuple[type, list[Any], bool]:
    """Resolve an OpenAPI parameter schema to ``(py_type, choices, is_list)``."""
    if "anyOf" in schema:
        non_null = [s for s in schema["anyOf"] if s.get("type") != "null"]
        if not non_null:
            return (str, [], False)
        if len(non_null) == 1:
            return _resolve_schema(non_null[0])
        choices: list[Any] = []
        is_list = False
        for s in non_null:
            if s.get("type") == "array":
                is_list = True
            for v in s.get("enum", []):
                if v not in choices:
                    choices.append(v)
        return (str, choices, is_list)

    if schema.get("type") == "array":
        items = schema.get("items", {}) or {}
        py_type, item_choices, _ = _resolve_schema(items)
        return (py_type, item_choices, True)

    if "const" in schema:
        return (type(schema["const"]), [schema["const"]], False)

    py_type = _TYPE_MAP.get(schema.get("type", "string"), str)
    return (py_type, list(schema.get("enum", [])), False)


def _is_json_arg(schema: dict[str, Any]) -> bool:
    """Whether a request-body field must be supplied as a raw JSON value."""
    if schema.get("type") == "object" or "properties" in schema:
        return True
    if schema.get("additionalProperties"):
        return True
    if schema.get("type") == "array":
        items = schema.get("items")
        return isinstance(items, dict) and _is_json_arg(items)
    return any(
        isinstance(member, dict) and _is_json_arg(member)
        for member in schema.get("anyOf", []) or []
    )


def parse_json_arg(raw: str) -> Any:
    """Decode a CLI argument value that carries a JSON document."""
    try:
        return json.loads(raw)
    except json.JSONDecodeError as exc:
        raise argparse.ArgumentTypeError(f"invalid JSON: {exc}") from exc


def request_body_parameters(
    body_schema: dict[str, Any] | None,
) -> list[dict[str, Any]]:
    """Flatten a dereferenced request-body schema into OpenAPI parameter objects."""
    if not isinstance(body_schema, dict) or body_schema.get("type") != "object":
        return []
    required = set(body_schema.get("required") or [])
    out: list[dict[str, Any]] = []
    for name, prop in (body_schema.get("properties") or {}).items():
        if not isinstance(prop, dict):
            continue
        out.append(
            {
                "name": name,
                "in": "body",
                "required": name in required,
                "description": prop.get("description") or prop.get("title"),
                "schema": prop,
                "_json_arg": _is_json_arg(prop),
            }
        )
    return out


def _provider_choices(schema: dict[str, Any]) -> list[Any]:
    """Union of per-provider ``choices`` lists found in OpenBB extension keys."""
    out: list[Any] = []
    for key, value in schema.items():
        if key in _OPENAPI_RESERVED_SCHEMA_KEYS:
            continue
        if isinstance(value, dict) and isinstance(value.get("choices"), list):
            for c in value["choices"]:
                if c not in out:
                    out.append(c)
    return out


def _escape_help(text: str | None) -> str | None:
    """Escape lone ``%`` so argparse's %-formatting validator accepts the help text."""
    if text is None:
        return None
    return text.replace("%", "%%")


def parameter_to_kwargs(param: dict[str, Any]) -> tuple[str, dict[str, Any]] | None:
    """Translate one OpenAPI parameter to ``(flag, add_argument kwargs)``."""
    name = param.get("name")
    if not name or name == "chart":
        return None

    schema = param.get("schema", {}) or {}
    py_type, enum_choices, is_list = _resolve_schema(schema)
    provider_choices = _provider_choices(schema)
    choices: list[Any] = list(enum_choices)
    for c in provider_choices:
        if c not in choices:
            choices.append(c)

    kwargs: dict[str, Any] = {
        "dest": name,
        "help": _escape_help(param.get("description") or schema.get("description")),
    }

    if param.get("_json_arg"):
        kwargs["type"] = parse_json_arg
        if param.get("required"):
            kwargs["required"] = True
        return (f"--{name}", kwargs)

    if py_type is bool:
        kwargs.update(action="store_true", default=bool(schema.get("default", False)))
        return (f"--{name}", kwargs)

    kwargs["type"] = py_type
    if "default" in schema:
        kwargs["default"] = schema["default"]
    if choices:
        kwargs["choices"] = choices
    if is_list:
        kwargs["nargs"] = "+"
    if param.get("required") and "default" not in kwargs:
        kwargs["required"] = True

    return (f"--{name}", kwargs)


def build_parser_from_operation(
    op: dict[str, Any], spec: dict[str, Any] | None = None
) -> argparse.ArgumentParser:
    """Build an ``ArgumentParser`` from an OpenAPI operation object."""
    parser = argparse.ArgumentParser(
        prog=op.get("operationId", "cmd"),
        description=(op.get("description") or op.get("summary") or "").strip() or None,
        add_help=False,
        formatter_class=argparse.RawTextHelpFormatter,
    )
    body_params = (
        request_body_parameters(extract_request_body_schema(spec, op))
        if spec is not None
        else []
    )
    for param in [*op.get("parameters", []), *body_params]:
        translated = parameter_to_kwargs(param)
        if translated is None:
            continue
        flag, kwargs = translated
        try:
            parser.add_argument(flag, **kwargs)
        except argparse.ArgumentError:
            continue
    return parser


def _strip_placeholders(segment: str) -> str:
    """Remove ``{placeholder}`` substrings from a URL segment."""
    out = segment
    while "{" in out and "}" in out:
        i = out.find("{")
        j = out.find("}", i)
        if j < 0:
            break
        out = out[:i] + out[j + 1 :]
    return out.strip(".")


def detect_api_prefix(spec: dict[str, Any]) -> str:
    """Compute the longest slash-segmented prefix shared by every path."""
    paths = list(spec.get("paths", {}).keys())
    if not paths:
        return "/api/v1"
    parts_lists = [p.split("/")[:-1] for p in paths]
    common: list[str] = parts_lists[0]
    for parts in parts_lists[1:]:
        new_common: list[str] = []
        for a, b in zip(common, parts):
            if a == b:
                new_common.append(a)
            else:
                break
        common = new_common
        if not common:
            break
    prefix = "/".join(common).rstrip("/")
    return prefix or ""


def url_to_command(url: str, api_prefix: str = "/api/v1") -> str:
    """Convert a URL path to a dotted command, dropping ``{path_params}``."""
    prefix_parts = [p for p in api_prefix.strip("/").split("/") if p]
    parts = [p for p in url.strip("/").split("/") if p]
    if parts[: len(prefix_parts)] == prefix_parts:
        parts = parts[len(prefix_parts) :]
    cleaned = [_strip_placeholders(p) for p in parts]
    return ".".join(p for p in cleaned if p)


def build_command_index(
    spec: dict[str, Any], *, api_prefix: str = "/api/v1"
) -> dict[str, argparse.ArgumentParser]:
    """Map each ``paths`` entry in the spec to an ArgumentParser by dotted command."""
    index: dict[str, argparse.ArgumentParser] = {}
    for url, methods in spec.get("paths", {}).items():
        op = methods.get("get") or methods.get("post")
        if not op:
            continue
        index[url_to_command(url, api_prefix=api_prefix)] = build_parser_from_operation(
            op, spec
        )
    return index


def build_router_map(
    spec: dict[str, Any], *, api_prefix: str = "/api/v1"
) -> dict[str, str]:
    """Classify every dotted path in ``spec`` as ``"menu"`` or ``"command"``."""
    out: dict[str, str] = {}
    for url, methods in spec.get("paths", {}).items():
        if not (methods.get("get") or methods.get("post")):
            continue
        cmd = url_to_command(url, api_prefix=api_prefix)
        if not cmd:
            continue
        parts = cmd.split(".")
        for i in range(1, len(parts)):
            out[".".join(parts[:i])] = "menu"
        out.setdefault(cmd, "command")
    return out


def build_reference(
    spec: dict[str, Any], *, api_prefix: str = "/api/v1"
) -> dict[str, Any]:
    """Mimic ``obb.reference`` — ``paths`` for commands, ``routers`` for menus."""
    prefix_parts = [p for p in api_prefix.strip("/").split("/") if p]
    tag_descriptions: dict[str, str] = {
        t["name"]: t.get("description", "")
        for t in spec.get("tags", [])
        if isinstance(t, dict) and t.get("name")
    }

    router_map = build_router_map(spec, api_prefix=api_prefix)

    paths_out: dict[str, dict[str, Any]] = {}
    routers_out: dict[str, dict[str, Any]] = {}
    for url, methods in spec.get("paths", {}).items():
        op = methods.get("get") or methods.get("post")
        if not op:
            continue
        parts = [p for p in url.strip("/").split("/") if p]
        if parts[: len(prefix_parts)] == prefix_parts:
            parts = parts[len(prefix_parts) :]
        non_template = [p for p in parts if not (p.startswith("{") and p.endswith("}"))]
        cli_path = "/" + "/".join(non_template)
        op_desc = (op.get("description") or op.get("summary") or "").strip()
        if cli_path not in paths_out or not paths_out[cli_path].get("description"):
            paths_out[cli_path] = {"description": op_desc}
        tags = op.get("tags") or []
        tag_desc = tag_descriptions.get(tags[0]) if tags else ""
        for i in range(1, len(non_template) + 1):
            dotted = ".".join(non_template[:i])
            if router_map.get(dotted) != "menu":
                continue
            menu_path = "/" + "/".join(non_template[:i]) + "/"
            if not routers_out.get(menu_path, {}).get("description"):
                routers_out[menu_path] = {"description": tag_desc or ""}
    return {"paths": paths_out, "routers": routers_out}


def _parse_spec_text(text: str, *, content_type: str = "") -> dict[str, Any]:
    """Parse a fetched spec body, choosing JSON or YAML by content sniff."""
    stripped = text.lstrip()
    if stripped.startswith(("{", "[")):
        return json.loads(text)
    if "yaml" in content_type or "yml" in content_type:
        return _yaml_load(text)
    if stripped[:8] in ("openapi:", "swagger:"):
        return _yaml_load(text)
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        return _yaml_load(text)


def _yaml_load(text: str) -> dict[str, Any]:
    """Parse a YAML document."""
    import yaml

    return yaml.safe_load(text)


_EMBEDDED_SPEC_MARKERS: tuple[str, ...] = (
    "var spec = ",
    "const spec = ",
    "let spec = ",
    "window.spec = ",
    '"spec":',
    "spec:",
)


def _find_matching_brace(text: str, start: int) -> int | None:
    """Return the index *after* the ``}`` that matches the ``{`` at ``start``."""
    if start >= len(text) or text[start] != "{":
        return None
    depth = 0
    in_string = False
    escape = False
    for i in range(start, len(text)):
        c = text[i]
        if escape:
            escape = False
            continue
        if c == "\\":
            escape = True
            continue
        if c == '"':
            in_string = not in_string
            continue
        if in_string:
            continue
        if c == "{":
            depth += 1
        elif c == "}":
            depth -= 1
            if depth == 0:
                return i + 1
    return None


def _extract_embedded_spec(html: str) -> dict[str, Any] | None:
    """Scan an HTML page for an embedded OpenAPI / Swagger spec."""
    for marker in _EMBEDDED_SPEC_MARKERS:
        idx = 0
        while True:
            i = html.find(marker, idx)
            if i < 0:
                break
            start = i + len(marker)
            while start < len(html) and html[start] in " \t\r\n":
                start += 1
            end = _find_matching_brace(html, start)
            if end is None:
                idx = i + 1
                continue
            try:
                obj = json.loads(html[start:end])
            except json.JSONDecodeError:
                idx = i + 1
                continue
            if isinstance(obj, dict) and ("openapi" in obj or "swagger" in obj):
                return obj
            idx = end
    return None


def fetch_openapi(
    base_url: str,
    *,
    timeout: float = 10.0,
    path: str | None = None,
    headers: dict[str, str] | None = None,
    query_params: dict[str, str] | None = None,
) -> dict[str, Any]:
    """Fetch the OpenAPI spec at ``base_url`` and return the parsed dict."""
    explicit_path = path is not None
    full_url = (
        path
        if path and (path.startswith("http://") or path.startswith("https://"))
        else f"{base_url.rstrip('/')}{path or '/openapi.json'}"
    )
    merged_headers = {"User-Agent": "openbb-cli/1.0"}
    if headers:
        merged_headers.update(headers)

    response = httpx.get(
        full_url,
        timeout=timeout,
        follow_redirects=True,
        headers=merged_headers,
        params=query_params or None,
    )
    if response.status_code < 400:
        try:
            return _ensure_openapi_dict(
                _parse_spec_text(
                    response.text,
                    content_type=response.headers.get("content-type", ""),
                ),
                full_url,
            )
        except (json.JSONDecodeError, ValueError):
            pass

    if explicit_path:
        response.raise_for_status()
        return _ensure_openapi_dict(
            _parse_spec_text(
                response.text,
                content_type=response.headers.get("content-type", ""),
            ),
            full_url,
        )

    landing_url = base_url.rstrip("/") + "/"
    landing = httpx.get(
        landing_url,
        timeout=timeout,
        follow_redirects=True,
        headers=merged_headers,
        params=query_params or None,
    )
    landing.raise_for_status()
    embedded = _extract_embedded_spec(landing.text)
    if embedded is not None:
        return embedded

    response.raise_for_status()
    return _ensure_openapi_dict(
        _parse_spec_text(
            response.text,
            content_type=response.headers.get("content-type", ""),
        ),
        full_url,
    )


def _ensure_openapi_dict(parsed: Any, source_url: str) -> dict[str, Any]:
    """Reject parsed bodies that aren't a JSON object (and thus can't be an OpenAPI doc)."""
    if not isinstance(parsed, dict):
        raise ValueError(
            f"{source_url!r} returned a "
            f"{type(parsed).__name__}, not an OpenAPI document. "
            "Pass --openapi-path to point at the real spec endpoint "
            "(e.g. /swagger/v1/swagger.json)."
        )
    return parsed