#!/usr/bin/env python3
"""Ontology - Typed knowledge graph for structured agent memory."""

import json
import sys
import os
import argparse
import uuid
from datetime import datetime, timezone
from pathlib import Path

GRAPH_FILE = os.environ.get("ONTOLOGY_GRAPH", "memory/ontology/graph.jsonl")
SCHEMA_FILE = os.environ.get("ONTOLOGY_SCHEMA", "memory/ontology/schema.yaml")


def ensure_dirs():
    Path(GRAPH_FILE).parent.mkdir(parents=True, exist_ok=True)
    if not Path(GRAPH_FILE).exists():
        Path(GRAPH_FILE).touch()


def load_graph():
    """Load all entities and relations from the append-only log."""
    entities = {}
    relations = []
    if not Path(GRAPH_FILE).exists():
        return entities, relations
    with open(GRAPH_FILE, "r") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                entry = json.loads(line)
            except json.JSONDecodeError:
                continue
            op = entry.get("op")
            if op == "create":
                ent = entry["entity"]
                entities[ent["id"]] = ent
            elif op == "update":
                eid = entry["id"]
                if eid in entities:
                    entities[eid]["properties"].update(entry.get("properties", {}))
                    entities[eid]["updated"] = entry.get("timestamp", datetime.now(timezone.utc).isoformat())
            elif op == "delete":
                eid = entry["id"]
                entities.pop(eid, None)
                relations = [r for r in relations if r["from"] != eid and r["to"] != eid]
            elif op == "relate":
                relations.append({
                    "from": entry["from"],
                    "rel": entry["rel"],
                    "to": entry["to"],
                    "properties": entry.get("properties", {})
                })
            elif op == "unrelate":
                relations = [r for r in relations
                             if not (r["from"] == entry["from"] and r["rel"] == entry["rel"] and r["to"] == entry["to"])]
    return entities, relations


def append_op(op_data):
    """Append an operation to the graph log."""
    ensure_dirs()
    op_data["timestamp"] = datetime.now(timezone.utc).isoformat()
    with open(GRAPH_FILE, "a") as f:
        f.write(json.dumps(op_data, ensure_ascii=False) + "\n")


def gen_id(type_name):
    """Generate a short ID based on type."""
    prefix = type_name.lower()[:4]
    short = uuid.uuid4().hex[:8]
    return f"{prefix}_{short}"


def cmd_create(args):
    props = json.loads(args.props) if args.props else {}
    eid = args.id or gen_id(args.type)
    entity = {
        "id": eid,
        "type": args.type,
        "properties": props,
        "created": datetime.now(timezone.utc).isoformat(),
        "updated": datetime.now(timezone.utc).isoformat()
    }
    append_op({"op": "create", "entity": entity})
    print(json.dumps({"status": "created", "id": eid, "entity": entity}, ensure_ascii=False, indent=2))


def cmd_update(args):
    props = json.loads(args.props) if args.props else {}
    entities, _ = load_graph()
    if args.id not in entities:
        print(json.dumps({"error": f"Entity {args.id} not found"}))
        sys.exit(1)
    append_op({"op": "update", "id": args.id, "properties": props})
    print(json.dumps({"status": "updated", "id": args.id}))


def cmd_delete(args):
    entities, _ = load_graph()
    if args.id not in entities:
        print(json.dumps({"error": f"Entity {args.id} not found"}))
        sys.exit(1)
    append_op({"op": "delete", "id": args.id})
    print(json.dumps({"status": "deleted", "id": args.id}))


def cmd_get(args):
    entities, relations = load_graph()
    if args.id not in entities:
        print(json.dumps({"error": f"Entity {args.id} not found"}))
        sys.exit(1)
    ent = entities[args.id]
    ent_relations = [r for r in relations if r["from"] == args.id or r["to"] == args.id]
    result = {**ent, "relations": ent_relations}
    print(json.dumps(result, ensure_ascii=False, indent=2))


def cmd_query(args):
    entities, _ = load_graph()
    results = []
    where = json.loads(args.where) if args.where else {}
    for eid, ent in entities.items():
        if args.type and ent.get("type") != args.type:
            continue
        match = True
        for k, v in where.items():
            if ent.get("properties", {}).get(k) != v:
                match = False
                break
        if match:
            results.append(ent)
    print(json.dumps({"count": len(results), "results": results}, ensure_ascii=False, indent=2))


def cmd_list(args):
    entities, _ = load_graph()
    results = []
    for eid, ent in entities.items():
        if args.type and ent.get("type") != args.type:
            continue
        results.append(ent)
    print(json.dumps({"count": len(results), "results": results}, ensure_ascii=False, indent=2))


def cmd_relate(args):
    entities, _ = load_graph()
    for eid in [args.from_id, args.to_id]:
        if eid not in entities:
            print(json.dumps({"error": f"Entity {eid} not found"}))
            sys.exit(1)
    props = json.loads(args.props) if args.props else {}
    append_op({"op": "relate", "from": args.from_id, "rel": args.rel, "to": args.to_id, "properties": props})
    print(json.dumps({"status": "related", "from": args.from_id, "rel": args.rel, "to": args.to_id}))


