"""ReflACT Model backend — Azure OpenAI wrapper with token tracking.

Provides optimizer/target dual-deployment chat functions and a global
TokenTracker for per-stage cost accounting. Previously llm/azure_openai.py.
"""
from __future__ import annotations

import json
import os
import subprocess
import threading
import time
from types import SimpleNamespace
from typing import Any
from openai import AzureOpenAI, OpenAI

# Sentinel value used as the api_version when the "openai_compatible"
# auth_mode is selected. Real Azure deployments never use this string,
# so it doubles as a marker for downstream type narrowing.
_OPENAI_COMPATIBLE_API_VERSION = "openai-compat"

# ── Configuration ─────────────────────────────────────────────────────────────

ENDPOINT = os.environ.get(
    "AZURE_OPENAI_ENDPOINT",
    "",  # Set via env var or config: e.g. "https://your-resource.openai.azure.com/"
)
API_VERSION = os.environ.get("AZURE_OPENAI_API_VERSION", "2024-12-01-preview")
API_KEY = os.environ.get(
    "AZURE_OPENAI_API_KEY",
    "",
)
AUTH_MODE = os.environ.get("AZURE_OPENAI_AUTH_MODE", "azure_cli").strip().lower()
AD_SCOPE = os.environ.get(
    "AZURE_OPENAI_AD_SCOPE",
    "https://cognitiveservices.azure.com/.default",
)
MANAGED_IDENTITY_CLIENT_ID = os.environ.get(
    "AZURE_OPENAI_MANAGED_IDENTITY_CLIENT_ID",
    "",
).strip()

OPTIMIZER_ENDPOINT = (
    os.environ.get("OPTIMIZER_AZURE_OPENAI_ENDPOINT")
    or os.environ.get("AZURE_OPENAI_OPTIMIZER_ENDPOINT")
    or ENDPOINT
)
TARGET_ENDPOINT = (
    os.environ.get("TARGET_AZURE_OPENAI_ENDPOINT")
    or os.environ.get("AZURE_OPENAI_TARGET_ENDPOINT")
    or ENDPOINT
)
OPTIMIZER_API_VERSION = (
    os.environ.get("OPTIMIZER_AZURE_OPENAI_API_VERSION")
    or os.environ.get("AZURE_OPENAI_OPTIMIZER_API_VERSION")
    or API_VERSION
)
TARGET_API_VERSION = (
    os.environ.get("TARGET_AZURE_OPENAI_API_VERSION")
    or os.environ.get("AZURE_OPENAI_TARGET_API_VERSION")
    or API_VERSION
)
OPTIMIZER_API_KEY = (
    os.environ.get("OPTIMIZER_AZURE_OPENAI_API_KEY")
    or os.environ.get("AZURE_OPENAI_OPTIMIZER_API_KEY")
    or API_KEY
)
TARGET_API_KEY = (
    os.environ.get("TARGET_AZURE_OPENAI_API_KEY")
    or os.environ.get("AZURE_OPENAI_TARGET_API_KEY")
    or API_KEY
)
OPTIMIZER_AUTH_MODE = (
    os.environ.get("OPTIMIZER_AZURE_OPENAI_AUTH_MODE")
    or os.environ.get("AZURE_OPENAI_OPTIMIZER_AUTH_MODE")
    or AUTH_MODE
).strip().lower()
TARGET_AUTH_MODE = (
    os.environ.get("TARGET_AZURE_OPENAI_AUTH_MODE")
    or os.environ.get("AZURE_OPENAI_TARGET_AUTH_MODE")
    or AUTH_MODE
).strip().lower()
OPTIMIZER_AD_SCOPE = (
    os.environ.get("OPTIMIZER_AZURE_OPENAI_AD_SCOPE")
    or os.environ.get("AZURE_OPENAI_OPTIMIZER_AD_SCOPE")
    or AD_SCOPE
)
TARGET_AD_SCOPE = (
    os.environ.get("TARGET_AZURE_OPENAI_AD_SCOPE")
    or os.environ.get("AZURE_OPENAI_TARGET_AD_SCOPE")
    or AD_SCOPE
)
OPTIMIZER_MANAGED_IDENTITY_CLIENT_ID = (
    os.environ.get("OPTIMIZER_AZURE_OPENAI_MANAGED_IDENTITY_CLIENT_ID")
    or os.environ.get("AZURE_OPENAI_OPTIMIZER_MANAGED_IDENTITY_CLIENT_ID")
    or MANAGED_IDENTITY_CLIENT_ID
).strip()
TARGET_MANAGED_IDENTITY_CLIENT_ID = (
    os.environ.get("TARGET_AZURE_OPENAI_MANAGED_IDENTITY_CLIENT_ID")
    or os.environ.get("AZURE_OPENAI_TARGET_MANAGED_IDENTITY_CLIENT_ID")
    or MANAGED_IDENTITY_CLIENT_ID
).strip()

