"""
Invariant MCP Server — Semantic Search for LLM Agents

WHAT THIS DOES:
Instead of reading every file to find relevant code, use these tools to:
1. FIND files matching your task (locate) — 14ms vs minutes of grep
2. UNDERSTAND file structure without reading content (semantic_map)
3. VERIFY connections exist before claiming them (prove_path)

RECOMMENDED WORKFLOW:

  1. status() — Check what's indexed, how many edges
  
  2. locate(issue_text) — Finds files ranked by interference score (2^n)
     Input: paste the error, issue, or task description
     Output: ranked files (score=32 means 5 concepts matched, very relevant)
     
  3. semantic_map(file) — Get file skeleton (10x cheaper than reading)
     Shows: key concepts, connections, line numbers
     
  4. prove_path(A, B) — Verify "A relates to B" before stating it
     Returns: exists=True/False + witness path
     
  5. ingest(file) — Add new files to the index

WHY USE THIS INSTEAD OF GREP/FILE READ:
- locate() gives ONE ranked list, not N separate grep results
- semantic_map() costs ~50 tokens, full file costs ~5000 tokens
- prove_path() prevents hallucinations about connections

THEORY:
σ-edges = proven in documents, λ-edges = inferred connections
"""
from __future__ import annotations

import os
import json
from pathlib import Path
from typing import Optional

try:
    from mcp.server.fastmcp import FastMCP  # type: ignore
except Exception:  # pragma: no cover
    FastMCP = None


class _NoMCP:
    """Fallback so SDK functions can be imported without the `mcp` package."""

    def tool(self):  # noqa: D401
        def decorator(fn):
            return fn

        return decorator

    def run(self, *args, **kwargs):
        raise RuntimeError(
            "MCP server runtime requires the `mcp` Python package, but it is not installed. "
            "Install it in your environment and retry."
        )


# Initialize MCP server (or a no-op wrapper for direct imports)
mcp = FastMCP("invariant") if FastMCP else _NoMCP()

# Globals (initialized on first use)
_physics = None
_overlay = None
_overlay_path = None
_halo_meta_cache: dict[str, dict] = {}
_overlay_index = None
_overlay_index_key: Optional[tuple] = None

# Persistent disk cache for Crystal responses (1M scale optimization)
import sqlite3
from pathlib import Path

class _DiskCache:
    """SQLite-backed persistent cache for Crystal HTTP responses."""
    
    def __init__(self, path: Optional[Path] = None):
        if path is None:
            path = Path.home() / ".invariant" / "crystal_cache.db"
        path.parent.mkdir(parents=True, exist_ok=True)
        self.conn = sqlite3.connect(str(path), check_same_thread=False)
        self.conn.execute("""
            CREATE TABLE IF NOT EXISTS cache (
                key TEXT PRIMARY KEY,
                value TEXT
            )
        """)
        self.conn.commit()
    
    def get(self, key: str) -> Optional[dict]:
        cur = self.conn.execute("SELECT value FROM cache WHERE key = ?", (key,))
        row = cur.fetchone()
        if row:
            import json
            return json.loads(row[0])
        return None
    
    def set(self, key: str, value: dict):
        import json
        self.conn.execute(
            "INSERT OR REPLACE INTO cache (key, value) VALUES (?, ?)",
            (key, json.dumps(value))
        )
        self.conn.commit()
    
    def get_many(self, keys: list) -> dict:
        """Batch get for multiple keys."""
        if not keys:
            return {}
        placeholders = ",".join("?" * len(keys))
        cur = self.conn.execute(
            f"SELECT key, value FROM cache WHERE key IN ({placeholders})",
            keys
        )
        import json
        return {row[0]: json.loads(row[1]) for row in cur.fetchall()}

_disk_cache: Optional[_DiskCache] = None

def _get_disk_cache() -> _DiskCache:
    global _disk_cache
    if _disk_cache is None:
        _disk_cache = _DiskCache()
    return _disk_cache


def _ensure_initialized():
    """Lazy initialization of physics and overlay."""
    global _physics, _overlay, _overlay_path
    
    if _physics is not None:
        return
    
    from invariant_sdk.physics import HaloPhysics
    from invariant_sdk.overlay import OverlayGraph
    
    # Connect to crystal server
    server_url = os.environ.get("INVARIANT_SERVER", "http://165.22.145.158:8080")
    _physics = HaloPhysics(server_url)
    
    # Load overlay if exists
    overlay_candidates = [
        Path("./.invariant/overlay.jsonl"),
        Path("./overlay.jsonl"),
    ]
    for candidate in overlay_candidates:
        if candidate.exists():
            _overlay = OverlayGraph.load(candidate)
            _overlay_path = candidate
            break
    
    if _overlay is None:
        _overlay = OverlayGraph()
        _overlay_path = Path("./.invariant/overlay.jsonl")


