# agent_loop_config.py - LLM / backend configuration helpers extracted from agent_loop.
from __future__ import annotations

import logging
import re
from typing import NamedTuple

from persistence import MUGI_DIR

_log = logging.getLogger(__name__)
from utils.app_prefs_store import load_app_prefs_cached


# Chats database: JSON file for tab/conversation list; SQLite for message checkpoints (same dir as Mugi config)
CHATS_DB_DIR = MUGI_DIR
APP_PREFS_FILENAME = "app_prefs.json"

# Default backend URL (LM Studio / OpenAI-compatible API). Overridden by app_prefs.json when present.
DEFAULT_BACKEND_BASE_URL = "http://127.0.0.1:1234/v1"

_pref_load_warned = False


def _normalize_backend_url(url: str) -> str:
    """Prepend http:// if the URL has no scheme, so e.g. 10.0.1.111 becomes http://10.0.1.111."""
    u = (url or "").strip()
    if not u:
        return u
    lower = u.lower()
    if lower.startswith("http://") or lower.startswith("https://"):
        return u
    return "http://" + u


def _local_mlx_openai_autostart_eligible(base_url: str) -> bool:
    """True when ``base_url`` targets the conventional local mlx-openai-server port on loopback.

    Autostart must not run for remote ``…:8000`` hosts — that would start a *local* server while the
    client expects a peer on the LAN.
    """
    u = (base_url or "").strip().lower()
    if ":8000" not in u:
        return False
    try:
        from urllib.parse import urlparse

        host = (urlparse(_normalize_backend_url(base_url)).hostname or "").lower()
    except Exception:
        host = ""
    return host in ("127.0.0.1", "localhost", "::1")


class MLXReadyResult(NamedTuple):
    ok: bool
    reason: str


def get_backend_base_url():
    """Return backend API base URL from ~/.mugi/app_prefs.json, or DEFAULT_BACKEND_BASE_URL."""
    try:
        prefs = load_app_prefs_cached()
        ext = prefs.get("use_external_backend") is True
        url = (prefs.get("backend_base_url") or "").strip()
        use_llmster = prefs.get("use_local_ollama") is True
        if use_llmster and not ext:
            norm = _normalize_backend_url(url) if url else ""
            if not norm or ":8000" in norm:
                if url and ":8000" in _normalize_backend_url(url):
                    _log.warning(
                        "backend_base_url %r targets MLX :8000 but use_local_ollama remaps to %s",
                        url,
                        DEFAULT_BACKEND_BASE_URL,
                    )
                return DEFAULT_BACKEND_BASE_URL
        if url:
            return _normalize_backend_url(url)
        if ext:
            # External mode with no URL yet — avoid falling back to local defaults.
            return ""
    except Exception:
        _log.error("Failed to load app prefs for backend_base_url; using default", exc_info=True)
    return DEFAULT_BACKEND_BASE_URL