OPTIMIZER_DEPLOYMENT = os.environ.get("OPTIMIZER_DEPLOYMENT", "gpt-4o")
TARGET_DEPLOYMENT = os.environ.get("TARGET_DEPLOYMENT", "gpt-4o")

REASONING_EFFORT: str | None = None

_AZ_CLI_TOKEN_CACHE: dict[str, dict[str, Any]] = {}

# Deployments that require Responses API
_RESPONSES_API_MODELS = {
    "gpt-5.3-codex", "gpt-5.1-codex", "gpt-5.2-codex",
    "gpt-5-codex", "codex-mini", "gpt-5.4-pro",
}


# ── Token Tracker ─────────────────────────────────────────────────────────────

class TokenTracker:
    """Thread-safe per-stage token counter."""

    def __init__(self) -> None:
        self._lock = threading.Lock()
        self._data: dict[str, dict] = {}

    def record(
        self, stage: str, prompt_tokens: int, completion_tokens: int,
    ) -> None:
        with self._lock:
            if stage not in self._data:
                self._data[stage] = {
                    "calls": 0,
                    "prompt_tokens": 0,
                    "completion_tokens": 0,
                }
            d = self._data[stage]
            d["calls"] += 1
            d["prompt_tokens"] += prompt_tokens
            d["completion_tokens"] += completion_tokens

    def summary(self) -> dict:
        with self._lock:
            out: dict = {}
            total_p = total_c = total_calls = 0
            for stage, d in sorted(self._data.items()):
                out[stage] = {
                    "calls": d["calls"],
                    "prompt_tokens": d["prompt_tokens"],
                    "completion_tokens": d["completion_tokens"],
                    "total_tokens": d["prompt_tokens"] + d["completion_tokens"],
                }
                total_p += d["prompt_tokens"]
                total_c += d["completion_tokens"]
                total_calls += d["calls"]
            out["_total"] = {
                "calls": total_calls,
                "prompt_tokens": total_p,
                "completion_tokens": total_c,
                "total_tokens": total_p + total_c,
            }
            return out

    def reset(self) -> None:
        with self._lock:
            self._data.clear()

    def stage_snapshot(self, stage: str) -> dict:
        """Return a copy of one stage's counters (or zeros if not tracked yet)."""
        with self._lock:
            d = self._data.get(stage, {})
            return {
                "calls": d.get("calls", 0),
                "prompt_tokens": d.get("prompt_tokens", 0),
                "completion_tokens": d.get("completion_tokens", 0),
                "total_tokens": d.get("prompt_tokens", 0) + d.get("completion_tokens", 0),
            }


tracker = TokenTracker()


# ── Client management ─────────────────────────────────────────────────────────

_optimizer_client: AzureOpenAI | OpenAI | None = None
_target_client: AzureOpenAI | OpenAI | None = None
_optimizer_lock = threading.Lock()
_target_lock = threading.Lock()


def _role_config(role: str) -> dict[str, str]:
    if role == "optimizer":
        return {
            "endpoint": OPTIMIZER_ENDPOINT,
            "api_version": OPTIMIZER_API_VERSION,
            "api_key": OPTIMIZER_API_KEY,
            "auth_mode": OPTIMIZER_AUTH_MODE,
            "ad_scope": OPTIMIZER_AD_SCOPE,
            "managed_identity_client_id": OPTIMIZER_MANAGED_IDENTITY_CLIENT_ID,
        }
    if role == "target":
        return {
            "endpoint": TARGET_ENDPOINT,
            "api_version": TARGET_API_VERSION,
            "api_key": TARGET_API_KEY,
            "auth_mode": TARGET_AUTH_MODE,
            "ad_scope": TARGET_AD_SCOPE,
            "managed_identity_client_id": TARGET_MANAGED_IDENTITY_CLIENT_ID,
        }
    raise ValueError(f"Unknown Azure OpenAI client role: {role!r}")


