# Copyright (C) 2026 GoodData Corporation

"""Config writer helpers for MCP installation."""

from __future__ import annotations

import re
from collections.abc import Sequence
from pathlib import Path
from typing import cast

import orjson

from aida_mcp.server.mcp_install_common import _make_server_entry

_AIDA_CODEX_SECTION_RE = re.compile(r"^\s*\[mcp_servers\.aida(?:\.[^\]]+)?\]\s*$")
_TOML_TABLE_RE = re.compile(r"^\s*\[[^\]]+\]\s*$")


def write_mcp_json_config(
    *,
    dest: Path,
    command_path: Path | None = None,
    args: Sequence[str] = (),
    url: str | None = None,
    workspace_root: str | None = None,
    config_root: str | None = None,
    extra_env: dict[str, str] | None = None,
) -> None:
    dest.parent.mkdir(parents=True, exist_ok=True)

    data: dict[str, object] = {}
    if dest.exists():
        try:
            loaded: object = orjson.loads(dest.read_text(encoding="utf-8"))
            if isinstance(loaded, dict):
                data = cast(dict[str, object], loaded)
        except Exception:
            data = {}

    servers_raw: object = data.get("mcpServers")
    if isinstance(servers_raw, dict):
        servers: dict[str, object] = cast(dict[str, object], servers_raw)
    else:
        servers = {}
        data["mcpServers"] = servers

    servers["aida"] = _make_server_entry(
        command=str(command_path) if command_path is not None else None,
        args=args,
        url=url,
        workspace_root=workspace_root,
        config_root=config_root,
        extra_env=extra_env,
    )
    dest.write_text(orjson.dumps(data, option=orjson.OPT_INDENT_2).decode("utf-8") + "\n", encoding="utf-8")


def _toml_escape(value: str) -> str:
    return value.replace("\\", "\\\\").replace('"', '\\"')


def _toml_string(value: str) -> str:
    return f'"{_toml_escape(value)}"'


def _toml_array(values: Sequence[str]) -> str:
    rendered = ", ".join(_toml_string(value) for value in values)
    return f"[{rendered}]"


def _render_codex_aida_section(
    *,
    command_path: Path | None,
    args: Sequence[str],
    url: str | None,
    workspace_root: str | None,
    config_root: str | None,
    extra_env: dict[str, str] | None = None,
) -> str:
    env: dict[str, str] = {}
    if extra_env:
        env.update(extra_env)
    if not url and workspace_root:
        env["AIDA_WORKSPACE_ROOT"] = workspace_root
    if config_root:
        env["AIDA_CONFIG_ROOT"] = config_root

    lines: list[str] = ["[mcp_servers.aida]"]
    if url:
        lines.append(f"url = {_toml_string(url)}")
    elif command_path is not None:
        lines.append(f"command = {_toml_string(str(command_path))}")
        lines.append(f"args = {_toml_array(args)}")
    else:
        raise ValueError("Either command_path or url must be provided for Codex MCP config.")
    if env:
        lines.append("")
        lines.append("[mcp_servers.aida.env]")
        for key in sorted(env.keys()):
            lines.append(f"{key} = {_toml_string(env[key])}")
    return "\n".join(lines) + "\n"


def _remove_existing_codex_aida_sections(config_text: str) -> str:
    if not config_text.strip():
        return ""
    lines = config_text.splitlines()
    kept: list[str] = []
    idx = 0
    while idx < len(lines):
        line = lines[idx]
        if _AIDA_CODEX_SECTION_RE.match(line.strip()):
            idx += 1
            while idx < len(lines):
                current = lines[idx].strip()
                if _TOML_TABLE_RE.match(current) and not _AIDA_CODEX_SECTION_RE.match(current):
                    break
                idx += 1
            continue
        kept.append(line)
        idx += 1

    while kept and not kept[-1].strip():
        kept.pop()
    return "\n".join(kept)


def write_codex_toml_config(
    *,
    dest: Path,
    command_path: Path | None = None,
    args: Sequence[str] = (),
    url: str | None = None,
    workspace_root: str | None = None,
    config_root: str | None = None,
    extra_env: dict[str, str] | None = None,
) -> None:
    dest.parent.mkdir(parents=True, exist_ok=True)
    existing = dest.read_text(encoding="utf-8") if dest.exists() else ""
    prefix = _remove_existing_codex_aida_sections(existing)
    section = _render_codex_aida_section(
        command_path=command_path,
        args=args,
        url=url,
        workspace_root=workspace_root,
        config_root=config_root,
        extra_env=extra_env,
    )
    output = prefix.rstrip() + "\n\n" + section if prefix.strip() else section
    dest.write_text(output, encoding="utf-8")
