"""SDK JSON config loader and converter (grouped sections only)."""

from __future__ import annotations

import json
import os
from pathlib import Path
from typing import Any


def _as_dict(value: Any) -> dict[str, Any]:
    return dict(value) if isinstance(value, dict) else {}


def _pick(src: dict[str, Any], *keys: str) -> Any:
    for key in keys:
        if key in src and src[key] is not None:
            return src[key]
    return None


def _normalize_picasso_config(raw: dict[str, Any]) -> dict[str, Any]:
    return {
        "branch": _pick(raw, "branch", "picasso_branch"),
        "liquid_prompt_entries": _pick(raw, "liquid_prompt_entries"),
        "add_feature_flags": _pick(raw, "add_feature_flags"),
        "baseline_feature_flags": _pick(raw, "baseline_feature_flags"),
        "picasso_version": _pick(raw, "picasso_version"),
        "run_leaderboard_after_optimization": _pick(raw, "run_leaderboard_after_optimization"),
    }


def _normalize_aipa_config(raw: dict[str, Any]) -> dict[str, Any]:
    return {
        "dataset_name": _pick(raw, "dataset_name"),
        "dataset_project_name": _pick(raw, "dataset_project_name"),
        "max_rows": _pick(raw, "max_rows"),
        "dataset_sampling_method": _pick(raw, "dataset_sampling_method"),
        "additional_cmd_args": _pick(raw, "additional_cmd_args"),
        "eval_name": _pick(raw, "eval_name"),
        "eval_project_name": _pick(raw, "eval_project_name"),
        "conversation_mode": _pick(raw, "conversation_mode"),
        "environment": _pick(raw, "environment"),
        "model": _pick(raw, "model"),
        "picasso_env": _pick(raw, "picasso_env"),
        "task_id_prefix": _pick(raw, "task_id_prefix"),
    }


def _normalize_braintrust_config(raw: dict[str, Any]) -> dict[str, Any]:
    return {
        "org": _pick(raw, "org", "braintrust_org"),
        "api_key": _pick(raw, "api_key", "braintrust_api_key"),
        "api_key_env": _pick(raw, "api_key_env"),
        "metrics_config": _pick(raw, "metrics_config"),
        "dimension_threshold": _pick(raw, "dimension_threshold"),
        "bad_case_limit_per_iteration": _pick(raw, "bad_case_limit_per_iteration"),
        "badcase_deep_analysis": _pick(raw, "badcase_deep_analysis"),
    }


def _normalize_model_config(raw: dict[str, Any]) -> dict[str, Any]:
    return {
        "type": _pick(raw, "type", "model_client_type"),
        "azure_endpoint": _pick(raw, "azure_endpoint"),
        "deployment": _pick(raw, "deployment", "model_deployment"),
        "azure_api_version": _pick(raw, "azure_api_version"),
        "papyrus_endpoint": _pick(raw, "papyrus_endpoint"),
        "papyrus_model": _pick(raw, "papyrus_model"),
        "papyrus_quota_id": _pick(raw, "papyrus_quota_id"),
        "papyrus_scope": _pick(raw, "papyrus_scope"),
        "local_endpoint": _pick(raw, "local_endpoint"),
        "local_model": _pick(raw, "local_model"),
        "use_structured_output": _pick(raw, "use_structured_output"),
    }


def _normalize_runtime_config(raw: dict[str, Any]) -> dict[str, Any]:
    return {
        "session_dir": _pick(raw, "session_dir"),
        "local_debug": _pick(raw, "local_debug"),
        "enable_run_logging": _pick(raw, "enable_run_logging"),
    }


def _normalize_optimize_config(raw: dict[str, Any]) -> dict[str, Any]:
    return {
        "max_rounds": _pick(raw, "max_rounds"),
        "termination_rule": _pick(raw, "termination_rule"),
        "termination_no_progress_rounds": _pick(raw, "termination_no_progress_rounds"),
        "rollback_strategy": _pick(raw, "rollback_strategy"),
        "baseline_experiment_id": _pick(raw, "baseline_experiment_id"),
        "baseline_experiment_link": _pick(raw, "baseline_experiment_link"),
        "selection_strategy": _pick(raw, "selection_strategy"),
        "pareto_max_selected": _pick(raw, "pareto_max_selected"),
    }


