#!/usr/bin/env python3
"""docs-lookup.py — Unified Lookup & Regex Passing Tool for docs_internal/.

Provides high-speed queries, regex passing, grammar validation, and surface
lookups for autonomous agents and operators interfacing with NetVM, Box,
and box.muse-dev.online.

Usage:
  docs-lookup.py search <query>
  docs-lookup.py surfaces [name]
  docs-lookup.py sentence [name]
  docs-lookup.py regex [name] [--test "<string>"]
  docs-lookup.py parse "<string>"
  docs-lookup.py get <collection> [key]
  docs-lookup.py cli [domain]
  docs-lookup.py overview
"""

import argparse
import json
import os
import re
import sys
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

NETVM_ROOT = Path(__file__).resolve().parent.parent
LOOKUP_INTERNAL = NETVM_ROOT / "lookup_internal"
if not LOOKUP_INTERNAL.exists() and (NETVM_ROOT / "docs_internal").exists():
    LOOKUP_INTERNAL = NETVM_ROOT / "docs_internal"
DOCS_INTERNAL = LOOKUP_INTERNAL

USE_COLOR = sys.stdout.isatty() and os.environ.get("NO_COLOR") is None

def _c(code: str, text: str) -> str:
    return f"\033[{code}m{text}\033[0m" if USE_COLOR else str(text)

def c_bold(s: str) -> str: return _c("1", s)
def c_dim(s: str) -> str: return _c("2", s)
def c_green(s: str) -> str: return _c("32", s)
def c_red(s: str) -> str: return _c("31", s)
def c_yellow(s: str) -> str: return _c("33", s)
def c_blue(s: str) -> str: return _c("34", s)
def c_cyan(s: str) -> str: return _c("36", s)
def c_magenta(s: str) -> str: return _c("35", s)


def load_json_file(filename: str) -> Dict[str, Any]:
    """Safely load a JSON file from docs_internal."""
    p = DOCS_INTERNAL / filename
    if not p.is_file():
        return {}
    try:
        with open(p, "r", encoding="utf-8") as f:
            return json.load(f)
    except Exception as e:
        print(f"Error reading {p}: {e}", file=sys.stderr)
        return {}


def load_manifest() -> Dict[str, Any]:
    return load_json_file("manifest.json")


def load_sentence_structures() -> Dict[str, Any]:
    return load_json_file("sentence_structure.json")


def load_regex_patterns() -> Dict[str, Any]:
    return load_json_file("regex_patterns.json")


def load_assistive_surfaces() -> Dict[str, Any]:
    return load_json_file("assistive_surfaces.json")


def load_cli_tools() -> Dict[str, Any]:
    return load_json_file("cli_tools.json")


def load_fleet_nodes() -> Dict[str, Any]:
    return load_json_file("fleet_nodes.json")


def compile_pattern(pat_entry: Dict[str, Any]) -> Tuple[Optional[re.Pattern], Optional[str]]:
    """Compile a regex entry with its configured flags."""
    raw_pat = pat_entry.get("pattern", "")
    flag_names = pat_entry.get("flags", [])
    flags = 0
    for fn in flag_names:
        if hasattr(re, fn):
            flags |= getattr(re, fn)
    try:
        return re.compile(raw_pat, flags), None
    except Exception as e:
        return None, str(e)


# ---------------------------------------------------------------------------
# Core Lookup Actions
# ---------------------------------------------------------------------------

def handle_overview(json_mode: bool = False):
    manifest = load_manifest()
    if json_mode:
        print(json.dumps(manifest, indent=2))
        return

    print(c_bold("\n=== docs_internal — Internal Agent & Operator Lookup Database ==="))
    print(c_dim(f"Location: {DOCS_INTERNAL} | Host: {manifest.get('host', 'box.muse-dev.online')}"))
    print(f"{manifest.get('description', '')}\n")

    print(c_cyan("Available Collections:"))
    collections = manifest.get("collections", [])
    for col in collections:
        cid = col.get("id")
        name = col.get("name")
        desc = col.get("description")
        jfile = col.get("json_file")
        mfile = col.get("md_file")
        print(f"  • {c_bold(cid):<20} {c_green(name)}")
        print(f"    {c_dim(desc)}")
        print(f"    {c_dim('Files:')} {c_yellow(jfile)} | {c_yellow(mfile)}")
        print()

    print(c_dim("Query commands: super docs [search|surfaces|sentence|regex|parse|get|cli]"))


