#!/usr/bin/env python3
"""
Aria-Brain MCP-Server (V5 Sprint 6, 2026-05-09)

Exposes Aria-Brain as MCP tools:
  - brain_search(query, top_n)       — Hybrid BM25-Search via aria-brain-search.py
  - memory_search(query)             — grep on /root/.claude/projects/-/memory/
  - read_brain_file(rel_path)        — direct read with path validation
  - list_brain_folders()             — show folder structure of 02-Wissen
  - brain_status()                   — quick stats (files, total bytes, last commit)

Uses FastMCP from official MCP Python SDK.
Run via: claude --mcp-server aria-brain (configured in /root/.claude/mcp.json).

Path security: nur Pfade unter /root/aria/brain/ akzeptiert.
"""

import os
import sys
import subprocess
import json
import re
from pathlib import Path
from typing import Optional

from mcp.server.fastmcp import FastMCP

BRAIN_ROOT = Path(os.environ.get("ARIA_BRAIN", "/root/aria/brain")).resolve()
MEMORY_DIR = Path("/root/.claude/projects/-/memory")
SEARCH_SCRIPT = Path("/root/aria/scripts/aria-brain-search.py")

# Import aria-atomic-write (symlink unter aria_atomic_write.py vorhanden)
sys.path.insert(0, "/root/aria/scripts")
try:
    from aria_atomic_write import atomic_write_text, file_lock  # type: ignore
    _HAS_ATOMIC = True
except ImportError:
    _HAS_ATOMIC = False

# Optional: Prompt-Injection-Scanner für Read-Pfad
try:
    from aria_prompt_injection_scanner import scan_content  # type: ignore
    _HAS_SCANNER = True
except ImportError:
    _HAS_SCANNER = False

mcp = FastMCP("aria-brain")


def _validate_path(rel_path: str) -> Optional[Path]:
    """Stelle sicher, dass rel_path unter BRAIN_ROOT bleibt (kein Path-Traversal)."""
    if not rel_path:
        return None
    candidate = (BRAIN_ROOT / rel_path).resolve()
    try:
        candidate.relative_to(BRAIN_ROOT)
    except ValueError:
        return None
    return candidate


@mcp.tool()
def brain_search(query: str, top_n: int = 5) -> str:
    """
    Hybrid-Search im Aria-Brain (BM25 + Recency + Frontmatter + Path-Match).
    Liefert Top-N Treffer mit Pfad, Score, kurzem Kontext.

    Nutze für: konzeptuelle Fragen, Pattern-Lookup, „wo steht etwas zu X".
    """
    if not SEARCH_SCRIPT.exists():
        return f"ERROR: Search script not found at {SEARCH_SCRIPT}"
    if not query.strip():
        return "ERROR: Empty query"

    top_n = max(1, min(top_n, 20))
    try:
        result = subprocess.run(
            ["python3", str(SEARCH_SCRIPT), query, "--top", str(top_n)],
            capture_output=True,
            text=True,
            timeout=10,
            env={**os.environ, "ARIA_BRAIN": str(BRAIN_ROOT)},
        )
        out = result.stdout.strip()
        if result.returncode != 0:
            return f"Search exited {result.returncode}: {result.stderr.strip()[:500]}"
        return out or "Keine Treffer."
    except subprocess.TimeoutExpired:
        return "ERROR: Search timeout (>10s)"
    except Exception as exc:
        return f"ERROR: {exc}"


@mcp.tool()
def memory_search(query: str, max_hits: int = 5) -> str:
    """
    Grep-Search auf Memory-Index (`/root/.claude/projects/-/memory/`).
    Liefert Filename + Description + Erste Zeilen pro Treffer.

    Nutze für: konkrete Memory-Lookup (LRN-Records, Feedback, Project-Notes,
    User-Profile-Fakten).
    """
    if not query.strip():
        return "ERROR: Empty query"
    if not MEMORY_DIR.is_dir():
        return f"ERROR: Memory dir not found at {MEMORY_DIR}"

    keywords = [w for w in re.split(r"\W+", query.lower()) if len(w) > 3]
    if not keywords:
        return "ERROR: Query contains only short words; need >3 chars per token"
    pattern = "|".join(re.escape(k) for k in keywords[:5])

    hits = []
    for f in sorted(MEMORY_DIR.glob("*.md")):
        try:
            content = f.read_text(errors="ignore")
        except OSError:
            continue
        if re.search(pattern, content, re.IGNORECASE):
            name = f.stem
            desc_match = re.search(r"^description:\s*(.+)$", content, re.MULTILINE)
            desc = desc_match.group(1).strip()[:240] if desc_match else ""
            body_lines = [l for l in content.splitlines()[5:13] if l.strip()][:3]
            snippet = "\n".join(body_lines)[:400]
            hits.append((name, desc, snippet))
            if len(hits) >= max_hits:
                break

    if not hits:
        return f"Keine Treffer für Pattern: {pattern}"
    return "\n\n".join(
        f"**{n}**\n  {d}\n  {s}" for n, d, s in hits
    )