def _get_halo_meta_cached(hashes, *, chunk_size: int = 4000) -> tuple[dict[str, dict], int]:
    """
    Meta-only Halo lookup with in-process cache + disk cache + chunking.
    
    Cache hierarchy:
      1. In-memory dict (_halo_meta_cache) - fastest
      2. Disk SQLite (_disk_cache) - survives restart
      3. HTTP to Crystal server - slowest

    Returns:
      (results_by_hash8, http_requests_made)
    """
    _ensure_initialized()
    if not _physics:
        return {}, 0

    global _halo_meta_cache
    disk_cache = _get_disk_cache()

    hashes_list = [str(h).lower() for h in hashes]
    
    # Step 1: Check memory cache
    missing_from_memory = [h for h in hashes_list if h and h not in _halo_meta_cache]
    
    # Step 2: Check disk cache for memory misses
    if missing_from_memory:
        disk_hits = disk_cache.get_many(missing_from_memory)
        _halo_meta_cache.update(disk_hits)  # Promote to memory
        missing_from_disk = [h for h in missing_from_memory if h not in disk_hits]
    else:
        missing_from_disk = []
    
    # Step 3: HTTP for disk misses
    http_requests = 0
    for start in range(0, len(missing_from_disk), int(chunk_size)):
        http_requests += 1
        chunk = missing_from_disk[start : start + int(chunk_size)]
        resp = _physics._client.get_halo_pages(chunk, limit=0) or {}
        _halo_meta_cache.update(resp)
        # Persist to disk for future sessions
        for h, data in resp.items():
            disk_cache.set(h, data)

    return {h: (_halo_meta_cache.get(h) or {}) for h in hashes_list}, http_requests


# ============================================================================
# TOOLS — Actions that LLM can take
# ============================================================================

@mcp.tool()
def status() -> str:
    """
    Check if Invariant is ready and what's indexed.
    
    CALL THIS FIRST to see:
    - How many files are indexed (overlay_docs)
    - How many connections exist (overlay_edges)
    - If crystal server is connected
    
    Example output:
    {
      "overlay_edges": 25402,
      "overlay_docs": 3,
      "overlay_labels": 2786
    }
    
    If overlay_edges = 0, run ingest() on your project first.
    """
    _ensure_initialized()
    
    info = {
        "crystal_id": _physics.crystal_id if _physics else "Not connected",
        "mean_mass": round(_physics.mean_mass, 4) if _physics else 0,
        "overlay_path": str(_overlay_path) if _overlay_path else None,
        "overlay_edges": _overlay.n_edges if _overlay else 0,
        "overlay_labels": len(_overlay.labels) if _overlay else 0,
        # NOTE: `OverlayGraph.sources` tracks overlay FILES, not indexed documents.
        # For progress/debugging we want unique provenance docs inside edges.
        "overlay_docs": (
            len({e.doc for edges in _overlay.edges.values() for e in edges if e.doc})
            if _overlay
            else 0
        ),
    }
    return json.dumps(info, indent=2)


@mcp.tool()
def locate(issue_text: str, max_results: int = 0) -> str:
    """
    Find relevant files with ALL occurrences of semantic anchors.
    
    USE INSTEAD OF: grep -rn, rg, find (repo-wide searches)
    
    HOW IT WORKS:
        1. Extracts semantic anchors (solid words with high Mass)
        2. Filters out Gas words (common words like 'def', 'the', 'import')
        3. Returns files + ALL occurrences with line content
    
    Example:
        locate("separability_matrix")
        → Returns: {
            file: "separable.py",
            occurrences: [
                {line: 32, content: "from .separable import separability_matrix"},
                {line: 293, content: "def separability_matrix(transform):"}
            ]
          }
        
        Agent sees all occurrences and picks which one to investigate!
    
    THEORY: All σ-observations have equal weight (Hierarchy Law).
            No magic ranking - agent has full control.
    
    Args:
        issue_text: Paste error message or function/class name to find
        max_results: How many files to return (0 = all with score > 1)
    
    Returns:
        JSON with files + occurrences list (line + content for each)
    """
    _ensure_initialized()
    
    import math
    from invariant_sdk.cli import hash8_hex
    from invariant_sdk.engine import OverlayIndex, locate_files, tokenize_query
    
    # Extract ALL words from issue text (universal tokenization)
    # Let the crystal classify them by mass (solid vs gas)
    # NO heuristics: no stopwords, no programming-specific patterns
    
    # Extract deterministic surface tokens (supports snake_case identifiers).
    unique_words = tokenize_query(issue_text)
    
    if not unique_words:
        return json.dumps({"error": "No words found in issue_text"})
    
    # Hash words and get mass from crystal
    word_hashes = {w: hash8_hex(f"Ġ{w}") for w in unique_words}
    
    # THEORY: Overlay contains σ-facts (ground truth from ingested docs)
    # Crystal is global knowledge (fallback for unknown words)
    # Check overlay FIRST, only call Crystal for words NOT in overlay
    
    solid_seeds = []
    gas_seeds = []
    unknown_words = []
    
    # Build (cached) overlay index once per overlay state.
    global _overlay_index, _overlay_index_key
    idx_key = (id(_overlay), _overlay.n_edges, len(_overlay.labels or {}))
    if _overlay_index is None or _overlay_index_key != idx_key:
        _overlay_index = OverlayIndex.build(_overlay)
        _overlay_index_key = idx_key


    # Step 1: Get mass for ALL words from Crystal first (for V.1 classification)
    # Theory: Overlay presence ≠ Solid. Must check mass against mean_mass.
    all_hashes = list(word_hashes.values())
    mass_by_hash: dict = {}
    try:
        batch_results, _http = _get_halo_meta_cached(all_hashes, chunk_size=4000)
        for h8 in all_hashes:
            result = batch_results.get(h8) or {}
            if result.get('exists'):
                meta = result.get('meta') or {}
                degree = int(meta.get('degree_total') or 0)
                mass = 1.0 / math.log(2 + max(0, degree)) if degree > 0 else 1.0
            else:
                mass = 1.0  # void words get max mass (will be grounded later)
            mass_by_hash[h8] = mass
    except Exception:
        # Fallback: assume all have high mass
        for h8 in all_hashes:
            mass_by_hash[h8] = 1.0
    
    # Step 2: Classify words using V.1 Condensation Law
    overlay_known = set()
    for word, h8 in word_hashes.items():
        mass = mass_by_hash.get(h8, 1.0)
        is_in_overlay = h8 in _overlay_index.known_hashes or h8 in (_overlay.labels or {})
        
        # V.1: phase determined by mass, not by overlay presence
        if mass >= _physics.mean_mass:
            # Solid anchor (high mass)
            solid_seeds.append((word, h8, mass))
            if is_in_overlay:
                overlay_known.add(word)
        else:
            # Gas (low mass) — even if in overlay!
            gas_seeds.append((word, h8, mass))
            if is_in_overlay:
                overlay_known.add(word)


    # Step 3: Handle void words (not in Crystal) — Grounding Principle
    # Words with mass=1.0 that aren't in overlay are unknown (typos)
    for word, h8 in word_hashes.items():
        mass = mass_by_hash.get(h8, 1.0)
        already_classified = any(w == word for w, _, _ in solid_seeds) or any(w == word for w, _, _ in gas_seeds)
        if already_classified:
            continue
        
        # Void word: not in Crystal batch results
        if _overlay and (h8 in _overlay.edges or h8 in _overlay.labels):
            # Grounding: void + overlay = solid anchor (rare real word)
            solid_seeds.append((word, h8, 1.0))
        else:
            # Void + no overlay = unknown (typo or garbage)
            unknown_words.append(word)
    
    # File discovery + bounded previews (shared engine with UI/CLI).
    # Pass _physics to enable Query Lensing (Halo neighbor expansion)
    locate_out = locate_files(
        issue_text,
        overlay=_overlay,
        index=_overlay_index,
        physics=_physics,  # Enable Query Lensing!
        max_results=int(max_results or 0),
        preview_files=5,
        preview_occurrences=8,
        resolve_doc_path=_find_doc_path,
    )

    if locate_out.get("error"):
        return json.dumps(locate_out, indent=2)

    results = locate_out.get("results") or []
    
    return json.dumps({
        "query_words": unique_words,
        "solid_count": len(solid_seeds),
        "gas_count": len(gas_seeds),
        "unknown_count": len(unknown_words),
        "files_found": len(results),
        "results": results,
    }, indent=2)