def handle_search(query: str, json_mode: bool = False):
    q = query.lower()
    results = []

    # Search in all JSON files
    for jpath in sorted(DOCS_INTERNAL.glob("*.json")):
        try:
            data = json.loads(jpath.read_text(encoding="utf-8"))
        except Exception:
            continue

        def recurse_search(obj, path=""):
            if isinstance(obj, dict):
                for k, v in obj.items():
                    subpath = f"{path}.{k}" if path else k
                    if q in str(k).lower():
                        results.append({
                            "file": jpath.name,
                            "type": "json_key",
                            "path": subpath,
                            "match": str(k),
                            "preview": str(v)[:160]
                        })
                    recurse_search(v, subpath)
            elif isinstance(obj, list):
                for idx, item in enumerate(obj):
                    recurse_search(item, f"{path}[{idx}]")
            elif isinstance(obj, str):
                if q in obj.lower():
                    results.append({
                        "file": jpath.name,
                        "type": "json_value",
                        "path": path,
                        "match": obj[:120],
                        "preview": obj[:240]
                    })

        recurse_search(data)

    # Search in Markdown files
    for mpath in sorted(DOCS_INTERNAL.glob("*.md")):
        try:
            content = mpath.read_text(encoding="utf-8")
        except Exception:
            continue
        lines = content.splitlines()
        for idx, line in enumerate(lines, 1):
            if q in line.lower():
                results.append({
                    "file": mpath.name,
                    "type": "markdown",
                    "line": idx,
                    "match": line.strip(),
                    "preview": line.strip()
                })

    if json_mode:
        print(json.dumps({"query": query, "count": len(results), "results": results}, indent=2))
        return

    print(c_bold(f"\nSearch results for '{query}' ({len(results)} matches):"))
    if not results:
        print(c_dim("  (no matching entries found)"))
        return

    for r in results[:40]:
        if r["type"] == "markdown":
            print(f"  [{c_yellow(r['file'])}:{c_cyan(str(r['line']))}] {r['match']}")
        else:
            print(f"  [{c_blue(r['file'])}:{c_magenta(r['path'])}] {r['preview']}")


def handle_surfaces(view_name: Optional[str] = None, json_mode: bool = False):
    data = load_assistive_surfaces()
    views = data.get("views", {})

    if view_name:
        key = view_name.lower().strip()
        v = views.get(key)
        if not v:
            for k, val in views.items():
                if key in k or key in val.get("name", "").lower():
                    v = val
                    key = k
                    break
        if not v:
            err = {"error": f"Surface '{view_name}' not found", "available": list(views.keys())}
            if json_mode:
                print(json.dumps(err, indent=2))
            else:
                print(c_red(f"Error: Surface '{view_name}' not found. Available: {', '.join(views.keys())}"))
            sys.exit(1)

        if json_mode:
            print(json.dumps({key: v}, indent=2))
            return

        print(c_bold(f"\n=== Assistive Surface: {v.get('name')} (`{key}`) ==="))
        print(c_dim(f"Host: {data.get('host')} | Tab ID: {v.get('tab_id')}"))
        print(f"{c_cyan('DOM Tab Selector:')}  {c_yellow(str(v.get('dom_tab_selector')))}")
        print(f"{c_cyan('DOM Pane Selector:')} {c_yellow(str(v.get('dom_pane_selector')))}")
        
        elements = v.get("key_elements", {})
        if elements:
            print(c_bold("\nKey DOM Elements / Selectors:"))
            for el_name, sel in elements.items():
                print(f"  • {c_magenta(el_name):<20} {c_green(sel)}")

        endpoints = v.get("api_endpoints", [])
        if endpoints:
            print(c_bold("\nAssociated REST API Endpoints:"))
            for ep in endpoints:
                print(f"  • {c_bold(ep.get('method'))} {c_cyan(ep.get('path'))}")
                print(f"    {c_dim(ep.get('description'))}")
                if ep.get("curl_example"):
                    print(f"    {c_dim('curl:')} {c_yellow(ep.get('curl_example'))}")

        recipe = v.get("assistive_recipe")
        if recipe:
            print(c_bold("\nAssistive Recipe:"))
            print(f"  {recipe}")
        print()
        return

    # All views
    if json_mode:
        print(json.dumps(data, indent=2))
        return

    print(c_bold(f"\n=== Assistive Surfaces for {data.get('host', 'box.muse-dev.online')} ==="))
    print(c_dim(f"Operator PIN: {data.get('auth', {}).get('pin')} | Session Cookie: {data.get('auth', {}).get('cookie_name')}"))
    print()

    for k, v in views.items():
        name = v.get("name", k)
        tab_sel = v.get("dom_tab_selector") or "(modal/overlay)"
        eps = [f"{ep.get('method')} {ep.get('path')}" for ep in v.get("api_endpoints", [])]
        ep_summary = ", ".join(eps) if eps else "(no direct endpoint)"
        print(f"  • {c_bold(k):<16} {c_cyan(name):<30} {c_yellow(tab_sel)}")
        print(f"    {c_dim('API:')} {ep_summary}")
        if v.get("assistive_recipe"):
            print(f"    {c_dim('Hint:')} {v.get('assistive_recipe')[:100]}...")
        print()