@mcp.tool()
def read_brain_file(rel_path: str) -> str:
    """
    Liest eine Datei aus dem Brain (`/root/aria/brain/`).
    Pfad muss relativ zum Brain-Root sein (z.B. `SOUL.md`, `02-Wissen/foo.md`).
    Path-Traversal wird verhindert.

    Liefert max 50 KB, danach truncated mit Marker.
    """
    target = _validate_path(rel_path)
    if not target:
        return f"ERROR: Invalid path '{rel_path}'. Must be relative under {BRAIN_ROOT}."
    if not target.exists():
        return f"ERROR: File not found: {target}"
    if not target.is_file():
        return f"ERROR: Not a regular file: {target}"
    try:
        content = target.read_text(errors="replace")
    except OSError as exc:
        return f"ERROR reading {target}: {exc}"

    if len(content) > 50_000:
        return content[:50_000] + f"\n\n[TRUNCATED at 50 KB. Total file size: {len(content)} bytes]"
    return content


@mcp.tool()
def list_brain_folders() -> str:
    """
    Zeigt Folder-Struktur von /root/aria/brain/ und 02-Wissen/-Sub-Folders.
    Plus File-Count pro Folder.
    """
    if not BRAIN_ROOT.is_dir():
        return f"ERROR: Brain root not found at {BRAIN_ROOT}"

    output = [f"Brain-Root: {BRAIN_ROOT}\n"]

    # Top-level files (nur .md)
    md_files = sorted(BRAIN_ROOT.glob("*.md"))
    output.append(f"Root-Files ({len(md_files)}):")
    for f in md_files:
        size = f.stat().st_size
        output.append(f"  {f.name} ({size}B)")

    # Top-level folders mit File-Count
    folders = sorted([d for d in BRAIN_ROOT.iterdir() if d.is_dir() and not d.name.startswith(".")])
    output.append(f"\nFolders ({len(folders)}):")
    for d in folders:
        files_count = sum(1 for _ in d.rglob("*.md"))
        output.append(f"  {d.name}/ — {files_count} markdown files")

        # 02-Wissen extra: zeige sub-folders
        if d.name == "02-Wissen":
            sub_folders = sorted([s for s in d.iterdir() if s.is_dir()])
            for s in sub_folders:
                sf_count = sum(1 for _ in s.glob("*.md"))
                output.append(f"    {s.name}/ — {sf_count} files")

    return "\n".join(output)


@mcp.tool()
def brain_status() -> str:
    """
    Schnelle Statistik: total files, total size, letzter Git-Commit.
    """
    if not BRAIN_ROOT.is_dir():
        return f"ERROR: Brain root not found at {BRAIN_ROOT}"

    total_files = 0
    total_bytes = 0
    for f in BRAIN_ROOT.rglob("*.md"):
        if ".git" in f.parts:
            continue
        try:
            total_files += 1
            total_bytes += f.stat().st_size
        except OSError:
            pass

    git_log = ""
    try:
        result = subprocess.run(
            ["git", "log", "-1", "--format=%h %ai %an %s"],
            cwd=BRAIN_ROOT,
            capture_output=True,
            text=True,
            timeout=3,
        )
        git_log = result.stdout.strip() or "(no git log)"
    except Exception as exc:
        git_log = f"(git log failed: {exc})"

    return (
        f"Brain-Status:\n"
        f"  Root: {BRAIN_ROOT}\n"
        f"  Total .md files: {total_files}\n"
        f"  Total size: {total_bytes} bytes ({total_bytes / 1024:.1f} KB)\n"
        f"  Last commit: {git_log}"
    )