def _make_token_provider(
    auth_mode: str,
    ad_scope: str,
    managed_identity_client_id: str,
):
    try:
        from azure.identity import (  # type: ignore[import-not-found]
            AzureCliCredential,
            ManagedIdentityCredential,
            get_bearer_token_provider,
        )
    except ImportError as e:
        if auth_mode == "azure_cli":
            return _make_azure_cli_token_provider(ad_scope)
        raise ImportError(
            "Azure AD auth requires azure-identity. Install it with `pip install azure-identity`."
        ) from e

    if auth_mode in {"managed_identity", "aad", "azure_ad"}:
        if managed_identity_client_id:
            credential = ManagedIdentityCredential(client_id=managed_identity_client_id)
        else:
            credential = ManagedIdentityCredential()
    elif auth_mode == "azure_cli":
        credential = AzureCliCredential()
    else:
        raise ValueError(
            "Unsupported Azure OpenAI auth mode "
            f"{auth_mode!r}; expected api_key, managed_identity, azure_ad, aad, or azure_cli."
        )
    return get_bearer_token_provider(credential, ad_scope)


def _make_azure_cli_token_provider(ad_scope: str):
    """Return an Azure CLI token provider compatible with AzureOpenAI.

    This fallback avoids requiring azure-identity in environments where `az`
    is already logged in. The SDK calls this provider whenever it needs a
    bearer token.
    """

    resource = ad_scope.removesuffix("/.default")

    def _provider() -> str:
        now = int(time.time())
        cache = _AZ_CLI_TOKEN_CACHE.setdefault(resource, {"token": "", "expires_on": 0})
        cached = str(cache.get("token") or "")
        expires_on = int(cache.get("expires_on") or 0)
        if cached and expires_on - now > 300:
            return cached

        raw = subprocess.check_output(
            [
                "az",
                "account",
                "get-access-token",
                "--resource",
                resource,
                "-o",
                "json",
            ],
            text=True,
            stderr=subprocess.STDOUT,
        )
        payload = json.loads(raw)
        token = str(payload["accessToken"])
        cache["token"] = token
        cache["expires_on"] = int(payload.get("expires_on") or now + 3000)
        return token

    return _provider


def _make_client(role: str) -> AzureOpenAI | OpenAI:
    cfg = _role_config(role)
    if not cfg["endpoint"]:
        raise ValueError(
            f"Azure OpenAI endpoint is not configured for {role}. "
            "Pass --azure_openai_endpoint https://your-resource.openai.azure.com/ "
            "or set AZURE_OPENAI_ENDPOINT in your environment."
        )
    auth_mode = cfg["auth_mode"]
    if auth_mode in {"openai_compatible", "compat", "openai"}:
        return OpenAI(
            base_url=cfg["endpoint"].rstrip("/"),
            api_key=cfg["api_key"] or "dummy",
            default_headers={"User-Agent": "SkillOpt"},
        )
    if auth_mode in {"api_key", "key"}:
        if not cfg["api_key"]:
            raise ValueError(
                f"Azure OpenAI API key is not configured for {role}. "
                "Set model.azure_openai_api_key in the config or export AZURE_OPENAI_API_KEY."
            )
        return AzureOpenAI(
            api_version=cfg["api_version"],
            azure_endpoint=cfg["endpoint"],
            api_key=cfg["api_key"],
        )
    return AzureOpenAI(
        api_version=cfg["api_version"],
        azure_endpoint=cfg["endpoint"],
        azure_ad_token_provider=_make_token_provider(
            auth_mode,
            cfg["ad_scope"],
            cfg["managed_identity_client_id"],
        ),
    )


def get_optimizer_client() -> AzureOpenAI | OpenAI:
    global _optimizer_client
    with _optimizer_lock:
        if _optimizer_client is None:
            _optimizer_client = _make_client("optimizer")
        return _optimizer_client