def handle_sentence(name: Optional[str] = None, json_mode: bool = False):
    data = load_sentence_structures()
    structs = data.get("structures", {})

    if name:
        key = name.lower().strip()
        s = structs.get(key)
        if not s:
            for k, val in structs.items():
                if key in k or key in val.get("name", "").lower() or key in val.get("protocol_tag", "").lower():
                    s = val
                    key = k
                    break
        if not s:
            err = {"error": f"Sentence structure '{name}' not found", "available": list(structs.keys())}
            if json_mode:
                print(json.dumps(err, indent=2))
            else:
                print(c_red(f"Error: Sentence structure '{name}' not found. Available: {', '.join(structs.keys())}"))
            sys.exit(1)

        if json_mode:
            print(json.dumps({key: s}, indent=2))
            return

        print(c_bold(f"\n=== Protocol Structure: {s.get('name')} (`{key}`) ==="))
        print(f"{c_cyan('Tag:')}        {c_bold(s.get('protocol_tag'))}")
        print(f"{c_cyan('Template:')}   {c_green(s.get('template'))}")
        print(f"{c_cyan('Lifecycle:')}  {c_yellow(s.get('lifecycle_transition', ''))}")
        print(f"\n{c_bold('Description:')}\n  {s.get('description')}")
        print(f"\n{c_bold('Required Fields:')} {', '.join(s.get('required_fields', []))}")
        if s.get("optional_fields"):
            print(f"{c_bold('Optional Fields:')} {', '.join(s.get('optional_fields', []))}")
        print(f"\n{c_bold('Example:')}\n  {c_cyan(s.get('example'))}")
        print(f"\n{c_bold('Reply Expectation:')}\n  {s.get('reply_expectation')}")
        print()
        return

    if json_mode:
        print(json.dumps(data, indent=2))
        return

    print(c_bold("\n=== Agent Sentence Structures & Conversational Contracts ==="))
    print(c_dim("Format rules enforced by response-harvester and self_main_loop.\n"))

    for k, s in structs.items():
        tag = s.get("protocol_tag", "")
        desc = s.get("description", "")
        print(f"  • {c_bold(k):<18} {c_green(tag):<25} {s.get('name')}")
        print(f"    {c_dim(desc)}")
        print(f"    {c_dim('Template:')} {c_cyan(s.get('template', ''))}")
        print()