@mcp.tool()
def brain_write(rel_path: str, content: str, frontmatter: Optional[dict] = None) -> str:
    """
    Schreibt eine Markdown-Datei ins Brain. Atomar (mkstemp + fsync + os.replace)
    + File-Lock (Hermes-Pattern aus KAR-125).

    Pfad-Validierung wie bei read_brain_file. Wenn `frontmatter` übergeben:
    YAML-Block wird auto-vorne eingefügt (nur wenn content keinen `---`-Header hat).

    Use-Cases:
    - Daily-Log-Append
    - Brain-Note neu anlegen
    - Memory-File aktualisieren

    KEINE Schreibrechte für: 04-Feedback/ (Kais' Eigentum), Schutzzonen
    (08-Finanzen, 09-Buch-Aria, 11-Legal — falls existent), .git/, .archive/.
    """
    target = _validate_path(rel_path)
    if not target:
        return f"ERROR: Invalid path '{rel_path}'. Must be relative under {BRAIN_ROOT}."

    # Schutzzonen
    rel = target.relative_to(BRAIN_ROOT).as_posix()
    forbidden_prefixes = ["04-Feedback/", "08-Finanzen/", "09-Buch-Aria/", "11-Legal/", ".git/"]
    forbidden_contains = [".archive/"]
    if any(rel.startswith(p) for p in forbidden_prefixes) or any(p in rel for p in forbidden_contains):
        return f"ERROR: write to protected zone refused: {rel}"

    if not _HAS_ATOMIC:
        return "ERROR: aria_atomic_write nicht importierbar (KAR-125 dependency)"

    target.parent.mkdir(parents=True, exist_ok=True)
    body = content
    if frontmatter and not content.lstrip().startswith("---"):
        fm_lines = ["---"]
        for k, v in frontmatter.items():
            if isinstance(v, (list, tuple)):
                v = json.dumps(list(v))
            fm_lines.append(f"{k}: {v}")
        fm_lines.append("---")
        body = "\n".join(fm_lines) + "\n\n" + content

    lock_path = target.with_suffix(target.suffix + ".lock")
    try:
        with file_lock(lock_path):
            atomic_write_text(target, body)
    except Exception as exc:
        return f"ERROR atomic write: {exc}"
    return f"WRITTEN: {target} ({len(body)} chars)"


@mcp.tool()
def brain_frontmatter_query(filter_yaml: str, max_hits: int = 20) -> str:
    """
    Sucht Brain-Files anhand Frontmatter-Filter. `filter_yaml` ist key:value-Pairs,
    je Zeile, z.B.:

        type: research
        status: aktiv
        tags: aria-radar

    Match-Logik:
    - String-Werte: case-insensitive substring oder exact match
    - `tags:` matched wenn Filter-Tag in der Tag-Liste der Note ist
    - Multiple Filter sind AND-verknüpft

    Liefert: Pfade + Title (aus Frontmatter) + die gematchten Felder.
    """
    filters = {}
    for line in filter_yaml.strip().splitlines():
        line = line.strip()
        if not line or ":" not in line:
            continue
        k, _, v = line.partition(":")
        filters[k.strip().lower()] = v.strip().lower()

    if not filters:
        return "ERROR: no filters provided. Format: 'type: research\\nstatus: aktiv'"

    hits = []
    for md in BRAIN_ROOT.rglob("*.md"):
        if any(part.startswith(".") or part in ("__pycache__",) for part in md.parts):
            continue
        try:
            text = md.read_text(encoding="utf-8", errors="replace")
        except OSError:
            continue
        if not text.lstrip().startswith("---"):
            continue
        fm_end = text.find("\n---", 4)
        if fm_end < 0:
            continue
        fm_block = text[4:fm_end].lower()
        fm_fields = {}
        for fline in fm_block.splitlines():
            if ":" in fline:
                fk, _, fv = fline.partition(":")
                fm_fields[fk.strip()] = fv.strip()
        matches_all = True
        matched = {}
        for fk, fv in filters.items():
            present = fm_fields.get(fk, "")
            if fk == "tags":
                if fv not in present:
                    matches_all = False
                    break
            else:
                if fv not in present:
                    matches_all = False
                    break
            matched[fk] = present
        if matches_all:
            rel = md.relative_to(BRAIN_ROOT).as_posix()
            title = ""
            m = re.search(r"^title:\s*(.+)$", text[:fm_end], re.MULTILINE | re.IGNORECASE)
            if m:
                title = m.group(1).strip().strip('"').strip("'")
            hits.append((rel, title, matched))
            if len(hits) >= max_hits:
                break

    if not hits:
        return f"Keine Treffer für Filter: {filters}"
    lines = [f"Found {len(hits)} matches:"]
    for rel, title, matched in hits:
        lines.append(f"  {rel}")
        if title:
            lines.append(f"    title: {title[:80]}")
        lines.append(f"    matched: {matched}")
    return "\n".join(lines)