def _resolve_braintrust_api_key(
    braintrust_config: dict[str, Any],
    raw: dict[str, Any],
    *,
    load_env_secrets: bool,
) -> str:
    existing = _pick(braintrust_config, "api_key")
    if isinstance(existing, str) and existing.strip():
        return existing.strip()

    api_key_env = _pick(raw, "braintrust_api_key_env")
    if not api_key_env:
        api_key_env = _pick(braintrust_config, "api_key_env")
    if isinstance(api_key_env, str) and api_key_env.strip():
        return os.getenv(api_key_env.strip(), "")

    if not load_env_secrets:
        return ""

    org = str(
        _pick(braintrust_config, "org")
        or "non-pme"
    ).strip().lower()
    if org == "bt-azure":
        return os.getenv("BRAINTRUST_AZURE_API_KEY", "")
    return os.getenv("BRAINTRUST_API_KEY", "")


def load_create_kwargs_from_file(
    config_path: str | Path,
    *,
    load_env_secrets: bool = True,
) -> dict[str, Any]:
    """Load SDK config file and return grouped section configs.

    Return shape:
    - picasso_config
    - aipa_config
    - braintrust_config
    - model_config
    - runtime_config
    - optimize_config
    """
    path = Path(config_path)
    content = path.read_text(encoding="utf-8-sig")
    raw = json.loads(content)
    if not isinstance(raw, dict):
        raise ValueError(f"Config file must be a JSON object: {path}")

    picasso_config = _normalize_picasso_config(_as_dict(raw.get("picasso_config")))
    aipa_config = _normalize_aipa_config(_as_dict(raw.get("aipa_config")))
    braintrust_config = _normalize_braintrust_config(_as_dict(raw.get("braintrust_config")))
    model_config = _normalize_model_config(_as_dict(raw.get("model_config")))
    runtime_config = _normalize_runtime_config(_as_dict(raw.get("runtime_config")))
    optimize_config = _normalize_optimize_config(_as_dict(raw.get("optimize_config")))

    resolved_key = _resolve_braintrust_api_key(
        braintrust_config,
        raw,
        load_env_secrets=load_env_secrets,
    )
    if resolved_key:
        braintrust_config["api_key"] = resolved_key

    return {
        "picasso_config": picasso_config,
        "aipa_config": aipa_config,
        "braintrust_config": braintrust_config,
        "model_config": model_config,
        "runtime_config": runtime_config,
        "optimize_config": optimize_config,
    }