def cmd_related(args):
    entities, relations = load_graph()
    results = []
    for r in relations:
        if r["from"] == args.id or r["to"] == args.id:
            if args.rel and r["rel"] != args.rel:
                continue
            other_id = r["to"] if r["from"] == args.id else r["from"]
            if other_id in entities:
                results.append({
                    "relation": r["rel"],
                    "direction": "outgoing" if r["from"] == args.id else "incoming",
                    "entity": entities[other_id]
                })
    print(json.dumps({"count": len(results), "results": results}, ensure_ascii=False, indent=2))


def cmd_validate(args):
    """Basic validation - check required fields if schema exists."""
    entities, relations = load_graph()
    errors = []
    # Simple validation without yaml dependency
    schema_path = Path(SCHEMA_FILE)
    if not schema_path.exists():
        print(json.dumps({"status": "ok", "note": "No schema file found, skipping type validation",
                          "entities": len(entities), "relations": len(relations)}))
        return
    # If schema exists, try basic checks
    print(json.dumps({"status": "ok", "entities": len(entities), "relations": len(relations),
                       "errors": errors}))


def cmd_schema_append(args):
    """Append schema data."""
    data = json.loads(args.data)
    schema_path = Path(SCHEMA_FILE)
    schema_path.parent.mkdir(parents=True, exist_ok=True)
    # Write as JSON for simplicity (no yaml dependency needed)
    existing = {}
    if schema_path.exists():
        try:
            with open(schema_path) as f:
                import yaml
                existing = yaml.safe_load(f) or {}
        except ImportError:
            # Fallback: read as JSON
            try:
                with open(schema_path) as f:
                    existing = json.load(f)
            except:
                existing = {}
    # Merge
    for key in data:
        if key in existing and isinstance(existing[key], dict) and isinstance(data[key], dict):
            existing[key].update(data[key])
        else:
            existing[key] = data[key]
    with open(schema_path, "w") as f:
        json.dump(existing, f, indent=2, ensure_ascii=False)
    print(json.dumps({"status": "schema_updated", "path": str(schema_path)}))


def cmd_stats(args):
    """Show graph statistics."""
    entities, relations = load_graph()
    type_counts = {}
    for ent in entities.values():
        t = ent.get("type", "unknown")
        type_counts[t] = type_counts.get(t, 0) + 1
    rel_counts = {}
    for r in relations:
        rel_counts[r["rel"]] = rel_counts.get(r["rel"], 0) + 1
    print(json.dumps({
        "total_entities": len(entities),
        "total_relations": len(relations),
        "entities_by_type": type_counts,
        "relations_by_type": rel_counts
    }, ensure_ascii=False, indent=2))


def main():
    parser = argparse.ArgumentParser(description="Ontology knowledge graph CLI")
    sub = parser.add_subparsers(dest="command")

    # create
    p = sub.add_parser("create", help="Create an entity")
    p.add_argument("--type", required=True)
    p.add_argument("--props", default="{}")
    p.add_argument("--id", default=None)

    # update
    p = sub.add_parser("update", help="Update entity properties")
    p.add_argument("--id", required=True)
    p.add_argument("--props", required=True)

    # delete
    p = sub.add_parser("delete", help="Delete an entity")
    p.add_argument("--id", required=True)

    # get
    p = sub.add_parser("get", help="Get entity by ID")
    p.add_argument("--id", required=True)

    # query
    p = sub.add_parser("query", help="Query entities by type and properties")
    p.add_argument("--type", default=None)
    p.add_argument("--where", default="{}")

    # list
    p = sub.add_parser("list", help="List entities")
    p.add_argument("--type", default=None)

    # relate
    p = sub.add_parser("relate", help="Create a relation")
    p.add_argument("--from", dest="from_id", required=True)
    p.add_argument("--rel", required=True)
    p.add_argument("--to", dest="to_id", required=True)
    p.add_argument("--props", default="{}")

    # related
    p = sub.add_parser("related", help="Find related entities")
    p.add_argument("--id", required=True)
    p.add_argument("--rel", default=None)

    # validate
    sub.add_parser("validate", help="Validate graph constraints")

    # schema-append
    p = sub.add_parser("schema-append", help="Append to schema")
    p.add_argument("--data", required=True)

    # stats
    sub.add_parser("stats", help="Show graph statistics")

    args = parser.parse_args()
    if not args.command:
        parser.print_help()
        sys.exit(1)

    commands = {
        "create": cmd_create,
        "update": cmd_update,
        "delete": cmd_delete,
        "get": cmd_get,
        "query": cmd_query,
        "list": cmd_list,
        "relate": cmd_relate,
        "related": cmd_related,
        "validate": cmd_validate,
        "schema-append": cmd_schema_append,
        "stats": cmd_stats,
    }
    commands[args.command](args)


if __name__ == "__main__":
    main()