EventHubAiRouter/router/path_resolve.py

"""Path resolve: deterministic index first, cheap find_* tools if gaps remain."""

from __future__ import annotations

import json
import logging
import re
from typing import Any

log = logging.getLogger("ai-router")

# Absolute Windows root for EventHubDevOps (Zed multi-root).
_DEFAULT_DEVOPS_ROOT = r"C:\Users\alexc\IdeaProjects\eventHub\EventHubDevOps"

# keyword groups → relative paths under DevOps root (first match wins per group).
# Order matters: more specific rules first.
_INDEX_RULES: list[tuple[tuple[str, ...], tuple[str, ...]]] = [
    (
        ("loadtest", "load-test", "load_test"),
        (r"ift\traefik\dynamic_conf.loadtest.yml",),
    ),
    (
        (
            "traefik",
            "dynamic_conf",
            "router",
            "middleware",
            "calentiq",
            "observer",
            "host(",
            "ift.calentiq",
            "stage.calentiq",
        ),
        (r"ift\traefik\dynamic_conf.yml",),
    ),
    (
        ("portainer",),
        (r"ift\docker-compose.portainer.yml",),
    ),
    (
        ("grafana", "dashboard"),
        (
            r"ift\observability\grafana\provisioning\dashboards\dashboard.yml",
            r"ift\observability\grafana\provisioning\datasources\prometheus.yml",
        ),
    ),
    (
        ("prometheus", "observability"),
        (r"ift\observability\prometheus.yml",),
    ),
    (
        ("admin", "compose.admin", "docker-compose.admin"),
        (r"ift\docker-compose.admin.yml",),
    ),
    (
        ("client", "compose.client", "docker-compose.client"),
        (r"ift\docker-compose.client.yml",),
    ),
    (
        ("compose", "swarm", "stack", "docker-compose", "service"),
        (r"ift\docker-compose.core.yml",),
    ),
]


def _join_under(root: str, rel: str) -> str:
    root = root.rstrip("\\/")
    rel = rel.replace("/", "\\").lstrip("\\")
    return root + "\\" + rel


def _norm(path: str) -> str:
    return path.replace("/", "\\").rstrip("\\").lower()


def _blob(plan: dict[str, Any] | None) -> str:
    return json.dumps(plan or {}, ensure_ascii=False).lower()


def devops_root_from_plan(plan: dict[str, Any] | None) -> str:
    blob = json.dumps(plan or {}, ensure_ascii=False)
    m = re.search(r"([A-Za-z]:\\[^\"'\n\r]*?EventHubDevOps)", blob)
    if m:
        return m.group(1)
    return _DEFAULT_DEVOPS_ROOT


def looks_devops(plan: dict[str, Any] | None) -> bool:
    b = _blob(plan)
    return any(
        k in b
        for k in (
            "traefik",
            "calentiq",
            "eventhubdevops",
            "docker-compose",
            "swarm",
            "observer",
            "ift.",
            "devops",
        )
    )


def match_index_paths(plan: dict[str, Any] | None) -> list[str]:
    """Return absolute paths from keyword index (may be empty)."""
    if not looks_devops(plan):
        return []
    root = devops_root_from_plan(plan)
    text = _blob(plan)
    found: list[str] = []
    seen: set[str] = set()
    for keys, rels in _INDEX_RULES:
        if not any(k in text for k in keys):
            continue
        for rel in rels:
            abs_p = _join_under(root, rel)
            nk = _norm(abs_p)
            if nk in seen:
                continue
            seen.add(nk)
            found.append(abs_p)
        # one rule group is usually enough for a focused plan
        if found:
            break
    return found