def get_target_client() -> AzureOpenAI | OpenAI:
    global _target_client
    with _target_lock:
        if _target_client is None:
            # When using qwen_chat backend, return an OpenAI client pointing to vLLM
            from skillopt.model.backend_config import get_target_backend
            if get_target_backend() == "qwen_chat":
                from skillopt.model import qwen_backend as _qwen
                target_config = _qwen.TARGET_CONFIG
                _target_client = OpenAI(
                    base_url=target_config.base_url,
                    api_key=target_config.api_key or "dummy",
                )
            else:
                _target_client = _make_client("target")
        return _target_client


def _needs_responses_api(deployment: str) -> bool:
    dep = deployment.lower()
    return any(dep == m or dep.startswith(m + "-") for m in _RESPONSES_API_MODELS)


# ── Core chat function ────────────────────────────────────────────────────────

def _chat_impl(
    client: AzureOpenAI | OpenAI,
    deployment: str,
    system: str,
    user: str,
    max_completion_tokens: int,
    retries: int,
    stage: str,
    reasoning_effort: str | None = None,
    timeout: int | None = None,
) -> tuple[str, dict]:
    """Call LLM, track tokens, return (text, usage_dict)."""
    last_err = None
    usage_info = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}

    for attempt in range(retries):
        try:
            if _needs_responses_api(deployment):
                kwargs: dict[str, Any] = {
                    "model": deployment,
                    "instructions": system,
                    "input": [{"role": "user", "content": user}],
                    "max_output_tokens": max_completion_tokens,
                }
                actual_effort = reasoning_effort or REASONING_EFFORT
                if actual_effort:
                    kwargs["reasoning"] = {"effort": actual_effort}
                if timeout is not None:
                    kwargs["timeout"] = timeout
                resp = client.responses.create(**kwargs)
                text = getattr(resp, "output_text", None) or ""
                if not text:
                    for item in getattr(resp, "output", None) or []:
                        for part in getattr(item, "content", []):
                            if getattr(part, "type", "") == "output_text":
                                text = part.text or ""
                if hasattr(resp, "usage") and resp.usage:
                    usage_info = {
                        "prompt_tokens": getattr(resp.usage, "input_tokens", 0) or 0,
                        "completion_tokens": getattr(resp.usage, "output_tokens", 0) or 0,
                        "total_tokens": (
                            (getattr(resp.usage, "input_tokens", 0) or 0)
                            + (getattr(resp.usage, "output_tokens", 0) or 0)
                        ),
                    }
            else:
                kwargs: dict[str, Any] = dict(
                    model=deployment,
                    messages=[
                        {"role": "system", "content": system},
                        {"role": "user", "content": user},
                    ],
                    max_completion_tokens=max_completion_tokens,
                )
                actual_effort = reasoning_effort or REASONING_EFFORT
                if actual_effort is not None:
                    kwargs["reasoning_effort"] = actual_effort
                if timeout is not None:
                    kwargs["timeout"] = timeout
                resp = client.chat.completions.create(**kwargs)
                text = resp.choices[0].message.content or ""
                if resp.usage:
                    usage_info = {
                        "prompt_tokens": resp.usage.prompt_tokens or 0,
                        "completion_tokens": resp.usage.completion_tokens or 0,
                        "total_tokens": resp.usage.total_tokens or 0,
                    }

            tracker.record(
                stage,
                usage_info["prompt_tokens"],
                usage_info["completion_tokens"],
            )
            return text, usage_info

        except Exception as e:  # noqa: BLE001
            last_err = e
            sleep = min(2 ** attempt, 30)
            time.sleep(sleep)

    raise RuntimeError(f"LLM call failed after {retries} retries: {last_err}")