THINK_SCRIPT = Path("/root/aria/scripts/aria-brain-think.py")

# Optional: import think module for in-process calls (faster than subprocess)
try:
    import importlib.util as _ilu
    _think_spec = _ilu.spec_from_file_location("aria_brain_think", THINK_SCRIPT)
    _think_mod = _ilu.module_from_spec(_think_spec)
    _think_spec.loader.exec_module(_think_mod)
    _HAS_THINK = True
except Exception as _think_import_err:
    _HAS_THINK = False
    _think_mod = None


@mcp.tool()
def brain_think(query: str, top_n: int = 7, no_llm: bool = False, stale_days: int = 90) -> str:
    """
    Synthesis + Gap-Analysis über das Aria Brain (KAR-622).

    Im Gegensatz zu brain_search (gibt Treffer-Liste) liefert brain_think:
    - Eine synthetisierte Antwort mit zitierten Quellen ([slug])
    - GAP ANALYSIS: STALE (veraltete Treffer), UNCITED (unbelegte Aspekte),
      CONTRADICTION (widersprüchliche Dokumente), MISSING (unbekannte Aspekte)
    - LLM-Synthese wenn ein Key verfügbar ist (Gemini-Flash → DeepSeek → GPT-4o-mini),
      sonst deterministisch extraktiv.

    Nutze für: "Was weiß das Brain über X?", "Gibt es Widersprüche zum Thema Y?",
    "Wo hat das Brain Lücken zu Z?".

    Args:
        query: Frage oder Thema für die Synthese
        top_n: Anzahl der Treffer (1-20, Default 7)
        no_llm: True = deterministischer Extraktions-Modus ohne LLM-Aufruf
        stale_days: Alter in Tagen ab dem eine Note als STALE gilt (Default 90)
    """
    if not THINK_SCRIPT.exists():
        return f"ERROR: Think script not found at {THINK_SCRIPT}"
    if not query.strip():
        return "ERROR: Empty query"

    top_n = max(1, min(top_n, 20))

    if _HAS_THINK and _think_mod:
        # In-process call (faster, no subprocess overhead)
        try:
            result = _think_mod.think(
                query=query,
                top=top_n,
                use_llm=not no_llm,
                stale_days=stale_days,
            )
            # Format as human-readable text for MCP
            lines = []
            mode_str = result["mode"]
            if result.get("llm_provider"):
                mode_str += f" ({result['llm_provider']}/{result['llm_model']})"
            lines.append(f"[think/{mode_str}] query: {query}\n")

            lines.append("ANSWER")
            lines.append(result["answer"])
            lines.append("")

            citations = result.get("citations", [])
            if citations:
                lines.append(f"CITATIONS ({len(citations)}):")
                for c in citations:
                    lines.append(f"  [{c['slug']}] brain/{c['path']} (score={c['score']:.3f})")
                lines.append("")

            gaps = result.get("gaps", {})
            stale = gaps.get("stale", [])
            if stale:
                lines.append(f"GAP STALE ({len(stale)}):")
                for s in stale:
                    lines.append(f"  {s['path']} — {s['age_days']}d — {s['reason']}")
                lines.append("")
            uncited = gaps.get("uncited", [])
            if uncited:
                lines.append(f"GAP UNCITED: {', '.join(uncited)}")
            contradiction = gaps.get("contradiction", [])
            if contradiction:
                lines.append(f"GAP CONTRADICTION ({len(contradiction)}):")
                for c in contradiction:
                    lines.append(f"  {' vs '.join(c['paths'])}: {c['signal']}")
                lines.append("")
            missing = gaps.get("missing", [])
            if missing:
                lines.append(f"GAP MISSING (Brain kennt nicht): {', '.join(missing)}")

            return "\n".join(lines)
        except Exception as exc:
            return f"ERROR in think (in-process): {exc}"

    # Subprocess fallback
    try:
        cmd = ["python3", str(THINK_SCRIPT), query, "--top", str(top_n),
               "--stale-days", str(stale_days)]
        if no_llm:
            cmd.append("--no-llm")
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            timeout=45,
            env={**os.environ, "ARIA_BRAIN": str(BRAIN_ROOT)},
        )
        out = result.stdout.strip()
        if result.returncode != 0:
            return f"Think exited {result.returncode}: {result.stderr.strip()[:500]}"
        return out or "Keine Antwort generiert."
    except subprocess.TimeoutExpired:
        return "ERROR: Think timeout (>45s)"
    except Exception as exc:
        return f"ERROR: {exc}"


if __name__ == "__main__":
    mcp.run()