def build_create_kwargs(config_bundle: dict[str, Any]) -> dict[str, Any]:
    """Convert grouped config sections into AutoPromptLab.create kwargs."""
    picasso = _as_dict(config_bundle.get("picasso_config"))
    aipa = _as_dict(config_bundle.get("aipa_config"))
    braintrust = _as_dict(config_bundle.get("braintrust_config"))
    model = _as_dict(config_bundle.get("model_config"))
    runtime = _as_dict(config_bundle.get("runtime_config"))
    optimize = _as_dict(config_bundle.get("optimize_config"))
    values: dict[str, Any] = {}

    if _pick(picasso, "branch") is not None:
        values["picasso_branch"] = _pick(picasso, "branch")
    if _pick(picasso, "liquid_prompt_entries") is not None:
        values["liquid_prompt_entries"] = _pick(picasso, "liquid_prompt_entries")
    if _pick(picasso, "add_feature_flags") is not None:
        values["add_feature_flags"] = _pick(picasso, "add_feature_flags")
    if _pick(picasso, "baseline_feature_flags") is not None:
        values["baseline_feature_flags"] = _pick(picasso, "baseline_feature_flags")
    if _pick(picasso, "picasso_version") is not None:
        values["picasso_version"] = _pick(picasso, "picasso_version")
    if _pick(picasso, "run_leaderboard_after_optimization") is not None:
        values["run_leaderboard_after_optimization"] = _pick(picasso, "run_leaderboard_after_optimization")

    # environment and model can be in picasso_config or aipa_config; picasso takes precedence
    if _pick(picasso, "environment") is not None:
        values["environment"] = _pick(picasso, "environment")
    if _pick(picasso, "model") is not None:
        values["model"] = _pick(picasso, "model")

    if _pick(aipa, "dataset_name") is not None:
        values["dataset_name"] = _pick(aipa, "dataset_name")
    if _pick(aipa, "dataset_project_name") is not None:
        values["dataset_project_name"] = _pick(aipa, "dataset_project_name")
    if _pick(aipa, "max_rows") is not None:
        values["max_dataset_rows"] = _pick(aipa, "max_rows")
    if _pick(aipa, "eval_name") is not None:
        values["eval_name"] = _pick(aipa, "eval_name")
    if _pick(aipa, "eval_project_name") is not None:
        values["eval_project_name"] = _pick(aipa, "eval_project_name")
    if _pick(aipa, "conversation_mode") is not None:
        values["conversation_mode"] = _pick(aipa, "conversation_mode")
    environment = _pick(aipa, "environment")
    if environment is not None:
        values["environment"] = environment

    model_value = _pick(aipa, "model")
    if model_value is not None:
        values["model"] = model_value
    if _pick(aipa, "dataset_sampling_method") is not None:
        values["dataset_sampling_method"] = _pick(aipa, "dataset_sampling_method")
    if _pick(aipa, "additional_cmd_args") is not None:
        values["additional_cmd_args"] = _pick(aipa, "additional_cmd_args")
    if _pick(aipa, "picasso_env") is not None:
        values["picasso_env"] = _pick(aipa, "picasso_env")
    if _pick(aipa, "task_id_prefix") is not None:
        values["task_id_prefix"] = _pick(aipa, "task_id_prefix")

    if _pick(braintrust, "org") is not None:
        values["braintrust_org"] = _pick(braintrust, "org")
    if _pick(braintrust, "api_key") is not None:
        values["braintrust_api_key"] = _pick(braintrust, "api_key")
    if _pick(braintrust, "metrics_config") is not None:
        values["metrics_config"] = _pick(braintrust, "metrics_config")
    if _pick(braintrust, "dimension_threshold") is not None:
        values["dimension_threshold"] = _pick(braintrust, "dimension_threshold")
    if _pick(braintrust, "bad_case_limit_per_iteration") is not None:
        values["bad_case_limit_per_iteration"] = _pick(braintrust, "bad_case_limit_per_iteration")
    if _pick(braintrust, "badcase_deep_analysis") is not None:
        values["badcase_deep_analysis"] = _pick(braintrust, "badcase_deep_analysis")

    if _pick(model, "type") is not None:
        values["model_client_type"] = _pick(model, "type")
    if _pick(model, "azure_endpoint") is not None:
        values["azure_endpoint"] = _pick(model, "azure_endpoint")
    if _pick(model, "deployment") is not None:
        values["model_deployment"] = _pick(model, "deployment")
    if _pick(model, "azure_api_version") is not None:
        values["azure_api_version"] = _pick(model, "azure_api_version")
    if _pick(model, "papyrus_endpoint") is not None:
        values["papyrus_endpoint"] = _pick(model, "papyrus_endpoint")
    if _pick(model, "papyrus_model") is not None:
        values["papyrus_model"] = _pick(model, "papyrus_model")
    if _pick(model, "papyrus_quota_id") is not None:
        values["papyrus_quota_id"] = _pick(model, "papyrus_quota_id")
    if _pick(model, "papyrus_scope") is not None:
        values["papyrus_scope"] = _pick(model, "papyrus_scope")
    if _pick(model, "local_endpoint") is not None:
        values["local_endpoint"] = _pick(model, "local_endpoint")
    if _pick(model, "local_model") is not None:
        values["local_model"] = _pick(model, "local_model")
    if _pick(model, "use_structured_output") is not None:
        values["use_structured_output"] = _pick(model, "use_structured_output")

    if _pick(runtime, "session_dir") is not None:
        values["session_dir"] = _pick(runtime, "session_dir")
    if _pick(runtime, "local_debug") is not None:
        values["local_debug"] = _pick(runtime, "local_debug")
    if _pick(runtime, "enable_run_logging") is not None:
        values["enable_run_logging"] = _pick(runtime, "enable_run_logging")

    if _pick(optimize, "max_rounds") is not None:
        values["max_rounds"] = _pick(optimize, "max_rounds")
    if _pick(optimize, "termination_rule") is not None:
        values["termination_rule"] = _pick(optimize, "termination_rule")
    if _pick(optimize, "rollback_strategy") is not None:
        values["rollback_strategy"] = _pick(optimize, "rollback_strategy")
    if _pick(optimize, "baseline_experiment_id") is not None:
        values["baseline_experiment_id"] = _pick(optimize, "baseline_experiment_id")
    if _pick(optimize, "baseline_experiment_link") is not None:
        values["baseline_experiment_link"] = _pick(optimize, "baseline_experiment_link")
    if _pick(optimize, "selection_strategy") is not None:
        values["selection_strategy"] = _pick(optimize, "selection_strategy")
    if _pick(optimize, "pareto_max_selected") is not None:
        values["pareto_max_selected"] = _pick(optimize, "pareto_max_selected")

    return values