def _ensure_local_mlx_server_ready(base_url: str, on_log=None) -> MLXReadyResult:
    """Best-effort local MLX autostart when MLX is enabled and the API URL is loopback :8000.

    ``mlx_enabled`` is the product switch: when false, Mugi's agent must not spawn mlx-openai-server.
    If ``backend_base_url`` still points at ``http://127.0.0.1:8000/v1`` with MLX off, that is an
    inconsistent preference state; fix it in Settings (point at LM Studio / another backend) or
    turn MLX back on — do not autostart here.
    """
    if not _local_mlx_openai_autostart_eligible(base_url):
        result = MLXReadyResult(False, "not_eligible")
        _log.debug("MLX autostart skipped: %s", result.reason)
        return result
    try:
        if not bool(_read_pref("mlx_enabled", False)):
            result = MLXReadyResult(False, "mlx_disabled")
            _log.debug("MLX autostart skipped: %s", result.reason)
            return result
        model_path = str(_read_pref("mlx_model", "") or "").strip()
        if not model_path:
            result = MLXReadyResult(False, "no_model")
            _log.debug("MLX autostart skipped: %s", result.reason)
            return result
        import app_mlx as _am

        mlx_venv_exists_fn = getattr(_am, "mlx_venv_exists", None)
        mlx_server_running_fn = getattr(_am, "mlx_server_running", None)
        start_mlx_server_fn = getattr(_am, "start_mlx_server", None)
        if not callable(mlx_venv_exists_fn) or not callable(mlx_server_running_fn) or not callable(start_mlx_server_fn):
            result = MLXReadyResult(False, "missing_hooks")
            _log.debug("MLX autostart skipped: %s", result.reason)
            return result
        stop_mlx_server_fn = getattr(_am, "stop_mlx_server", lambda: None)
        build_fp_fn = getattr(_am, "build_mlx_runtime_fingerprint", None)
        get_last_fp_fn = getattr(_am, "get_last_mlx_runtime_fingerprint", None)
        fp_match_fn = getattr(_am, "runtime_fingerprint_matches", None)
        if not mlx_venv_exists_fn():
            result = MLXReadyResult(False, "no_venv")
            _log.debug("MLX autostart skipped: %s", result.reason)
            return result
        apply_patches_fn = getattr(_am, "mlx_prepare_venv_runtime", None)
        runtime_needs_restart = bool(apply_patches_fn()) if callable(apply_patches_fn) else False
        context_length = get_llm_context_length()
        intended_fp = (
            build_fp_fn(model_path, context_length=context_length)
            if callable(build_fp_fn)
            else {}
        )
        _server_up = mlx_server_running_fn()
        did_restart = False
        _log.debug("MLX autostart: server_up=%s restart=%s model=%s ctx=%d",
                    _server_up, runtime_needs_restart, model_path[-80:], context_length)
        if _server_up:
            if runtime_needs_restart:
                if on_log:
                    on_log(
                        "AGENT",
                        "MLX runtime files were patched; restarting local server so chat uses the fixes…",
                    )
                stop_mlx_server_fn()
                did_restart = True
            else:
                fingerprint_ok = False
                running_fp = {}
                if callable(build_fp_fn) and callable(get_last_fp_fn) and callable(fp_match_fn):
                    running_fp = get_last_fp_fn()
                    fingerprint_ok = fp_match_fn(intended_fp, running_fp) if running_fp else False
                _log.debug("MLX fingerprint: ok=%s intended=%s running=%s",
                            fingerprint_ok,
                            str(intended_fp.get("model_path", ""))[-80:],
                            str((running_fp or {}).get("model_path", ""))[-80:])
                if fingerprint_ok:
                    return MLXReadyResult(True, "already_running")
                if on_log:
                    on_log("WARN", "Local MLX server is running with a stale model/parser; restarting before request...")
                stop_mlx_server_fn()
                did_restart = True
        if on_log:
            on_log("AGENT", "Starting local MLX server for first request...")
        ok, msg = start_mlx_server_fn(model_path, context_length=context_length)
        if ok and did_restart:
            reprobe_fn = getattr(_am, "_run_and_persist_mlx_native_tool_probe", None)
            if callable(reprobe_fn):
                reprobe_fn()
            else:
                clear_probe_fn = getattr(_am, "clear_mlx_last_tool_probe", None)
                if callable(clear_probe_fn):
                    clear_probe_fn()
        if on_log:
            level = "AGENT" if ok else "WARN"
            on_log(level, f"Local MLX server {'ready' if ok else 'failed to start'}: {msg}")
        if ok:
            return MLXReadyResult(True, "restarted" if did_restart else "started")
        result = MLXReadyResult(False, "start_failed")
        _log.debug("MLX autostart failed: %s (%s)", result.reason, msg)
        return result
    except Exception as e:
        if on_log:
            on_log("WARN", f"Local MLX autostart skipped: {e}")
        result = MLXReadyResult(False, "exception")
        _log.debug("MLX autostart skipped: %s (%s)", result.reason, e)
        return result


def get_llmster_model():
    """Return selected llmster (LM Studio) model from app_prefs; default empty (server may use first loaded)."""
    try:
        prefs = load_app_prefs_cached()
        return (prefs.get("selected_llmster_model") or "").strip()
    except Exception:
        return ""


def _get_mlx_active_adapter():
    """Return active MLX LoRA adapter name from app_prefs, or empty string."""
    try:
        prefs = load_app_prefs_cached()
        return (prefs.get("mlx_active_adapter") or "").strip()
    except Exception:
        return ""


_GEMMA4_PARAM_RE = re.compile(r"gemma[_-]?4[_-](\d+)b", re.IGNORECASE)


