#!/usr/bin/env python3
"""kg.py — deterministic helpers for Clawdbot's file-based knowledge graph.

Data model:
- Workspace-relative root: life/areas/
- Entity path: life/areas/<kind>/<slug>/
- items.json: JSON array of atomic facts
- summary.md: short snapshot derived from active facts

This script is intentionally dependency-free (stdlib only).
"""

from __future__ import annotations

import argparse
import datetime as _dt
import json
import os
import re
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

WORKSPACE_ROOT = Path(__file__).resolve().parents[2]  # .../skills/knowledge-graph
# workspace root is one more up
WORKSPACE_ROOT = WORKSPACE_ROOT.parent  # .../clawd
KG_ROOT = WORKSPACE_ROOT / "life" / "areas"

ID_RE = re.compile(r"^(?P<prefix>[a-z0-9-]+)-(?P<num>\d{3})$")


def today_date() -> str:
    return _dt.date.today().isoformat()


def now_iso() -> str:
    return _dt.datetime.now().astimezone().isoformat(timespec="seconds")


def norm_entity(entity: str) -> Tuple[str, str]:
    """Parse `people/safa` into (kind, slug)."""
    entity = entity.strip().strip("/")
    parts = [p for p in entity.split("/") if p]
    if len(parts) != 2:
        raise SystemExit(f"--entity must be like 'people/safa' (got {entity!r})")
    kind, slug = parts
    if not re.fullmatch(r"[a-z0-9-]+", kind):
        raise SystemExit(f"Invalid kind {kind!r} (use [a-z0-9-]+)")
    if not re.fullmatch(r"[a-z0-9-]+", slug):
        raise SystemExit(f"Invalid slug {slug!r} (use [a-z0-9-]+)")
    return kind, slug


def entity_dir(kind: str, slug: str) -> Path:
    return KG_ROOT / kind / slug


def load_items(path: Path) -> List[Dict[str, Any]]:
    if not path.exists():
        return []
    try:
        data = json.loads(path.read_text("utf-8"))
    except json.JSONDecodeError as e:
        raise SystemExit(f"Invalid JSON in {path}: {e}")
    if not isinstance(data, list):
        raise SystemExit(f"items.json must be a JSON array: {path}")
    for i, it in enumerate(data):
        if not isinstance(it, dict):
            raise SystemExit(f"items.json entry {i} is not an object: {path}")
    return data


def write_items(path: Path, items: List[Dict[str, Any]]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(items, indent=2, ensure_ascii=False) + "\n", "utf-8")


def next_id(slug: str, items: List[Dict[str, Any]]) -> str:
    max_num = 0
    for it in items:
        raw = str(it.get("id", ""))
        m = ID_RE.match(raw)
        if not m:
            continue
        if m.group("prefix") != slug:
            continue
        max_num = max(max_num, int(m.group("num")))
    return f"{slug}-{max_num + 1:03d}"


def add_fact(
    *,
    kind: str,
    slug: str,
    fact: str,
    category: str,
    source: str,
    timestamp: str,
) -> str:
    ed = entity_dir(kind, slug)
    items_path = ed / "items.json"
    items = load_items(items_path)

    new_id = next_id(slug, items)
    item = {
        "id": new_id,
        "fact": fact.strip(),
        "category": category,
        "timestamp": timestamp,
        "source": source,
        "status": "active",
        "supersededBy": None,
    }
    items.append(item)
    write_items(items_path, items)
    return new_id


def supersede_fact(
    *,
    kind: str,
    slug: str,
    old_id: str,
    new_fact: str,
    category: str,
    source: str,
    timestamp: str,
) -> str:
    ed = entity_dir(kind, slug)
    items_path = ed / "items.json"
    items = load_items(items_path)

    # find old
    old = None
    for it in items:
        if it.get("id") == old_id:
            old = it
            break
    if old is None:
        raise SystemExit(f"Could not find old id {old_id!r} in {items_path}")

    new_id = next_id(slug, items)
    new_item = {
        "id": new_id,
        "fact": new_fact.strip(),
        "category": category,
        "timestamp": timestamp,
        "source": source,
        "status": "active",
        "supersededBy": None,
    }

    # mark old
    old["status"] = "superseded"
    old["supersededBy"] = new_id

    items.append(new_item)
    write_items(items_path, items)
    return new_id