@mcp.tool()
def semantic_map(file_path: str) -> str:
    """
    Get file structure without reading the whole file.
    
    USE INSTEAD OF: reading the entire file into context
    COST: ~50 tokens vs ~5000 tokens for full file
    
    Returns:
        - anchors: key concepts with importance (mass)
        - edges: connections between concepts with line numbers
    
    Example use case:
        After locate() finds "auth.py", use semantic_map("auth.py")
        to see its structure before deciding which lines to read.
    
    Args:
        file_path: Path to the file
    """
    _ensure_initialized()
    
    path = Path(file_path)
    if not path.exists():
        return json.dumps({"error": f"File not found: {file_path}"})
    
    result = {
        "file": file_path,
        "type": path.suffix,
    }
    
    # Count total lines for context
    try:
        with open(path, 'r', encoding='utf-8') as f:
            result["lines_total"] = sum(1 for _ in f)
    except Exception:
        result["lines_total"] = 0

    
    # Get all edges from this doc in overlay
    doc_name = path.name
    edges_from_doc = []
    nodes_in_doc = set()
    
    for src, edge_list in _overlay.edges.items():
        for edge in edge_list:
            if edge.doc and (edge.doc == doc_name or edge.doc.endswith(f"/{doc_name}")):
                src_label = _overlay.get_label(src) or src[:8]
                tgt_label = _overlay.get_label(edge.tgt) or edge.tgt[:8]
                edges_from_doc.append({
                    "src": src_label,
                    "tgt": tgt_label,
                    "line": edge.line,
                    "ring": edge.ring,
                    "witness": edge.witness,  # Structural evidence bitmask
                })
                nodes_in_doc.add(src_label)
                nodes_in_doc.add(tgt_label)
    
    # Sort by line number for reading order
    edges_from_doc.sort(key=lambda e: e.get("line") or 0)
    
    # Get mass info for key concepts
    anchors = []
    if _physics and nodes_in_doc:
        # Collect hashes for batch lookup
        hash_to_label = {}
        for node_label in list(nodes_in_doc)[:20]:
            for h, l in _overlay.labels.items():
                if l == node_label:
                    hash_to_label[h] = node_label
                    break
        
        if hash_to_label:
            try:
                import math
                batch_results, _http = _get_halo_meta_cached(hash_to_label.keys(), chunk_size=4000)
                for h8, label in hash_to_label.items():
                    res = batch_results.get(h8) or {}
                    if res.get('exists'):
                        meta = res.get('meta') or {}
                        degree_total = int(meta.get('degree_total') or 0)
                        mass = 1.0 / math.log(2 + max(0, degree_total)) if degree_total > 0 else 0
                        phase = "solid" if mass > _physics.mean_mass else "gas"
                        anchors.append({
                            "word": label,
                            "mass": round(mass, 4),
                            "phase": phase,
                        })
            except Exception:
                pass
    
    anchors.sort(key=lambda a: a["mass"], reverse=True)
    
    result["total_edges"] = len(edges_from_doc)
    result["unique_concepts"] = len(nodes_in_doc)
    result["anchors"] = anchors[:10]  # Top 10 heavy concepts
    result["edges"] = edges_from_doc[:30]  # First 30 edges (in order)
    
    return json.dumps(result, indent=2)