def mlx_active_adapter_matches_openai_model_id(adapter_name: str, openai_model_id: str) -> bool:
    """Return False when prefs clearly wire a Gemma-4 LoRA from one backbone size onto another.

    The mlx-vlm server accepts a per-request ``extra_body.adapter_path``; a mismatched
    evolution/LoRA vs the loaded model often **hangs or drops the HTTP stream** (no chunks
    until idle timeout) instead of returning a clean 4xx, so we gate here before every chat request.
    """
    a = (adapter_name or "").strip()
    m = (openai_model_id or "").strip()
    if not a:
        return True
    ma = _GEMMA4_PARAM_RE.search(a)
    mm = _GEMMA4_PARAM_RE.search(m)
    if ma and mm and ma.group(1) != mm.group(1):
        _log.warning(
            "LoRA adapter %r targets a different Gemma-4 backbone size than model %r — skipping adapter",
            a,
            m,
        )
        return False
    return True


_MIN_CONTEXT_LENGTH = 16384  # ~16k minimum across all models/engines
_DEFAULT_CONTEXT_LENGTH = 32768
_MAX_CONTEXT_LENGTH = 1_048_576  # 1M for models that support it
_DEFAULT_TEMPERATURE = 0.7
_DEFAULT_MAX_TOKENS = 8192
_DEFAULT_RESEARCH_MAX_SEARCH_ROUNDS = 8
_DEFAULT_RESEARCH_MAX_READ_URLS = 8
_DEFAULT_RESEARCH_MIN_SCRATCH_APPENDS = 2
_DEFAULT_RESEARCH_MIN_SCRATCH_FILE_SIZE = 2000
# Doctrine-pure progress forcing: after this many consecutive research rounds that gather
# (plan/search/read) without producing a deliverable (append_research_scratch or
# generate_final_research_report), force the worker to CLOSE from gathered evidence. The
# harness never synthesizes on the model's behalf — it forces the *model* to write its
# findings or emit an honest blocked handoff. Capped attempts keep an honest give-up reachable.
_DEFAULT_RESEARCH_FORCE_CLOSE_NONPROGRESS_ROUNDS = 4
_DEFAULT_RESEARCH_FORCE_CLOSE_MAX_ATTEMPTS = 2
_DEFAULT_MEMORY_RETRIEVAL_PROFILE = "conservative"
_DEFAULT_MEMORY_SEMANTIC_STRICTNESS = "warn_promotion"
# Defaults for get_agent_loop_limits(); also used by mugi_backend GET /preferences.
DEFAULT_AGENT_MAX_TOOL_ROUNDS = 35
DEFAULT_AGENT_GUARD_RETRIES_TOTAL = 8
DEFAULT_TASK_PLAN_GUARD_RETRIES_TOTAL = 6
DEFAULT_TASK_PLAN_CONTINUATION_RETRIES = 3


def _read_pref(key, default=None):
    global _pref_load_warned
    try:
        prefs = load_app_prefs_cached()
        return prefs.get(key, default)
    except Exception:
        if not _pref_load_warned:
            _pref_load_warned = True
            _log.warning(
                "Failed to load app prefs; using defaults for agent loop settings",
                exc_info=True,
            )
        return default


def _int_pref(key: str, default: int, lo: int, hi: int) -> int:
    val = _read_pref(key, default)
    try:
        return max(lo, min(int(val), hi))
    except Exception:
        return default


def _float_pref(key: str, default: float, lo: float, hi: float) -> float:
    val = _read_pref(key, default)
    try:
        return max(lo, min(float(val), hi))
    except Exception:
        return default


def get_mlx_openai_model_id() -> str:
    """Return the OpenAI API ``model`` id for mlx-openai-server.

    LM Studio accepts the placeholder ``local-model``; MLX registers whatever path was passed to
    ``mlx-openai-server launch --model-path`` (typically ``realpath``). Empty string if no
    ``mlx_model`` is configured in app prefs.
    """
    try:
        from app_mlx import _normalize_mlx_model_identity

        path = str(_read_pref("mlx_model", "") or "").strip()
        return (_normalize_mlx_model_identity(path) or "").strip()
    except Exception:
        return ""