def _chat_messages_impl(
    client: AzureOpenAI | OpenAI,
    deployment: str,
    messages: list[dict[str, Any]],
    max_completion_tokens: int,
    retries: int,
    stage: str,
    reasoning_effort: str | None = None,
    *,
    tools: list[dict[str, Any]] | None = None,
    tool_choice: str | dict[str, Any] | None = None,
    return_message: bool = False,
    timeout: int | None = None,
) -> tuple[Any, dict]:
    """Call the model with a pre-built message list."""
    last_err = None
    usage_info = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}

    for attempt in range(retries):
        try:
            if _needs_responses_api(deployment):
                input_items, instructions = _messages_to_responses_input(messages)
                kwargs: dict[str, Any] = {
                    "model": deployment,
                    "input": input_items,
                    "max_output_tokens": max_completion_tokens,
                }
                if instructions:
                    kwargs["instructions"] = instructions
                actual_effort = reasoning_effort or REASONING_EFFORT
                if actual_effort:
                    kwargs["reasoning"] = {"effort": actual_effort}
                if tools:
                    kwargs["tools"] = [_chat_tool_to_responses_tool(tool) for tool in tools]
                    if tool_choice is not None:
                        kwargs["tool_choice"] = tool_choice
                if timeout is not None:
                    kwargs["timeout"] = timeout
                resp = client.responses.create(**kwargs)
                message, text = _responses_to_chat_message(resp)
                if hasattr(resp, "usage") and resp.usage:
                    usage_info = {
                        "prompt_tokens": getattr(resp.usage, "input_tokens", 0) or 0,
                        "completion_tokens": getattr(resp.usage, "output_tokens", 0) or 0,
                        "total_tokens": (
                            (getattr(resp.usage, "input_tokens", 0) or 0)
                            + (getattr(resp.usage, "output_tokens", 0) or 0)
                        ),
                    }
            else:
                kwargs = dict(
                    model=deployment,
                    messages=messages,
                    max_completion_tokens=max_completion_tokens,
                )
                actual_effort = reasoning_effort or REASONING_EFFORT
                if tools:
                    kwargs["tools"] = tools
                    if tool_choice is not None:
                        kwargs["tool_choice"] = tool_choice
                    # Some models (e.g. gpt-5.5) don't support reasoning_effort with function tools
                elif actual_effort is not None:
                    kwargs["reasoning_effort"] = actual_effort
                if timeout is not None:
                    kwargs["timeout"] = timeout
                resp = client.chat.completions.create(**kwargs)
                message = resp.choices[0].message
                text = message.content or ""
                if resp.usage:
                    usage_info = {
                        "prompt_tokens": resp.usage.prompt_tokens or 0,
                        "completion_tokens": resp.usage.completion_tokens or 0,
                        "total_tokens": resp.usage.total_tokens or 0,
                    }
            tracker.record(
                stage,
                usage_info["prompt_tokens"],
                usage_info["completion_tokens"],
            )
            return (message if return_message else text), usage_info
        except Exception as e:  # noqa: BLE001
            last_err = e
            sleep = min(2 ** attempt, 30)
            time.sleep(sleep)

    raise RuntimeError(f"LLM message call failed after {retries} retries: {last_err}")


def _chat_tool_to_responses_tool(tool: dict[str, Any]) -> dict[str, Any]:
    """Convert a Chat Completions function tool to Responses API format."""
    if tool.get("type") == "function" and isinstance(tool.get("function"), dict):
        fn = tool["function"]
        return {
            "type": "function",
            "name": fn.get("name", ""),
            "description": fn.get("description", ""),
            "parameters": fn.get("parameters", {"type": "object", "properties": {}}),
        }
    return tool


def _messages_to_responses_input(messages: list[dict[str, Any]]) -> tuple[list[dict[str, Any]], str]:
    """Convert chat-style messages, including tool results, to Responses input."""
    instructions: list[str] = []
    input_items: list[dict[str, Any]] = []
    for message in messages:
        role = message.get("role")
        content = message.get("content") or ""
        if role == "system":
            if content:
                instructions.append(str(content))
            continue
        if role == "tool":
            input_items.append({
                "type": "function_call_output",
                "call_id": str(message.get("tool_call_id", "")),
                "output": str(content),
            })
            continue
        if role == "assistant":
            if content:
                input_items.append({"role": "assistant", "content": str(content)})
            for tool_call in message.get("tool_calls") or []:
                function = tool_call.get("function", {}) or {}
                input_items.append({
                    "type": "function_call",
                    "call_id": str(tool_call.get("id", "")),
                    "name": str(function.get("name", "")),
                    "arguments": str(function.get("arguments", "{}") or "{}"),
                })
            continue
        if role in {"user", "developer"}:
            input_items.append({"role": "user", "content": str(content)})
    return input_items, "\n\n".join(instructions)