def apply_deterministic_path_index(
    plan: dict[str, Any] | None,
) -> tuple[dict[str, Any] | None, list[str]]:
    """Fill plan.paths from index. Returns (plan, unresolved_queries).

    unresolved_queries empty ⇒ index covered the goal; else cheap find_* needed.
    """
    if not isinstance(plan, dict):
        return plan, []
    out = dict(plan)
    tasks = out.get("subtasks")
    if not isinstance(tasks, list) or not tasks:
        # still try goal-level match
        indexed = match_index_paths(out)
        if indexed:
            out["subtasks"] = [
                {
                    "id": "1",
                    "prompt": str(out.get("user_goal") or "edit")[:200],
                    "worker_tier": "medium_code",
                    "paths": indexed,
                    "edit_goal": "",
                    "constraints": [],
                }
            ]
            out["path_resolve"] = "index"
            return out, []
        if looks_devops(out):
            q = _queries_from_text(_blob(out))
            out["path_resolve_queries"] = q
            return out, q
        return out, []

    indexed = match_index_paths(out)
    unresolved: list[str] = []
    fixed: list[Any] = []
    for item in tasks:
        if not isinstance(item, dict):
            fixed.append(item)
            continue
        it = dict(item)
        paths = [
            p
            for p in (it.get("paths") or [])
            if isinstance(p, str) and p.strip()
        ]
        # Prefer index hits when devops; replace empty/weak paths
        if indexed:
            it["paths"] = indexed
            it["path_source"] = "index"
        elif not paths:
            q = _queries_from_text(
                str(it.get("prompt") or "")
                + " "
                + str(it.get("edit_goal") or "")
                + " "
                + _blob(out)
            )
            unresolved.extend(q)
            it["paths"] = []
            it["path_source"] = "missing"
        else:
            it["path_source"] = "planner"
        fixed.append(it)
    out["subtasks"] = fixed
    if indexed:
        out["path_resolve"] = "index"
        out["path_resolve_queries"] = []
        log.info("path_index hit → %s", indexed)
        return out, []
    # Planner had paths but index missed — if devops, still verify via find
    if looks_devops(out) and not indexed:
        unresolved = unresolved or _queries_from_text(_blob(out))
        out["path_resolve_queries"] = unresolved
        out["path_resolve"] = "needed"
        return out, unresolved
    out["path_resolve_queries"] = unresolved
    if unresolved:
        out["path_resolve"] = "needed"
    return out, unresolved


def _queries_from_text(text: str) -> list[str]:
    low = text.lower()
    qs: list[str] = []
    if any(k in low for k in ("traefik", "router", "calentiq", "observer", "dynamic")):
        qs.append("dynamic_conf.yml")
    if any(k in low for k in ("compose", "swarm", "docker")):
        qs.append("docker-compose.core.yml")
    if "portainer" in low:
        qs.append("docker-compose.portainer.yml")
    if "grafana" in low:
        qs.append("grafana")
    if "prometheus" in low:
        qs.append("prometheus.yml")
    if not qs:
        qs.append("dynamic_conf.yml")
    # unique preserve order
    out: list[str] = []
    for q in qs:
        if q not in out:
            out.append(q)
    return out[:3]


def enrich_plan_from_discovery_tools(
    plan: dict[str, Any] | None,
    messages: list[dict[str, Any]] | None,
) -> dict[str, Any] | None:
    """Pull absolute paths from find_path / list_directory tool results into plan."""
    if not isinstance(plan, dict):
        return plan
    # Never overwrite a good deterministic index hit with noisy find parses
    if str(plan.get("path_resolve") or "") == "index":
        return plan
    tasks = plan.get("subtasks") or []
    if (
        isinstance(tasks, list)
        and tasks
        and isinstance(tasks[0], dict)
        and str(tasks[0].get("path_source") or "") == "index"
    ):
        return plan
    found = _paths_from_discovery_messages(messages)
    if not found:
        return plan
    out = dict(plan)
    tasks = list(out.get("subtasks") or [])
    if not tasks:
        out["subtasks"] = [
            {
                "id": "1",
                "prompt": "edit",
                "worker_tier": "medium_code",
                "paths": found[:3],
                "path_source": "find",
            }
        ]
    else:
        fixed = []
        for item in tasks:
            if not isinstance(item, dict):
                fixed.append(item)
                continue
            it = dict(item)
            it["paths"] = found[:3]
            it["path_source"] = "find"
            fixed.append(it)
        out["subtasks"] = fixed
    out["path_resolve"] = "find"
    out["path_resolve_queries"] = []
    log.info("path_resolve find → %s", found[:3])
    return out


def _is_workspace_path(path: str) -> bool:
    """Accept only real EventHub project paths — never /etc or null:/etc artifacts."""
    key = _norm(path)
    if "eventhubdevops" not in key and "eventhub" not in key:
        return False
    if "\\etc\\" in key or key.startswith("etc\\"):
        return False
    if "\\nginx\\" in key:
        return False
    # Reject 1-letter drive artifacts from YAML like null:/etc → l:\etc
    if re.match(r"^[a-z]:\\etc\\", key):
        return False
    base = key.rsplit("\\", 1)[-1]
    if base in ("traefik.yml", "traefik.yaml", "docker-compose.yml"):
        if "\\ift\\" not in key and "\\stage\\" not in key:
            return False
    return True