def probe_openai_compatible_first_model_id(base_url: str, *, timeout_sec: float = 2.5) -> str:
    """Best-effort probe for the loaded model id of a local OpenAI-compatible server.

    Used when ``mlx_model`` is unset but a server is already running with a loaded weights
    path registered as the model id (avoids sending ``local-model``, which MLX rejects).

    Prefers the mlx-vlm ``GET /health.loaded_model`` (the actually-loaded model) and only
    falls back to ``GET /v1/models`` ``data[0].id``. mlx-vlm's ``/v1/models`` lists the
    whole HF cache, so the first entry is NOT necessarily the loaded model — using it would
    point chat at the wrong model and trigger an eviction/reload.
    """
    bu = (base_url or "").strip().rstrip("/")
    if not bu:
        return ""
    timeout = max(0.5, float(timeout_sec))
    try:
        import httpx

        # /health lives at the server root, not under /v1.
        root = bu[:-3].rstrip("/") if bu.endswith("/v1") else bu
        try:
            hresp = httpx.get(
                root + "/health",
                headers={"Accept": "application/json"},
                timeout=timeout,
            )
            if hresp.status_code == 200:
                loaded = str((hresp.json() or {}).get("loaded_model") or "").strip()
                if loaded:
                    return loaded
        except Exception:
            _log.debug("OpenAI /health probe failed for base_url=%r", bu, exc_info=True)

        resp = httpx.get(
            bu + "/models",
            headers={"Accept": "application/json"},
            timeout=timeout,
        )
        resp.raise_for_status()
        body = resp.json()
        data = body.get("data")
        if not isinstance(data, list):
            return ""
        for item in data:
            if not isinstance(item, dict):
                continue
            mid = str(item.get("id") or "").strip()
            if mid:
                return mid
    except Exception:
        _log.debug("OpenAI model probe failed for base_url=%r", bu, exc_info=True)
    return ""


def get_mlx_native_tool_probe_ok() -> bool | None:
    """Return latest MLX native-tool streaming probe status from app prefs.

    Returns:
      - True/False when a probe result exists
      - None when no probe has been recorded yet
    """
    try:
        probe = _read_pref("mlx_last_tool_probe", None)
        if not isinstance(probe, dict):
            return None
        if "native_tool_calls_ok" in probe:
            return bool(probe.get("native_tool_calls_ok"))
        if "native_tool_calls" in probe:
            return bool(probe.get("native_tool_calls"))
        return None
    except Exception:
        return None


def get_llm_context_length():
    return _int_pref("llm_context_length", _DEFAULT_CONTEXT_LENGTH, _MIN_CONTEXT_LENGTH, _MAX_CONTEXT_LENGTH)


def get_research_context_params(context_length: int) -> dict:
    """Return research/trim params scaled to context window. Call with get_llm_context_length()."""
    if context_length <= 32_000:
        return {"research_result_max_chars": 10_000, "research_keep_last_n": 1, "cap_fraction": 0.70}
    if context_length <= 64_000:
        return {"research_result_max_chars": 15_000, "research_keep_last_n": 2, "cap_fraction": 0.75}
    if context_length <= 96_000:
        return {"research_result_max_chars": 18_000, "research_keep_last_n": 2, "cap_fraction": 0.78}
    if context_length <= 128_000:
        return {"research_result_max_chars": 20_000, "research_keep_last_n": 3, "cap_fraction": 0.80}
    if context_length <= 256_000:
        return {"research_result_max_chars": 25_000, "research_keep_last_n": 4, "cap_fraction": 0.82}
    if context_length <= 512_000:
        return {"research_result_max_chars": 30_000, "research_keep_last_n": 5, "cap_fraction": 0.84}
    # 1M+
    return {"research_result_max_chars": 35_000, "research_keep_last_n": 6, "cap_fraction": 0.85}


def get_llm_temperature():
    return _float_pref("llm_temperature", _DEFAULT_TEMPERATURE, 0.0, 2.0)


def get_llm_max_tokens():
    return _int_pref("llm_max_tokens", _DEFAULT_MAX_TOKENS, 256, 32768)


def get_vision_image_format():
    """Return how to send images to the vision model: 'base64' (default) or 'url'.
    Some backends (e.g. LM Studio) require base64; others accept http(s) URLs. Read from app_prefs.json.
    LM Studio (port 1234) only accepts base64; we override to 'base64' when that backend is in use."""
    base_url = get_backend_base_url()
    if "1234" in (base_url or ""):
        return "base64"
    try:
        prefs = load_app_prefs_cached()
        fmt = (prefs.get("vision_image_format") or "base64").strip().lower()
        return fmt if fmt in ("base64", "url") else "base64"
    except Exception:
        return "base64"