def build_summary(kind: str, slug: str, items: List[Dict[str, Any]], max_lines: int = 12) -> str:
    active = [it for it in items if it.get("status") == "active"]

    # stable-ish ordering: newest first by timestamp, then id
    def sort_key(it: Dict[str, Any]):
        ts = str(it.get("timestamp") or "")
        return (ts, str(it.get("id") or ""))

    active.sort(key=sort_key, reverse=True)

    lines = [f"# {slug}", "", f"Updated: {today_date()}", ""]
    for it in active[:max_lines]:
        fact = str(it.get("fact") or "").strip()
        if not fact:
            continue
        lines.append(f"- {fact}")

    if len(active) > max_lines:
        lines.append(f"- …and {len(active) - max_lines} more active facts")

    lines.append("")
    return "\n".join(lines)


def summarize_entity(kind: str, slug: str) -> Path:
    ed = entity_dir(kind, slug)
    items_path = ed / "items.json"
    items = load_items(items_path)

    summary = build_summary(kind, slug, items)
    out = ed / "summary.md"
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(summary, "utf-8")
    return out


def main() -> None:
    p = argparse.ArgumentParser(prog="kg.py")
    sub = p.add_subparsers(dest="cmd", required=True)

    p_add = sub.add_parser("add", help="Add an active fact to an entity")
    p_add.add_argument("--entity", required=True, help="e.g. people/safa")
    p_add.add_argument("--fact", required=True)
    p_add.add_argument("--category", required=True)
    p_add.add_argument("--source", default="conversation")
    p_add.add_argument("--timestamp", default=today_date())
    p_add.add_argument("--summarize", action="store_true", help="Also regenerate summary.md")

    p_sup = sub.add_parser("supersede", help="Supersede an old fact by id and add a new one")
    p_sup.add_argument("--entity", required=True)
    p_sup.add_argument("--old", required=True, dest="old_id")
    p_sup.add_argument("--fact", required=True, dest="new_fact")
    p_sup.add_argument("--category", required=True)
    p_sup.add_argument("--source", default="conversation")
    p_sup.add_argument("--timestamp", default=today_date())
    p_sup.add_argument("--summarize", action="store_true")

    p_sum = sub.add_parser("summarize", help="Regenerate summary.md from active facts")
    p_sum.add_argument("--entity", required=True)

    args = p.parse_args()

    kind, slug = norm_entity(args.entity)

    if args.cmd == "add":
        new_id = add_fact(
            kind=kind,
            slug=slug,
            fact=args.fact,
            category=args.category,
            source=args.source,
            timestamp=args.timestamp,
        )
        if args.summarize:
            summarize_entity(kind, slug)
        print(new_id)
        return

    if args.cmd == "supersede":
        new_id = supersede_fact(
            kind=kind,
            slug=slug,
            old_id=args.old_id,
            new_fact=args.new_fact,
            category=args.category,
            source=args.source,
            timestamp=args.timestamp,
        )
        if args.summarize:
            summarize_entity(kind, slug)
        print(new_id)
        return

    if args.cmd == "summarize":
        out = summarize_entity(kind, slug)
        print(str(out))
        return

    raise SystemExit(f"Unknown cmd: {args.cmd}")


def list_entities(kind: Optional[str] = None) -> List[Dict[str, Any]]:
    """List all entities, optionally filtered by kind."""
    results = []
    search_root = KG_ROOT / kind if kind else KG_ROOT
    if not search_root.exists():
        return results
    
    for entity_dir in search_root.rglob("*"):
        if entity_dir.is_dir() and (entity_dir / "items.json").exists():
            rel_path = entity_dir.relative_to(KG_ROOT)
            parts = str(rel_path).split(os.sep)
            if len(parts) >= 2:
                results.append({
                    "kind": parts[0],
                    "slug": parts[1],
                    "path": str(rel_path)
                })
    return results


def query_facts(query: str, kind: Optional[str] = None, slug: Optional[str] = None) -> List[Dict[str, Any]]:
    """Search facts across entities."""
    results = []
    query_lower = query.lower()
    
    entities = list_entities(kind)
    for ent in entities:
        if slug and ent["slug"] != slug:
            continue
        
        items_path = KG_ROOT / ent["path"] / "items.json"
        items = load_items(items_path)
        
        for it in items:
            if it.get("status") == "active":
                fact = str(it.get("fact", "")).lower()
                if query_lower in fact:
                    results.append({
                        "entity": f"{ent['kind']}/{ent['slug']}",
                        "id": it.get("id"),
                        "fact": it.get("fact"),
                        "category": it.get("category"),
                        "timestamp": it.get("timestamp"),
                        "source": it.get("source")
                    })
    return results