def _responses_to_chat_message(resp: Any) -> tuple[Any, str]:
    """Convert Responses output into the subset of Chat message API we use."""
    text = getattr(resp, "output_text", None) or ""
    tool_calls: list[dict[str, Any]] = []
    for item in getattr(resp, "output", None) or []:
        item_type = getattr(item, "type", "")
        if item_type == "function_call":
            tool_calls.append({
                "id": getattr(item, "call_id", "") or getattr(item, "id", ""),
                "type": "function",
                "function": {
                    "name": getattr(item, "name", ""),
                    "arguments": getattr(item, "arguments", "") or "{}",
                },
            })
        elif item_type == "message" and not text:
            content_parts = getattr(item, "content", []) or []
            for part in content_parts:
                if getattr(part, "type", "") == "output_text":
                    text += getattr(part, "text", "") or ""
    return SimpleNamespace(content=text, tool_calls=tool_calls), text


# ── Public API ────────────────────────────────────────────────────────────────

def configure_azure_openai(
    *,
    endpoint: str | None = None,
    api_version: str | None = None,
    api_key: str | None = None,
    auth_mode: str | None = None,
    ad_scope: str | None = None,
    managed_identity_client_id: str | None = None,
    optimizer_endpoint: str | None = None,
    optimizer_api_version: str | None = None,
    optimizer_api_key: str | None = None,
    optimizer_auth_mode: str | None = None,
    optimizer_ad_scope: str | None = None,
    optimizer_managed_identity_client_id: str | None = None,
    target_endpoint: str | None = None,
    target_api_version: str | None = None,
    target_api_key: str | None = None,
    target_auth_mode: str | None = None,
    target_ad_scope: str | None = None,
    target_managed_identity_client_id: str | None = None,
) -> None:
    global ENDPOINT, API_VERSION, API_KEY, AUTH_MODE, AD_SCOPE, MANAGED_IDENTITY_CLIENT_ID
    global OPTIMIZER_ENDPOINT, OPTIMIZER_API_VERSION, OPTIMIZER_API_KEY, OPTIMIZER_AUTH_MODE
    global OPTIMIZER_AD_SCOPE, OPTIMIZER_MANAGED_IDENTITY_CLIENT_ID
    global TARGET_ENDPOINT, TARGET_API_VERSION, TARGET_API_KEY, TARGET_AUTH_MODE
    global TARGET_AD_SCOPE, TARGET_MANAGED_IDENTITY_CLIENT_ID
    global _optimizer_client, _target_client

    def _clean(value: str | None, *, lower: bool = False) -> str | None:
        if value is None:
            return None
        str_value = str(value).strip()
        if not str_value:
            return None
        if lower:
            str_value = str_value.lower()
        return str_value

    def _set(global_name: str, value: str | None, env_key: str) -> None:
        if value is None:
            return
        globals()[global_name] = value
        os.environ[env_key] = value

    shared_endpoint = _clean(endpoint)
    shared_api_version = _clean(api_version)
    shared_api_key = _clean(api_key)
    shared_auth_mode = _clean(auth_mode, lower=True)
    shared_ad_scope = _clean(ad_scope)
    shared_managed_identity_client_id = _clean(managed_identity_client_id)

    # Auto-configure for openai_compatible mode
    if shared_auth_mode in {"openai_compatible", "compat", "openai"}:
        if shared_api_version is None:
            shared_api_version = _OPENAI_COMPATIBLE_API_VERSION

    _set("ENDPOINT", shared_endpoint, "AZURE_OPENAI_ENDPOINT")
    _set("API_VERSION", shared_api_version, "AZURE_OPENAI_API_VERSION")
    _set("API_KEY", shared_api_key, "AZURE_OPENAI_API_KEY")
    _set("AUTH_MODE", shared_auth_mode, "AZURE_OPENAI_AUTH_MODE")
    _set("AD_SCOPE", shared_ad_scope, "AZURE_OPENAI_AD_SCOPE")
    _set(
        "MANAGED_IDENTITY_CLIENT_ID",
        shared_managed_identity_client_id,
        "AZURE_OPENAI_MANAGED_IDENTITY_CLIENT_ID",
    )

    resolved_optimizer_endpoint = _clean(optimizer_endpoint) or shared_endpoint
    resolved_optimizer_api_version = _clean(optimizer_api_version) or shared_api_version
    resolved_optimizer_api_key = _clean(optimizer_api_key) or shared_api_key
    resolved_optimizer_auth_mode = _clean(optimizer_auth_mode, lower=True) or shared_auth_mode
    resolved_optimizer_ad_scope = _clean(optimizer_ad_scope) or shared_ad_scope
    resolved_optimizer_mi = (
        _clean(optimizer_managed_identity_client_id)
        or shared_managed_identity_client_id
    )
    
    # Auto-configure for openai_compatible mode
    if resolved_optimizer_auth_mode in {"openai_compatible", "compat", "openai"}:
        if resolved_optimizer_api_version is None:
            resolved_optimizer_api_version = _OPENAI_COMPATIBLE_API_VERSION
    
    resolved_target_endpoint = _clean(target_endpoint) or shared_endpoint
    resolved_target_api_version = _clean(target_api_version) or shared_api_version
    resolved_target_api_key = _clean(target_api_key) or shared_api_key
    resolved_target_auth_mode = _clean(target_auth_mode, lower=True) or shared_auth_mode
    resolved_target_ad_scope = _clean(target_ad_scope) or shared_ad_scope
    resolved_target_mi = (
        _clean(target_managed_identity_client_id)
        or shared_managed_identity_client_id
    )
    
    # Auto-configure for openai_compatible mode
    if resolved_target_auth_mode in {"openai_compatible", "compat", "openai"}:
        if resolved_target_api_version is None:
            resolved_target_api_version = _OPENAI_COMPATIBLE_API_VERSION

    _set("OPTIMIZER_ENDPOINT", resolved_optimizer_endpoint, "OPTIMIZER_AZURE_OPENAI_ENDPOINT")
    _set(
        "OPTIMIZER_API_VERSION",
        resolved_optimizer_api_version,
        "OPTIMIZER_AZURE_OPENAI_API_VERSION",
    )
    _set("OPTIMIZER_API_KEY", resolved_optimizer_api_key, "OPTIMIZER_AZURE_OPENAI_API_KEY")
    _set("OPTIMIZER_AUTH_MODE", resolved_optimizer_auth_mode, "OPTIMIZER_AZURE_OPENAI_AUTH_MODE")
    _set("OPTIMIZER_AD_SCOPE", resolved_optimizer_ad_scope, "OPTIMIZER_AZURE_OPENAI_AD_SCOPE")
    _set(
        "OPTIMIZER_MANAGED_IDENTITY_CLIENT_ID",
        resolved_optimizer_mi,
        "OPTIMIZER_AZURE_OPENAI_MANAGED_IDENTITY_CLIENT_ID",
    )
    _set("TARGET_ENDPOINT", resolved_target_endpoint, "TARGET_AZURE_OPENAI_ENDPOINT")
    _set(
        "TARGET_API_VERSION",
        resolved_target_api_version,
        "TARGET_AZURE_OPENAI_API_VERSION",
    )
    _set("TARGET_API_KEY", resolved_target_api_key, "TARGET_AZURE_OPENAI_API_KEY")
    _set("TARGET_AUTH_MODE", resolved_target_auth_mode, "TARGET_AZURE_OPENAI_AUTH_MODE")
    _set("TARGET_AD_SCOPE", resolved_target_ad_scope, "TARGET_AZURE_OPENAI_AD_SCOPE")
    _set(
        "TARGET_MANAGED_IDENTITY_CLIENT_ID",
        resolved_target_mi,
        "TARGET_AZURE_OPENAI_MANAGED_IDENTITY_CLIENT_ID",
    )

    with _optimizer_lock:
        _optimizer_client = None
    with _target_lock:
        _target_client = None