@mcp.tool()
def prove_path(source: str, target: str, max_hops: int = 5, mode: str = "physics") -> str:
    """
    Verify a connection exists before claiming it.
    
    USE BEFORE: stating "A is related to B" or "A affects B"
    PREVENTS: hallucinating connections that don't exist
    
    Example:
        prove_path("user", "database")
        → {"exists": true, "path": ["user", "auth", "database"], "ring": "sigma"}
        
        prove_path("coffee", "database")  
        → {"exists": false} — don't claim this connection!
    
    Ring types:
        "sigma" = proven in documents (strong evidence)
        "lambda" = inferred from language patterns (weaker)
    
    Args:
        source: First concept (e.g., "user", "authentication")
        target: Second concept to check connection to
        max_hops: Search depth. Most real connections are within 3 hops.
                  Use higher values only for exploring distant connections.
        mode: "physics" (default) = path exists
              "typed" = path exists AND all edges pass infer_SEQ
    """
    _ensure_initialized()
    
    from invariant_sdk.cli import hash8_hex
    
    # Hash the concepts
    src_hash = hash8_hex(f"Ġ{source.lower()}")
    tgt_hash = hash8_hex(f"Ġ{target.lower()}")
    
    # BFS for path (track edges for typed mode)
    visited = {src_hash}
    queue = [(src_hash, [source], [])]  # (hash, path, edges_in_path)
    
    for _ in range(max_hops):
        if not queue:
            break
        
        next_queue = []
        for current, path, edges_in_path in queue:
            # Check overlay edges
            for edge in _overlay.edges.get(current, []):
                if edge.tgt == tgt_hash:
                    # Found!
                    final_path = path + [_overlay.get_label(edge.tgt) or target]
                    all_edges = edges_in_path + [edge]
                    
                    # Typed mode: verify all edges pass infer_SEQ
                    typed_verified = True
                    if mode == "typed":
                        from invariant_sdk.operators import infer_SEQ, InferResult
                        for e in all_edges:
                            # dt_observed=1 for ADJACENT edges
                            result = infer_SEQ(e, dt_observed=1)
                            if result != InferResult.TRUE:
                                typed_verified = False
                                break
                    
                    return json.dumps({
                        "exists": True,
                        "ring": edge.ring,
                        "path": final_path,
                        "doc": edge.doc,
                        "line": edge.line,
                        "provenance": f"{edge.doc}:{edge.line}" if edge.doc and edge.line else None,
                        "mode": mode,
                        "typed_verified": typed_verified if mode == "typed" else None,
                    }, indent=2)
                
                if edge.tgt not in visited:
                    visited.add(edge.tgt)
                    label = _overlay.get_label(edge.tgt) or edge.tgt[:8]
                    next_queue.append((edge.tgt, path + [label], edges_in_path + [edge]))
            
            # Check halo edges (if physics available)
            if _physics:
                try:
                    neighbors = _physics.get_neighbors(current, limit=50)
                    for n in neighbors:
                        n_hash = n.get("hash8")
                        if n_hash == tgt_hash:
                            final_path = path + [_overlay.get_label(n_hash) or target]
                            # λ paths have no σ-witness — can't be typed verified
                            return json.dumps({
                                "exists": True,
                                "ring": "lambda",  # From halo = ghost edge
                                "path": final_path,
                                "doc": None,
                                "line": None,
                                "provenance": None,
                                "mode": mode,
                                "typed_verified": False if mode == "typed" else None,  # λ can't pass
                            }, indent=2)
                        
                        if n_hash and n_hash not in visited:
                            visited.add(n_hash)
                            label = _overlay.get_label(n_hash) or n_hash[:8]
                            # λ edges have no edge object, so pass empty list
                            next_queue.append((n_hash, path + [label], edges_in_path))
                except Exception:
                    pass
        
        queue = next_queue
    
    return json.dumps({
        "exists": False,
        "ring": None,
        "path": None,
        "message": f"No path found from '{source}' to '{target}' within {max_hops} hops",
    }, indent=2)


@mcp.tool()
def prove_paths_batch(pairs: list) -> str:
    """
    Verify multiple concept connections at once (batch version of prove_path).
    
    More efficient than calling prove_path multiple times.
    
    Args:
        pairs: List of [source, target] pairs to verify, e.g. [["user", "auth"], ["api", "database"]]
    
    Returns:
        JSON with results for each pair: {pair: [src, tgt], exists: bool, ring: str|null}
    """
    _ensure_initialized()
    
    results = []
    for pair in pairs:
        if len(pair) != 2:
            results.append({"pair": pair, "error": "Invalid pair format"})
            continue
        
        src, tgt = pair
        result = json.loads(prove_path(src, tgt, max_hops=4))
        results.append({
            "pair": [src, tgt],
            "exists": result.get("exists", False),
            "ring": result.get("ring"),
            "path": result.get("path"),
            "provenance": result.get("provenance"),
        })
    
    return json.dumps({
        "total": len(results),
        "proven": sum(1 for r in results if r.get("exists")),
        "results": results,
    }, indent=2)