def handle_regex(name: Optional[str] = None, test_str: Optional[str] = None, json_mode: bool = False):
    data = load_regex_patterns()
    pats = data.get("patterns", {})

    if name:
        key = name.lower().strip()
        p = pats.get(key)
        if not p:
            for k, val in pats.items():
                if key in k or key in val.get("name", "").lower():
                    p = val
                    key = k
                    break
        if not p:
            err = {"error": f"Regex pattern '{name}' not found", "available": list(pats.keys())}
            if json_mode:
                print(json.dumps(err, indent=2))
            else:
                print(c_red(f"Error: Pattern '{name}' not found. Available: {', '.join(pats.keys())}"))
            sys.exit(1)

        compiled, comp_err = compile_pattern(p)

        test_result = None
        if test_str is not None:
            if compiled:
                m = compiled.search(test_str)
                test_result = {
                    "matched": bool(m),
                    "match_span": m.span() if m else None,
                    "matched_text": m.group(0) if m else None,
                    "named_groups": m.groupdict() if m else {},
                    "groups": list(m.groups()) if m else []
                }
            else:
                test_result = {"matched": False, "error": comp_err}

        if json_mode:
            out = {key: p, "compiled_ok": bool(compiled)}
            if test_str is not None:
                out["test_evaluation"] = test_result
            print(json.dumps(out, indent=2))
            return

        print(c_bold(f"\n=== Regex Pattern: {p.get('name')} (`{key}`) ==="))
        print(f"{c_cyan('Pattern:')}      {c_yellow(p.get('pattern'))}")
        print(f"{c_cyan('Flags:')}        {', '.join(p.get('flags', [])) or '(none)'}")
        print(f"{c_cyan('Description:')}  {p.get('description')}")
        print(f"{c_cyan('Usage:')}        {c_dim(p.get('usage', ''))}")

        ng = p.get("named_groups", {})
        if ng:
            print(c_bold("\nNamed Groups:"))
            for gname, gdesc in ng.items():
                print(f"  • {c_magenta(gname):<16} {gdesc}")

        if test_str is not None:
            print(c_bold("\nTest Execution Result:"))
            print(f"  Input: {c_dim(test_str)}")
            if test_result.get("matched"):
                print(f"  Verdict: {c_green('✔ MATCHED')}")
                print(f"  Matched Text: {c_cyan(test_result['matched_text'])}")
                if test_result["named_groups"]:
                    print(f"  Extracted Tokens:")
                    for k_grp, v_grp in test_result["named_groups"].items():
                        print(f"    - {c_magenta(k_grp)}: {c_yellow(str(v_grp))}")
            else:
                print(f"  Verdict: {c_red('✖ NO MATCH')}")
                if comp_err:
                    print(f"  Compile Error: {comp_err}")
        print()
        return

    # List patterns
    if json_mode:
        print(json.dumps(data, indent=2))
        return

    print(c_bold("\n=== Master Regex Patterns & Passing Dictionary ==="))
    print(c_dim("Patterns tested and calibrated across response-harvester, dm, and loop engines.\n"))

    for k, p in pats.items():
        print(f"  • {c_bold(k):<18} {c_green(p.get('name'))}")
        print(f"    {c_dim('Pattern:')} {c_yellow(p.get('pattern'))}")
        print(f"    {c_dim(p.get('description'))}")
        print()


def handle_parse(candidate_str: str, json_mode: bool = False):
    """Pass candidate_str through all registered regexes and extract tokens."""
    data = load_regex_patterns()
    pats = data.get("patterns", {})

    matches = []
    for k, p in pats.items():
        compiled, err = compile_pattern(p)
        if not compiled:
            continue
        m = compiled.search(candidate_str)
        if m:
            matches.append({
                "pattern_key": k,
                "pattern_name": p.get("name"),
                "matched_text": m.group(0),
                "named_groups": m.groupdict(),
                "span": m.span()
            })

    if json_mode:
        print(json.dumps({
            "input": candidate_str,
            "matched_patterns_count": len(matches),
            "matches": matches
        }, indent=2))
        return

    print(c_bold(f"\n=== Regex Parse Analysis ==="))
    print(f"Input: {c_dim(candidate_str)}\n")

    if not matches:
        print(c_yellow("  ⚠ No registered regex pattern matched this string."))
        return

    print(c_green(f"Matched {len(matches)} pattern(s):"))
    for match in matches:
        print(f"\n  • Pattern: {c_bold(match['pattern_name'])} (`{c_cyan(match['pattern_key'])}`)")
        print(f"    Matched Chunk: {c_yellow(match['matched_text'])}")
        ng = match["named_groups"]
        if ng:
            print(f"    Extracted Tokens:")
            for gk, gv in ng.items():
                print(f"      - {c_magenta(gk)}: {c_green(str(gv))}")
    print()


def handle_get(collection: str, key: Optional[str] = None, json_mode: bool = False):
    col_map = {
        "manifest": "manifest.json",
        "sentence": "sentence_structure.json",
        "sentence_structure": "sentence_structure.json",
        "regex": "regex_patterns.json",
        "regex_patterns": "regex_patterns.json",
        "surfaces": "assistive_surfaces.json",
        "assistive_surfaces": "assistive_surfaces.json",
        "cli": "cli_tools.json",
        "cli_tools": "cli_tools.json",
        "fleet": "fleet_nodes.json",
        "fleet_nodes": "fleet_nodes.json",
    }

    fname = col_map.get(collection.lower().strip())
    if not fname:
        print(c_red(f"Error: Unknown collection '{collection}'. Available: {', '.join(col_map.keys())}"), file=sys.stderr)
        sys.exit(1)

    data = load_json_file(fname)
    if key:
        # Check top level or primary container
        val = None
        for primary in ["structures", "patterns", "views", "domains", "nodes", "collections"]:
            if primary in data and isinstance(data[primary], dict) and key in data[primary]:
                val = data[primary][key]
                break
        if val is None and key in data:
            val = data[key]
        if val is None:
            print(c_red(f"Error: Key '{key}' not found in {fname}"), file=sys.stderr)
            sys.exit(1)
        data = {key: val}

    print(json.dumps(data, indent=2))