def chat_optimizer(
    system: str,
    user: str,
    max_completion_tokens: int = 16384,
    retries: int = 5,
    stage: str = "optimizer",
    reasoning_effort: str | None = None,
    timeout: int | None = None,
) -> tuple[str, dict]:
    """Call the optimizer model.  Returns (response_text, usage_dict)."""
    return _chat_impl(
        get_optimizer_client(), OPTIMIZER_DEPLOYMENT,
        system, user, max_completion_tokens, retries, stage, reasoning_effort, timeout,
    )


def chat_with_deployment(
    deployment: str,
    system: str,
    user: str,
    max_completion_tokens: int = 16384,
    retries: int = 5,
    stage: str = "custom",
    reasoning_effort: str | None = None,
    timeout: int | None = None,
) -> tuple[str, dict]:
    """Call an arbitrary deployment using the shared Azure client."""
    return _chat_impl(
        get_optimizer_client(),
        deployment,
        system,
        user,
        max_completion_tokens,
        retries,
        stage,
        reasoning_effort,
        timeout,
    )


def chat_target(
    system: str,
    user: str,
    max_completion_tokens: int = 16384,
    retries: int = 5,
    stage: str = "target",
    reasoning_effort: str | None = None,
    timeout: int | None = None,
) -> tuple[str, dict]:
    """Call the target model.  Returns (response_text, usage_dict)."""
    return _chat_impl(
        get_target_client(), TARGET_DEPLOYMENT,
        system, user, max_completion_tokens, retries, stage, reasoning_effort, timeout,
    )