@mcp.tool()
def search_concept(concept: str, limit: int = 20) -> str:
    """
    Find all documents and locations where a concept appears.
    
    Use this to understand where a term is used across the project.
    
    Args:
        concept: Word or phrase to search for
        limit: Maximum results (default 20)
    
    Returns:
        JSON with all occurrences: doc, line, related concepts
    """
    _ensure_initialized()
    
    from invariant_sdk.cli import hash8_hex
    
    concept_hash = hash8_hex(f"Ġ{concept.lower()}")
    occurrences = []
    
    # Find edges where this concept is source or target
    for src, edges in _overlay.edges.items():
        for edge in edges:
            src_label = _overlay.get_label(src) or ""
            tgt_label = _overlay.get_label(edge.tgt) or ""
            
            if concept.lower() in src_label.lower() or concept.lower() in tgt_label.lower():
                occurrences.append({
                    "doc": edge.doc,
                    "line": edge.line,
                    "src": src_label,
                    "tgt": tgt_label,
                    "ring": edge.ring,
                })
            
            if len(occurrences) >= limit:
                break
        if len(occurrences) >= limit:
            break
    
    # Group by document
    by_doc = {}
    for occ in occurrences:
        doc = occ.get("doc") or "unknown"
        if doc not in by_doc:
            by_doc[doc] = []
        by_doc[doc].append(occ)
    
    return json.dumps({
        "concept": concept,
        "total_occurrences": len(occurrences),
        "documents": len(by_doc),
        "by_document": by_doc,
    }, indent=2)


@mcp.tool()
def list_docs() -> str:
    """
    List all indexed documents with their stats.
    
    Use this to see what's in the knowledge base.
    
    Returns:
        JSON with documents: path, edge count, key concepts
    """
    _ensure_initialized()
    
    docs = {}
    for src, edges in _overlay.edges.items():
        for edge in edges:
            doc = edge.doc or "unknown"
            if doc not in docs:
                docs[doc] = {"edges": 0, "concepts": set()}
            docs[doc]["edges"] += 1
            
            src_label = _overlay.get_label(src)
            tgt_label = _overlay.get_label(edge.tgt)
            if src_label:
                docs[doc]["concepts"].add(src_label)
            if tgt_label:
                docs[doc]["concepts"].add(tgt_label)
    
    result = []
    for doc, info in sorted(docs.items(), key=lambda x: x[1]["edges"], reverse=True):
        result.append({
            "doc": doc,
            "edges": info["edges"],
            "concepts": len(info["concepts"]),
            "top_concepts": list(info["concepts"])[:5],
        })
    
    return json.dumps({
        "total_documents": len(result),
        "total_edges": sum(d["edges"] for d in result),
        "documents": result,
    }, indent=2)



@mcp.tool()
def list_conflicts() -> str:
    """
    Get all detected conflicts in the overlay.
    
    Conflicts arise when the same edge (A → B) appears with different
    weights or from different documents. This is critical for legal/compliance.
    
    Returns:
        JSON list of conflicts with sources and details
    """
    _ensure_initialized()
    
    conflicts = []
    for old_edge, new_edge in _overlay.conflicts:
        conflicts.append({
            "old": {
                "doc": old_edge.doc,
                "weight": old_edge.weight,
                "line": old_edge.line,
            },
            "new": {
                "doc": new_edge.doc,
                "weight": new_edge.weight,
                "line": new_edge.line,
            },
            "target": _overlay.get_label(old_edge.tgt) or old_edge.tgt[:8],
        })
    
    return json.dumps({
        "total": len(conflicts),
        "conflicts": conflicts,
    }, indent=2)


@mcp.tool()
def context(doc: str, line: int, ctx_hash: Optional[str] = None) -> str:
    """
    Get semantic context around a specific line in a document.
    
    Uses Anchor Integrity Protocol for self-healing:
    - If ctx_hash matches at line: fresh (exact match)
    - If ctx_hash found nearby: relocated (file changed, we found it)
    - If ctx_hash not found: broken (content deleted/changed significantly)
    
    Args:
        doc: Document path
        line: Line number (1-indexed)
        ctx_hash: Optional semantic checksum for verification
    
    Returns:
        JSON with content, status (fresh/relocated/broken/unchecked), actual_line
    """
    _ensure_initialized()
    from invariant_sdk.tokenize import tokenize_with_lines
    
    path = _find_doc_path(doc)
    if not path:
        return json.dumps({"error": f"Document not found: {doc}", "status": "broken"})
    
    try:
        text = path.read_text(encoding='utf-8')
        lines = text.split('\n')
        
        if line < 1 or line > len(lines):
            return json.dumps({"error": f"Line {line} out of range", "status": "broken"})
        
        # Tokenize for hash verification
        tokens = tokenize_with_lines(text)
        
        status = "unchecked"
        actual_line = line
        
        if ctx_hash:
            # Verify hash at expected line
            line_hashes = _compute_hashes_at_line(tokens, line)
            if ctx_hash in line_hashes:
                status = "fresh"
            else:
                # Scan ±50 lines for relocated content
                found = None
                for offset in range(1, 51):
                    for check in [line - offset, line + offset]:
                        if 1 <= check <= len(lines):
                            if ctx_hash in _compute_hashes_at_line(tokens, check):
                                found = check
                                break
                    if found:
                        break
                
                if found:
                    status = "relocated"
                    actual_line = found
                else:
                    status = "broken"
        
        # Extract semantic block
        target_idx = actual_line - 1
        start_idx = target_idx
        end_idx = target_idx
        
        # Find block boundaries
        while start_idx > 0 and (target_idx - start_idx) < 5:
            if not lines[start_idx - 1].strip():
                break
            start_idx -= 1
        
        while end_idx < len(lines) - 1 and (end_idx - target_idx) < 5:
            if not lines[end_idx + 1].strip():
                break
            end_idx += 1
        
        block = lines[start_idx:end_idx + 1]
        
        return json.dumps({
            "doc": doc,
            "requested_line": line,
            "actual_line": actual_line,
            "status": status,
            "block_start": start_idx + 1,
            "block_end": end_idx + 1,
            "content": "\n".join(block),
        }, indent=2)
        
    except Exception as e:
        return json.dumps({"error": str(e), "status": "broken"})