def get_entity_facts(kind: str, slug: str) -> List[Dict[str, Any]]:
    """Get all facts for a specific entity."""
    items_path = entity_dir(kind, slug) / "items.json"
    items = load_items(items_path)
    return [it for it in items if it.get("status") == "active"]


def mcp_server() -> None:
    """Run MCP server over stdio."""
    import sys
    
    def send_response(id: Any, result: Any) -> None:
        response = {"jsonrpc": "2.0", "id": id, "result": result}
        print(json.dumps(response), flush=True)
    
    def send_error(id: Any, code: int, message: str) -> None:
        response = {"jsonrpc": "2.0", "id": id, "error": {"code": code, "message": message}}
        print(json.dumps(response), flush=True)
    
    def handle_request(req: Dict[str, Any]) -> None:
        method = req.get("method")
        params = req.get("params", {})
        req_id = req.get("id")
        
        try:
            if method == "initialize":
                send_response(req_id, {
                    "protocolVersion": "[REDACTED_PHONE]",
                    "serverInfo": {"name": "knowledge-graph", "version": "1.0.0"},
                    "capabilities": {
                        "tools": {
                            "listChanged": False
                        }
                    }
                })
            
            elif method == "tools/list":
                send_response(req_id, {
                    "tools": [
                        {
                            "name": "kg_add",
                            "description": "Add a new fact to an entity",
                            "inputSchema": {
                                "type": "object",
                                "properties": {
                                    "entity": {"type": "string", "description": "Entity path like 'people/bruce'"},
                                    "fact": {"type": "string", "description": "The fact to add"},
                                    "category": {"type": "string", "description": "Category of the fact"},
                                    "source": {"type": "string", "description": "Source of the fact", "default": "conversation"}
                                },
                                "required": ["entity", "fact", "category"]
                            }
                        },
                        {
                            "name": "kg_supersede",
                            "description": "Supersede an old fact and add a new one",
                            "inputSchema": {
                                "type": "object",
                                "properties": {
                                    "entity": {"type": "string", "description": "Entity path like 'people/bruce'"},
                                    "old_id": {"type": "string", "description": "ID of the fact to supersede"},
                                    "fact": {"type": "string", "description": "The new fact"},
                                    "category": {"type": "string", "description": "Category of the new fact"},
                                    "source": {"type": "string", "description": "Source", "default": "conversation"}
                                },
                                "required": ["entity", "old_id", "fact", "category"]
                            }
                        },
                        {
                            "name": "kg_summarize",
                            "description": "Regenerate summary.md for an entity",
                            "inputSchema": {
                                "type": "object",
                                "properties": {
                                    "entity": {"type": "string", "description": "Entity path like 'people/bruce'"}
                                },
                                "required": ["entity"]
                            }
                        },
                        {
                            "name": "kg_query",
                            "description": "Search facts across entities",
                            "inputSchema": {
                                "type": "object",
                                "properties": {
                                    "query": {"type": "string", "description": "Search query"},
                                    "kind": {"type": "string", "description": "Optional kind filter", "default": None},
                                    "slug": {"type": "string", "description": "Optional slug filter", "default": None}
                                },
                                "required": ["query"]
                            }
                        },
                        {
                            "name": "kg_list_entities",
                            "description": "List all entities in the knowledge base",
                            "inputSchema": {
                                "type": "object",
                                "properties": {
                                    "kind": {"type": "string", "description": "Optional kind filter", "default": None}
                                }
                            }
                        },
                        {
                            "name": "kg_get_facts",
                            "description": "Get all facts for a specific entity",
                            "inputSchema": {
                                "type": "object",
                                "properties": {
                                    "entity": {"type": "string", "description": "Entity path like 'people/bruce'"}
                                },
                                "required": ["entity"]
                            }
                        }
                    ]
                })
            
            elif method == "tools/call":
                tool_name = params.get("name")
                args = params.get("arguments", {})
                
                if tool_name == "kg_add":
                    kind, slug = norm_entity(args["entity"])
                    new_id = add_fact(
                        kind=kind,
                        slug=slug,
                        fact=args["fact"],
                        category=args["category"],
                        source=args.get("source", "conversation"),
                        timestamp=now_iso()
                    )
                    summarize_entity(kind, slug)
                    send_response(req_id, {"content": [{"type": "text", "text": f"Added fact {new_id}"}]})
                
                elif tool_name == "kg_supersede":
                    kind, slug = norm_entity(args["entity"])
                    new_id = supersede_fact(
                        kind=kind,
                        slug=slug,
                        old_id=args["old_id"],
                        new_fact=args["fact"],
                        category=args["category"],
                        source=args.get("source", "conversation"),
                        timestamp=now_iso()
                    )
                    summarize_entity(kind, slug)
                    send_response(req_id, {"content": [{"type": "text", "text": f"Superseded with {new_id}"}]})
                
                elif tool_name == "kg_summarize":
                    kind, slug = norm_entity(args["entity"])
                    out = summarize_entity(kind, slug)
                    send_response(req_id, {"content": [{"type": "text", "text": f"Summary regenerated: {out}"}]})
                
                elif tool_name == "kg_query":
                    results = query_facts(
                        args["query"],
                        args.get("kind"),
                        args.get("slug")
                    )
                    text = json.dumps(results, indent=2) if results else "No matching facts found."
                    send_response(req_id, {"content": [{"type": "text", "text": text}]})
                
                elif tool_name == "kg_list_entities":
                    results = list_entities(args.get("kind"))
                    text = json.dumps(results, indent=2) if results else "No entities found."
                    send_response(req_id, {"content": [{"type": "text", "text": text}]})
                
                elif tool_name == "kg_get_facts":
                    kind, slug = norm_entity(args["entity"])
                    results = get_entity_facts(kind, slug)
                    text = json.dumps(results, indent=2) if results else "No facts found."
                    send_response(req_id, {"content": [{"type": "text", "text": text}]})
                
                else:
                    send_error(req_id, -32601, f"Unknown tool: {tool_name}")
            
            else:
                send_error(req_id, -32601, f"Unknown method: {method}")
        
        except Exception as e:
            send_error(req_id, -32603, str(e))
    
    # Main loop
    for line in sys.stdin:
        line = line.strip()
        if not line:
            continue
        try:
            req = json.loads(line)
            handle_request(req)
        except json.JSONDecodeError:
            pass