def get_agent_loop_limits() -> dict:
    """Return tool-round and guardrail caps from app_prefs with safe clamps.

    Keys in ``~/.mugi/app_prefs.json``:

    - ``agent_max_tool_rounds`` — max tool-calling rounds per run (default 35; clamped 1–500).
    - ``agent_guard_retries_total`` — shared budget for guardrail-injected retries, including
      task-plan continuation and max-rounds summary (default 8; clamped 0–100).
    - ``task_plan_continuation_retries`` — nudges when a task plan still has open work before
      allowing a final answer (default 3; clamped 0–50).
    - ``task_plan_guard_retries_total`` — isolated retry budget for task-plan incomplete
      guard only (default 6; clamped 0–50).
    """
    return {
        "max_tool_rounds": _int_pref(
            "agent_max_tool_rounds", DEFAULT_AGENT_MAX_TOOL_ROUNDS, 1, 500
        ),
        "guard_retries_total": _int_pref(
            "agent_guard_retries_total", DEFAULT_AGENT_GUARD_RETRIES_TOTAL, 0, 100
        ),
        "task_plan_guard_retries_total": _int_pref(
            "task_plan_guard_retries_total",
            DEFAULT_TASK_PLAN_GUARD_RETRIES_TOTAL,
            0,
            50,
        ),
        "task_plan_continuation_retries": _int_pref(
            "task_plan_continuation_retries", DEFAULT_TASK_PLAN_CONTINUATION_RETRIES, 0, 50
        ),
    }


def get_research_limits() -> dict:
    """Return research gating/loop limits from prefs with safe defaults."""
    return {
        "max_search_rounds": _int_pref(
            "research_max_search_rounds", _DEFAULT_RESEARCH_MAX_SEARCH_ROUNDS, 1, 50
        ),
        "max_read_urls": _int_pref(
            "research_max_read_urls", _DEFAULT_RESEARCH_MAX_READ_URLS, 1, 100
        ),
        "min_scratch_appends": _int_pref(
            "research_min_scratch_appends", _DEFAULT_RESEARCH_MIN_SCRATCH_APPENDS, 1, 10
        ),
        "min_scratch_file_size": _int_pref(
            "research_min_scratch_file_size", _DEFAULT_RESEARCH_MIN_SCRATCH_FILE_SIZE, 200, 200_000
        ),
        "strict_cycle": bool(_read_pref("research_strict_cycle", False)),
        "llm_followups": bool(_read_pref("research_llm_followups", False)),
        "force_close_nonprogress_rounds": _int_pref(
            "research_force_close_nonprogress_rounds",
            _DEFAULT_RESEARCH_FORCE_CLOSE_NONPROGRESS_ROUNDS,
            2,
            50,
        ),
        "force_close_max_attempts": _int_pref(
            "research_force_close_max_attempts",
            _DEFAULT_RESEARCH_FORCE_CLOSE_MAX_ATTEMPTS,
            1,
            10,
        ),
    }


def get_memory_policy_prefs() -> dict:
    """Return memory tiering/promotion preferences from app_prefs with safe defaults."""
    profile_raw = str(_read_pref("memory_retrieval_profile", _DEFAULT_MEMORY_RETRIEVAL_PROFILE) or "").strip().lower()
    if profile_raw not in {"conservative", "balanced"}:
        profile_raw = _DEFAULT_MEMORY_RETRIEVAL_PROFILE

    strictness_raw = str(_read_pref("memory_semantic_strictness", _DEFAULT_MEMORY_SEMANTIC_STRICTNESS) or "").strip().lower()
    if strictness_raw not in {"off", "insights_sections", "enforce_promotion", "warn_promotion"}:
        strictness_raw = _DEFAULT_MEMORY_SEMANTIC_STRICTNESS

    semantic_search_enabled = bool(_read_pref("memory_semantic_search_enabled", True))

    return {
        "memory_retrieval_profile": profile_raw,
        "memory_semantic_strictness": strictness_raw,
        "memory_semantic_search_enabled": semantic_search_enabled,
    }


def get_agent_loop_router_config() -> dict:
    """Optional narrow-model router / planner-executor ([agent_loop_router]; disabled by default)."""
    try:
        from core.config import get_config

        raw = (get_config() or {}).get("agent_loop_router")
        return dict(raw) if isinstance(raw, dict) else {}
    except Exception:
        _log.warning("Failed to load agent_loop_router config; router disabled", exc_info=True)
        return {}


def get_mcp_enabled() -> bool:
    """Whether the agent should load and expose MCP tools (from ~/.mugi/app_prefs.json)."""
    try:
        prefs = load_app_prefs_cached()
        return prefs.get("mcp_enabled", True) is not False
    except Exception:
        _log.warning("Failed to load app prefs for mcp_enabled; defaulting to disabled (fail-closed)")
        return False