@mcp.tool()
def scoped_grep(pattern: str, files: str, max_matches: int = 20) -> str:
    """
    Search for EXACT pattern in specific files. Returns line numbers + context.
    
    **USE INSTEAD OF:** grep -rn, rg (repo-wide searches)
    **USE AFTER:** locate() to get file list, then grep ONLY those files
    
    Theory: Scoped search (file-specific) has higher info-gain than repo-wide.
    This tool limits output to avoid token waste.
    
    Args:
        pattern: Exact string or regex pattern to search
        files: Comma-separated file paths (from locate results)
               Example: "pkg/utils.py,tests/test_utils.py"
        max_matches: Maximum matches to return (default 20, max 50)
    
    Returns:
        JSON with matches: [{file, line, content, context_before, context_after}]
    
    Example workflow:
        1. locate("ASCIIValidator test") → finds validators.py, test_validators.py
        2. scoped_grep("ASCIIUsernameValidator", "validators.py,test_validators.py")
           → exact line numbers where pattern appears
    """
    import re
    
    _ensure_initialized()
    
    # Parse files
    file_list = [f.strip() for f in files.split(",") if f.strip()]
    if not file_list:
        return json.dumps({"error": "No files specified"})
    
    # Limit max_matches
    max_matches = min(max(1, max_matches), 50)
    
    matches = []
    files_searched = 0
    files_with_matches = 0
    
    for file_path in file_list:
        # Resolve path
        path = _find_doc_path(file_path)
        if not path:
            continue
        
        try:
            text = path.read_text(encoding='utf-8')
        except Exception:
            continue
        
        files_searched += 1
        lines = text.split('\n')
        file_has_match = False
        
        for line_num, line_content in enumerate(lines, 1):
            try:
                if re.search(pattern, line_content):
                    file_has_match = True
                    
                    # Get context (1 line before/after)
                    context_before = lines[line_num - 2] if line_num > 1 else ""
                    context_after = lines[line_num] if line_num < len(lines) else ""
                    
                    matches.append({
                        "file": file_path,
                        "line": line_num,
                        "content": line_content.strip(),
                        "context_before": context_before.strip()[:100],
                        "context_after": context_after.strip()[:100],
                    })
                    
                    if len(matches) >= max_matches:
                        break
            except re.error:
                # Invalid regex, try literal search
                if pattern in line_content:
                    file_has_match = True
                    context_before = lines[line_num - 2] if line_num > 1 else ""
                    context_after = lines[line_num] if line_num < len(lines) else ""
                    
                    matches.append({
                        "file": file_path,
                        "line": line_num,
                        "content": line_content.strip(),
                        "context_before": context_before.strip()[:100],
                        "context_after": context_after.strip()[:100],
                    })
                    
                    if len(matches) >= max_matches:
                        break
        
        if file_has_match:
            files_with_matches += 1
        
        if len(matches) >= max_matches:
            break
    
    return json.dumps({
        "pattern": pattern,
        "files_searched": files_searched,
        "files_with_matches": files_with_matches,
        "total_matches": len(matches),
        "truncated": len(matches) >= max_matches,
        "matches": matches,
    }, indent=2)