def main() -> None:
    p = argparse.ArgumentParser(prog="kg.py")
    sub = p.add_subparsers(dest="cmd", required=True)

    p_add = sub.add_parser("add", help="Add an active fact to an entity")
    p_add.add_argument("--entity", required=True, help="e.g. people/safa")
    p_add.add_argument("--fact", required=True)
    p_add.add_argument("--category", required=True)
    p_add.add_argument("--source", default="conversation")
    p_add.add_argument("--timestamp", default=today_date())
    p_add.add_argument("--summarize", action="store_true", help="Also regenerate summary.md")

    p_sup = sub.add_parser("supersede", help="Supersede an old fact by id and add a new one")
    p_sup.add_argument("--entity", required=True)
    p_sup.add_argument("--old", required=True, dest="old_id")
    p_sup.add_argument("--fact", required=True, dest="new_fact")
    p_sup.add_argument("--category", required=True)
    p_sup.add_argument("--source", default="conversation")
    p_sup.add_argument("--timestamp", default=today_date())
    p_sup.add_argument("--summarize", action="store_true")

    p_sum = sub.add_parser("summarize", help="Regenerate summary.md from active facts")
    p_sum.add_argument("--entity", required=True)
    
    p_mcp = sub.add_parser("mcp", help="Run MCP server")

    args = p.parse_args()

    if args.cmd == "mcp":
        mcp_server()
        return

    kind, slug = norm_entity(args.entity)

    if args.cmd == "add":
        new_id = add_fact(
            kind=kind,
            slug=slug,
            fact=args.fact,
            category=args.category,
            source=args.source,
            timestamp=args.timestamp,
        )
        if args.summarize:
            summarize_entity(kind, slug)
        print(new_id)
        return

    if args.cmd == "supersede":
        new_id = supersede_fact(
            kind=kind,
            slug=slug,
            old_id=args.old_id,
            new_fact=args.new_fact,
            category=args.category,
            source=args.source,
            timestamp=args.timestamp,
        )
        if args.summarize:
            summarize_entity(kind, slug)
        print(new_id)
        return

    if args.cmd == "summarize":
        out = summarize_entity(kind, slug)
        print(str(out))
        return

    raise SystemExit(f"Unknown cmd: {args.cmd}")


if __name__ == "__main__":
    main()