def handle_cli(domain: Optional[str] = None, json_mode: bool = False):
    data = load_cli_tools()
    domains = data.get("domains", {})

    if domain:
        d = domains.get(domain.lower().strip())
        if not d:
            print(c_red(f"Error: CLI domain '{domain}' not found. Available: {', '.join(domains.keys())}"), file=sys.stderr)
            sys.exit(1)
        if json_mode:
            print(json.dumps({domain: d}, indent=2))
            return
        print(c_bold(f"\n=== CLI Tool Domain: super {domain} / box {domain} ==="))
        print(f"Summary: {d.get('summary')}\n")
        for sub in d.get("subcommands", []):
            print(f"  • {c_green(sub['cmd'])}")
            print(f"    {c_dim(sub['desc'])}\n")
        return

    if json_mode:
        print(json.dumps(data, indent=2))
        return

    print(c_bold("\n=== Unified NetVM & Box CLI Tool Catalog ==="))
    print(c_dim("Powered by super-cli.py and bin/muse wrapper.\n"))
    for dom_k, dom_v in domains.items():
        print(f"  • {c_bold(dom_k):<16} {c_cyan(dom_v.get('summary'))}")
        for sub in dom_v.get("subcommands", [])[:2]:
            print(f"    - {c_dim(sub['cmd'])}")
        print()


# ---------------------------------------------------------------------------
# CLI Argument Parser
# ---------------------------------------------------------------------------

def build_parser():
    common = argparse.ArgumentParser(add_help=False)
    common.add_argument("--json", action="store_true", help="Output machine-readable JSON")

    parser = argparse.ArgumentParser(
        description="docs-lookup — Internal Agent & Operator Documentation Query Engine",
        parents=[common],
        formatter_class=argparse.RawDescriptionHelpFormatter
    )

    sub = parser.add_subparsers(dest="action")

    p_search = sub.add_parser("search", parents=[common], help="Full-text search across docs_internal")
    p_search.add_argument("query", help="Search keyword or phrase")

    p_surfaces = sub.add_parser("surfaces", parents=[common], help="Lookup assistive surfaces for box.muse-dev.online")
    p_surfaces.add_argument("name", nargs="?", default=None, help="Surface or tab name")

    p_sentence = sub.add_parser("sentence", parents=[common], help="Lookup agent sentence structures & protocols")
    p_sentence.add_argument("name", nargs="?", default=None, help="Structure name (e.g. work_order, result)")

    p_regex = sub.add_parser("regex", parents=[common], help="Lookup and test regex patterns")
    p_regex.add_argument("name", nargs="?", default=None, help="Pattern name (e.g. work_order, verb)")
    p_regex.add_argument("--test", dest="test_str", default=None, help="Test string to evaluate against pattern")

    p_parse = sub.add_parser("parse", parents=[common], help="Parse an agent utterance through all regex patterns")
    p_parse.add_argument("string", help="String/utterance to parse")

    p_get = sub.add_parser("get", parents=[common], help="Query raw JSON collection and key")
    p_get.add_argument("collection", help="Collection name (sentence, regex, surfaces, cli, fleet)")
    p_get.add_argument("key", nargs="?", default=None, help="Optional specific key")

    p_cli = sub.add_parser("cli", parents=[common], help="Lookup CLI commands and syntax")
    p_cli.add_argument("domain", nargs="?", default=None, help="CLI domain (fleet, dm, job, etc.)")

    p_overview = sub.add_parser("overview", parents=[common], help="Database overview and collections manifest")

    return parser



def main():
    parser = build_parser()
    args = parser.parse_args()

    act = args.action
    json_mode = getattr(args, "json", False)

    if not act or act == "overview":
        handle_overview(json_mode)
    elif act == "search":
        handle_search(args.query, json_mode)
    elif act == "surfaces":
        handle_surfaces(args.name, json_mode)
    elif act == "sentence":
        handle_sentence(args.name, json_mode)
    elif act == "regex":
        handle_regex(args.name, args.test_str, json_mode)
    elif act == "parse":
        handle_parse(args.string, json_mode)
    elif act == "get":
        handle_get(args.collection, args.key, json_mode)
    elif act == "cli":
        handle_cli(args.domain, json_mode)
    else:
        parser.print_help()


if __name__ == "__main__":
    main()