@mcp.tool()
def ingest(file_path: str) -> str:
    """
    Index files or folders into the local overlay.
    
    This creates σ-facts (grounded observations) from documents.
    
    **FOR SWE-BENCH:** Point this at the repository root to index all Python files.
    
    Args:
        file_path: Path to file OR folder to ingest
                   Folders: recursively finds .py, .md, .txt files
    
    Returns:
        JSON with stats: files processed, edges added, anchors found
    
    Example:
        ingest("/path/to/repo")  # Indexes the entire repo
        ingest("utils.py")        # Indexes single file
    """
    _ensure_initialized()
    import hashlib
    import math
    import time
    from invariant_sdk.tokenize import tokenize_with_lines
    
    from invariant_sdk.cli import hash8_hex
    
    path = Path(file_path)
    if not path.exists():
        return json.dumps({"error": f"Path not found: {file_path}"})
    
    # Collect files (support both single file and folder)
    if path.is_file():
        files = [path]
    else:
        # Helper: Check if file is text (UTF-8 decodable)
        # Theory: Observable property, not heuristic - works for ALL languages
        def is_text_file(file_path: Path) -> bool:
            if not file_path.is_file():
                return False
            try:
                # Read first 512 bytes to check (optimization)
                with open(file_path, 'rb') as f:
                    sample = f.read(512)
                # Try decode as UTF-8
                sample.decode('utf-8', errors='strict')
                return True
            except (UnicodeDecodeError, OSError):
                return False
        
        # Find ALL files (no extension filtering)
        all_files = [f for f in path.rglob("*") if f.is_file()]
        
        # Filter using .gitignore if present (Theory: Explicit user declaration)
        gitignore_path = path / '.gitignore'
        if gitignore_path.exists():
            try:
                import pathspec
                gitignore_text = gitignore_path.read_text(encoding='utf-8')
                spec = pathspec.PathSpec.from_lines('gitwildmatch', gitignore_text.splitlines())
                files_after_gitignore = [f for f in all_files if not spec.match_file(str(f.relative_to(path)))]
            except Exception:
                files_after_gitignore = all_files
        else:
            files_after_gitignore = all_files

        # Always ignore protocol / build artifacts (prevents self-indexing .invariant, etc.)
        def is_default_ignored(p: Path) -> bool:
            try:
                rel = p.relative_to(path)
            except Exception:
                rel = p
            for part in rel.parts:
                if part in {".git", ".invariant", "__pycache__", ".venv", "venv", "node_modules", "dist", "build"}:
                    return True
                if part.endswith(".egg-info"):
                    return True
            return False

        files_after_gitignore = [f for f in files_after_gitignore if not is_default_ignored(f)]
        
        # Filter to text files only (UTF-8 decodable)
        files = [f for f in files_after_gitignore if is_text_file(f)]
        
        if not files:
            return json.dumps({"error": f"No text files found in {file_path}"})

    t0 = time.perf_counter()
    try:
        print(f"[invariant] ingest: scanning {len(files)} files", flush=True)
    except Exception:
        pass
    
    # OPTIMIZATION: Mega-batch all words from all files (Theory: MDL - 1 request < N requests)
    # Step 1: Collect all unique words from ALL files first
    all_words_set = set()
    file_words_map = {}  # file_path -> list of unique words
    
    for scanned_i, file_path_obj in enumerate(files, 1):
        if scanned_i == 1 or (scanned_i % 100 == 0) or scanned_i == len(files):
            try:
                print(f"[invariant] ingest: scanned {scanned_i}/{len(files)}", flush=True)
            except Exception:
                pass
        try:
            text = file_path_obj.read_text(encoding='utf-8')
        except Exception:
            continue
        
        # Tokenize
        tokens = tokenize_with_lines(text)
        
        if len(tokens) < 2:
            continue
        
        words = [w for w, _ in tokens]
        unique_words = list(dict.fromkeys(words))  # L0: Crystal decides Solid vs Gas via Phase Separation
        
        if len(unique_words) < 2:
            continue
        
        file_words_map[file_path_obj] = (tokens, unique_words)
        all_words_set.update(unique_words)

    if not all_words_set:
        return json.dumps({"error": "No words found in any files"})

    t_scan = time.perf_counter()
    try:
        print(
            f"[invariant] ingest: vocabulary {len(all_words_set)} words from {len(file_words_map)} files",
            flush=True,
        )
    except Exception:
        pass
    
    # Step 2: ONE mega-batch request for ALL words
    all_words_list = list(all_words_set)
    word_to_hash = {w: hash8_hex(f"Ġ{w}") for w in all_words_list}
    
    try:
        # Mega-batch meta lookup (chunked + cached)
        batch_results, http_requests = _get_halo_meta_cached(word_to_hash.values(), chunk_size=4000)
    except Exception as e:
        return json.dumps({"error": f"Crystal server error: {e}"})

    t_meta = time.perf_counter()
    try:
        print(f"[invariant] ingest: crystal meta ok ({http_requests} HTTP)", flush=True)
    except Exception:
        pass
    
    mean_mass = _physics.mean_mass
    
    # Step 3: Process each file using cached batch_results
    total_files = len(file_words_map)
    total_edges = 0
    total_anchors_found = 0
    files_processed = 0
    files_details = []  # Track progress for each file
    
    for file_path_obj, (tokens, unique_words) in file_words_map.items():
        # Use cached batch_results (no new HTTP request)
        candidates = []
        for word in unique_words:
            h8 = word_to_hash.get(word)
            if not h8:
                continue
            result = batch_results.get(h8) or {}
            if not result.get('exists'):
                # Unknown words = Local Anchors (solid)
                candidates.append((word, h8, 1.0, 0, False))
                continue
            meta = result.get('meta') or {}
            degree_total = int(meta.get('degree_total') or 0)
            mass = 1.0 / math.log(2 + max(0, degree_total)) if degree_total > 0 else 0
            candidates.append((word, h8, mass, degree_total, True))
        
        # === INVARIANT VII: σ-PRESENCE WINS (RUNTIME_CONTRACT v1.7) ===
        # σ-index MUST contain ALL unique tokens observed in document.
        # Filtering happens via α-coefficient at query time, not at ingest.
        #
        # Only exclusion: Hub-words (degree > √N_vocab) which are noise.
        # This ensures names like "Eric", "Bass", "Arnold" are ALWAYS indexed.
        
        # Hub threshold: √N_vocab (derived, not magic)
        n_labels = int((_physics.meta or {}).get("n_labels") or 150000)
        hub_threshold = math.sqrt(max(1, n_labels))  # ≈ 387 for 150k vocab
        
        # Index ALL tokens, except Hubs
        solid = {
            (w, h8)
            for (w, h8, mass, degree, exists) in candidates
            if (
                not exists  # OOV → always index (local discovery)
                or degree <= hub_threshold  # Not a hub → index
            )
        }
        
        # Fallback only if filter killed everything (shouldn't happen)
        if len(solid) < 2:
            # Take all non-hub words unconditionally
            solid = {
                (w, h8) 
                for (w, h8, m, deg, exists) in candidates 
                if not exists or deg <= hub_threshold
            }
            # Still nothing? Take top by mass
            if len(solid) < 2:
                top = sorted(candidates, key=lambda x: x[2], reverse=True)[:64]
                solid = {(w, h8) for (w, h8, _m, _deg, _exists) in top}
        
        if len(solid) >= 2:
            anchors = list(solid)
        else:
            top = sorted(candidates, key=lambda x: x[2], reverse=True)[:64]
            top_set = {h8 for (_w, h8, _m, _deg, _exists) in top}
            anchors = [(w, h8) for (w, h8, _m, _deg, _exists) in candidates if h8 in top_set]
        
        if len(anchors) < 2:
            continue
        
        anchor_words = {w for w, _ in anchors}
        
        # Collect occurrences
        def compute_ctx_hash(idx: int, k: int = 2) -> str:
            start = max(0, idx - k)
            end = min(len(tokens), idx + k + 1)
            window = [tokens[i][0] for i in range(start, end)]
            normalized = ' '.join(w.lower() for w in window)
            return hashlib.sha256(normalized.encode('utf-8')).hexdigest()[:8]
        
        occurrences = []
        for idx, (word, line_num) in enumerate(tokens):
            if word in anchor_words:
                h8 = word_to_hash.get(word) or hash8_hex(f"Ġ{word}")
                occurrences.append((word, h8, line_num, compute_ctx_hash(idx)))
        
        if len(occurrences) < 2:
            continue
        
        
        # doc_name = path to file relative to cwd
        # file_path_obj is already the correct relative path (from path.rglob)
        doc_name = str(file_path_obj)
        edges_added = 0
        
        for i in range(len(occurrences) - 1):
            src_word, src_h8, _, _ = occurrences[i]
            tgt_word, tgt_h8, tgt_line, tgt_ctx = occurrences[i + 1]
            
            _overlay.add_edge(
                src_h8, tgt_h8,
                weight=1.0,
                doc=doc_name,
                ring="sigma",  # All document edges are σ (facts)
                phase="solid",
                line=tgt_line,
                ctx_hash=tgt_ctx,
            )
            _overlay.define_label(src_h8, src_word)
            _overlay.define_label(tgt_h8, tgt_word)
            edges_added += 1
        
        total_edges += edges_added
        total_anchors_found += len(anchor_words)
        files_processed += 1

        if files_processed == 1 or (files_processed % 100 == 0) or files_processed == total_files:
            try:
                print(f"[invariant] ingest: indexed {files_processed}/{total_files}", flush=True)
            except Exception:
                pass
        
        # Track first 10 files for progress visibility
        if len(files_details) < 10:
            files_details.append({
                "file": doc_name,
                "edges": edges_added,
                "anchors": len(anchor_words)
            })
    
    
    # Save overlay
    _overlay_path.parent.mkdir(parents=True, exist_ok=True)
    _overlay.save(_overlay_path)

    t_save = time.perf_counter()
    
    return json.dumps({
        "success": True,
        "path": file_path,
        "total_files": total_files,
        "files_processed": files_processed,
        "total_edges": total_edges,
        "total_anchors": total_anchors_found,
        "unique_words_processed": len(all_words_set),
        "http_requests": http_requests,  # Mega-batch optimization (chunked)
        "files_sample": files_details,  # First 10 files
        "overlay_path": str(_overlay_path),
        "timing_s": {
            "scan": round(t_scan - t0, 3),
            "crystal_meta": round(t_meta - t_scan, 3),
            "index": round(t_save - t_meta, 3),
            "total": round(t_save - t0, 3),
        },
    }, indent=2)