def _paths_from_discovery_messages(
    messages: list[dict[str, Any]] | None,
) -> list[str]:
    """Extract existing file paths from find/list tool results only."""
    call_ids: set[str] = set()
    for msg in messages or []:
        if not isinstance(msg, dict) or msg.get("role") != "assistant":
            continue
        for tc in msg.get("tool_calls") or []:
            if not isinstance(tc, dict):
                continue
            name = str((tc.get("function") or {}).get("name") or "").lower()
            if not any(k in name for k in ("find", "list", "search", "grep", "glob")):
                continue
            # Do not treat read_file / edit_file as discovery
            if any(k in name for k in ("read", "edit", "write", "create", "delete")):
                continue
            cid = str(tc.get("id") or "")
            if cid:
                call_ids.add(cid)
    # No discovery tool calls → do not scrape read_file bodies (YAML has /etc/… paths)
    if not call_ids:
        return []
    paths: list[str] = []
    seen: set[str] = set()
    for msg in messages or []:
        if not isinstance(msg, dict) or msg.get("role") != "tool":
            continue
        if str(msg.get("tool_call_id") or "") not in call_ids:
            continue
        content = str(msg.get("content") or "")
        low = content.lower()
        if "not found" in low and len(content) < 200:
            continue
        for m in re.finditer(
            r"([A-Za-z]:[\\/][^\s\"'<>\]\n\r]+\.(?:yml|yaml|toml|md|json|conf))",
            content,
        ):
            p = m.group(1).replace("/", "\\")
            if not _is_workspace_path(p):
                continue
            nk = _norm(p)
            if nk in seen:
                continue
            seen.add(nk)
            paths.append(p)
        # Also accept plain relative hits that Zed returns under EventHubDevOps
        for m in re.finditer(
            r"((?:EventHubDevOps[\\/](?:ift|stage)[\\/][^\s\"'<>\]\n\r]+\.(?:yml|yaml)))",
            content,
            re.I,
        ):
            p = _join_under(
                _DEFAULT_DEVOPS_ROOT,
                m.group(1).split("EventHubDevOps", 1)[-1].lstrip("\\/"),
            )
            if not _is_workspace_path(p):
                continue
            nk = _norm(p)
            if nk in seen:
                continue
            seen.add(nk)
            paths.append(p)
    return paths


def discovery_tools_only(tools: list[Any], *, max_tools: int = 3) -> list[Any]:
    prefer_keys = ("find_path", "find", "list_directory", "list_dir", "grep", "search", "glob")
    ranked: list[Any] = []
    for key in prefer_keys:
        for t in tools or []:
            if not isinstance(t, dict):
                continue
            fn = t.get("function") if isinstance(t.get("function"), dict) else {}
            name = str(fn.get("name") or t.get("name") or "").lower()
            if key in name and t not in ranked:
                ranked.append(t)
        if len(ranked) >= max_tools:
            break
    return ranked[:max_tools]


def pick_path_resolve_model(cfg: dict[str, Any]) -> str:
    """Cheap model for find_* only — allow a-simple (unlike edit executor)."""
    return str(cfg.get("path_resolve_model") or "a-simple")


def path_resolve_needed(plan: dict[str, Any] | None) -> bool:
    if not isinstance(plan, dict):
        return False
    qs = plan.get("path_resolve_queries") or []
    if qs:
        return True
    return str(plan.get("path_resolve") or "") == "needed"


def force_find_path_kickstart(
    *,
    model: str,
    tools: list[Any],
    queries: list[str],
) -> dict[str, Any]:
    """Synthetic find_path so Zed searches instead of inventing read_file."""
    from agent_hier import _synthetic_tool_completion

    tool_name = "find_path"
    for t in tools or []:
        if not isinstance(t, dict):
            continue
        fn = t.get("function") if isinstance(t.get("function"), dict) else {}
        name = str(fn.get("name") or "")
        low = name.lower()
        if "find" in low or "glob" in low or ("list" in low and "dir" in low):
            tool_name = name
            break
    q = (queries[0] if queries else "dynamic_conf.yml").strip()
    args: dict[str, Any]
    if "list" in tool_name.lower():
        args = {
            "path": _join_under(_DEFAULT_DEVOPS_ROOT, r"ift\traefik"),
        }
    else:
        args = {"query": q, "path": _DEFAULT_DEVOPS_ROOT}
        # Zed variants
        args["glob"] = f"**/{q}" if "." in q else f"**/*{q}*"
    return _synthetic_tool_completion(
        model=model, tool_name=tool_name, arguments=args
    )