def chat_optimizer_messages(
    messages: list[dict[str, Any]],
    max_completion_tokens: int = 16384,
    retries: int = 5,
    stage: str = "optimizer",
    reasoning_effort: str | None = None,
    *,
    tools: list[dict[str, Any]] | None = None,
    tool_choice: str | dict[str, Any] | None = None,
    return_message: bool = False,
    timeout: int | None = None,
) -> tuple[Any, dict]:
    """Call the optimizer model with a pre-built chat message list."""
    return _chat_messages_impl(
        get_optimizer_client(),
        OPTIMIZER_DEPLOYMENT,
        messages,
        max_completion_tokens,
        retries,
        stage,
        reasoning_effort,
        tools=tools,
        tool_choice=tool_choice,
        return_message=return_message,
        timeout=timeout,
    )


def chat_messages_with_deployment(
    deployment: str,
    messages: list[dict[str, Any]],
    max_completion_tokens: int = 16384,
    retries: int = 5,
    stage: str = "custom",
    reasoning_effort: str | None = None,
    *,
    tools: list[dict[str, Any]] | None = None,
    tool_choice: str | dict[str, Any] | None = None,
    return_message: bool = False,
    timeout: int | None = None,
) -> tuple[Any, dict]:
    """Call an arbitrary deployment with a pre-built chat message list."""
    return _chat_messages_impl(
        get_optimizer_client(),
        deployment,
        messages,
        max_completion_tokens,
        retries,
        stage,
        reasoning_effort,
        tools=tools,
        tool_choice=tool_choice,
        return_message=return_message,
        timeout=timeout,
    )


def chat_target_messages(
    messages: list[dict[str, Any]],
    max_completion_tokens: int = 16384,
    retries: int = 5,
    stage: str = "target",
    reasoning_effort: str | None = None,
    *,
    tools: list[dict[str, Any]] | None = None,
    tool_choice: str | dict[str, Any] | None = None,
    return_message: bool = False,
    timeout: int | None = None,
) -> tuple[Any, dict]:
    """Call the target model with a pre-built chat message list."""
    return _chat_messages_impl(
        get_target_client(),
        TARGET_DEPLOYMENT,
        messages,
        max_completion_tokens,
        retries,
        stage,
        reasoning_effort,
        tools=tools,
        tool_choice=tool_choice,
        return_message=return_message,
        timeout=timeout,
    )


def get_token_summary() -> dict:
    """Return per-stage and total token usage."""
    return tracker.summary()


def reset_token_tracker() -> None:
    tracker.reset()


def set_target_deployment(deployment: str) -> None:
    """Change target deployment at runtime."""
    global _target_client, TARGET_DEPLOYMENT
    TARGET_DEPLOYMENT = deployment
    os.environ["TARGET_DEPLOYMENT"] = deployment
    os.environ["AZURE_OPENAI_DEPLOYMENT"] = deployment
    with _target_lock:
        _target_client = None
    try:
        import llm_client as _legacy
        _legacy.DEPLOYMENT = deployment
        _legacy._client = None
    except Exception:
        pass


def set_reasoning_effort(effort: str | None) -> None:
    """Set reasoning effort for all LLM calls. None = off."""
    global REASONING_EFFORT
    REASONING_EFFORT = effort if effort else None


def get_reasoning_effort() -> str | None:
    """Return the process-wide reasoning effort for direct Azure client users."""
    return REASONING_EFFORT


def set_optimizer_deployment(deployment: str) -> None:
    """Change optimizer deployment at runtime."""
    global _optimizer_client, OPTIMIZER_DEPLOYMENT
    OPTIMIZER_DEPLOYMENT = deployment
    os.environ["OPTIMIZER_DEPLOYMENT"] = deployment
    with _optimizer_lock:
        _optimizer_client = None