# ============================================================================
# HELPERS
# ============================================================================

def _find_doc_path(doc: str) -> Optional[Path]:
    """Find document in project. Supports both full paths and basename fallback."""
    # Try exact path first (new overlay format with relative path)
    candidates = [
        Path(doc),
        Path(".") / doc,
        Path(".invariant") / doc,
        Path(".invariant/uploads") / doc,
        Path("docs") / doc,
    ]
    for c in candidates:
        if c.exists() and c.is_file():
            return c
    
    # Fallback: recursive search for old overlays with only basename
    # This handles case where doc="separable.py" but file is at "astropy/modeling/separable.py"
    basename = Path(doc).name
    if basename != doc:
        return None  # Already tried full path, don't search again
    
    # Limit search to avoid infinite loops in huge projects
    try:
        for found in Path(".").rglob(basename):
            if found.is_file():
                return found
    except Exception:
        pass
    
    return None


def _compute_hashes_at_line(tokens: list, target_line: int, k: int = 2) -> list:
    """Compute all ctx_hashes for tokens at a given line."""
    import hashlib
    
    line_tokens = [(i, t) for i, t in enumerate(tokens) if t[1] == target_line]
    hashes = []
    
    for anchor_idx, _ in line_tokens:
        start = max(0, anchor_idx - k)
        end = min(len(tokens), anchor_idx + k + 1)
        window = [tokens[i][0] for i in range(start, end)]
        normalized = ' '.join(window)
        h = hashlib.sha256(normalized.encode('utf-8')).hexdigest()[:8]
        hashes.append(h)
    
    return hashes


# ============================================================================
# MAIN
# ============================================================================

def main():
    """Run the MCP server."""
    mcp.run(transport='stdio')


if __name__ == "__main__":
    main()
