"""
Phase 86: Zero-Trust Execution — Sandbox Foreign Agents Running Local Tools.

Trust levels: untrusted, basic, verified, trusted, admin.

Storage:
  state/broker/zero_trust/contexts/{agent_id}.json
  state/broker/zero_trust/audit.jsonl
"""

import json
import logging
import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

from brynq.server import BROKER_DIR, _write_json_sync, _read_jsonl_sync, _append_jsonl_sync

logger = logging.getLogger(__name__)

# ── Paths ──────────────────────────────────────────────────────────────────

ZT_DIR = BROKER_DIR / "zero_trust"
CONTEXTS_DIR = ZT_DIR / "contexts"
AUDIT_FILE = ZT_DIR / "audit.jsonl"

for _d in (ZT_DIR, CONTEXTS_DIR):
    _d.mkdir(parents=True, exist_ok=True)

# ── Trust Levels & Default Policies ──────────────────────────────────────

TRUST_LEVELS = ["untrusted", "basic", "verified", "trusted", "admin"]

_RO_TOOLS = ["broker_read", "broker_status", "broker_search", "broker_channels", "broker_inbox"]
_BASIC_TOOLS = _RO_TOOLS + ["broker_post", "broker_list_tasks", "broker_view_profile"]
_ADMIN_DENY = ["broker_pardon", "broker_violations", "broker_broadcast", "broker_archive_channel"]

DEFAULT_POLICIES: Dict[str, Dict[str, Any]] = {
    "untrusted": {"allowed_tools": _RO_TOOLS, "denied_tools": [], "max_calls_per_minute": 10,
                  "can_write": False, "can_execute": False, "description": "Read-only access"},
    "basic": {"allowed_tools": _BASIC_TOOLS, "denied_tools": [], "max_calls_per_minute": 30,
              "can_write": True, "can_execute": False, "description": "Read + safe writes"},
    "verified": {"allowed_tools": ["*"], "denied_tools": _ADMIN_DENY, "max_calls_per_minute": 60,
                 "can_write": True, "can_execute": True, "description": "Most tools except admin"},
    "trusted": {"allowed_tools": ["*"], "denied_tools": ["broker_pardon"],
                "max_calls_per_minute": 120, "can_write": True, "can_execute": True,
                "description": "All tools except pardons"},
    "admin": {"allowed_tools": ["*"], "denied_tools": [], "max_calls_per_minute": 0,
              "can_write": True, "can_execute": True, "description": "Unrestricted access"},
}


# ── Helpers ──────────────────────────────────────────────────────────────

def _ctx_path(agent_id: str) -> Path:
    return CONTEXTS_DIR / f"{agent_id}.json"


def _load_ctx(agent_id: str) -> Optional[dict]:
    p = _ctx_path(agent_id)
    if not p.exists():
        return None
    try:
        return json.loads(p.read_text("utf-8"))
    except (json.JSONDecodeError, OSError):
        return None


def _save_ctx(ctx: dict) -> None:
    _write_json_sync(_ctx_path(ctx["agent_id"]), ctx)


def _log_audit(entry: dict) -> None:
    entry.setdefault("timestamp", datetime.now(timezone.utc).isoformat())
    entry.setdefault("id", uuid.uuid4().hex[:12])
    _append_jsonl_sync(AUDIT_FILE, entry)


# ── Public API ───────────────────────────────────────────────────────────

def get_policy(trust_level: str) -> dict:
    """Return the default tool-access policy for a trust level."""
    if trust_level not in TRUST_LEVELS:
        return {"error": f"Unknown trust level: {trust_level}"}
    return {**DEFAULT_POLICIES[trust_level], "trust_level": trust_level}


def create_execution_context(
    agent_id: str,
    trust_level: str = "untrusted",
    allowed_tools: Optional[List[str]] = None,
    denied_tools: Optional[List[str]] = None,
) -> dict:
    """Define what a foreign agent can do on this node."""
    if trust_level not in TRUST_LEVELS:
        return {"error": f"Invalid trust level: {trust_level}. Must be one of {TRUST_LEVELS}"}

    policy = DEFAULT_POLICIES[trust_level]
    ctx = {
        "agent_id": agent_id,
        "trust_level": trust_level,
        "allowed_tools": allowed_tools if allowed_tools is not None else list(policy["allowed_tools"]),
        "denied_tools": denied_tools if denied_tools is not None else list(policy["denied_tools"]),
        "max_calls_per_minute": policy["max_calls_per_minute"],
        "can_write": policy["can_write"],
        "can_execute": policy["can_execute"],
        "created_at": datetime.now(timezone.utc).isoformat(),
        "updated_at": datetime.now(timezone.utc).isoformat(),
        "total_calls": 0,
        "total_denied": 0,
        "last_call_at": None,
    }
    _save_ctx(ctx)
    _log_audit({"event": "context_created", "agent_id": agent_id,
                "trust_level": trust_level})
    return ctx


def check_tool_access(agent_id: str, tool_name: str) -> Tuple[bool, str]:
    """Check whether agent_id may invoke tool_name. Returns (allowed, reason)."""
    ctx = _load_ctx(agent_id)
    if ctx is None:
        return False, "No execution context found for agent"

    denied = ctx.get("denied_tools", [])
    if tool_name in denied:
        return False, f"Tool '{tool_name}' is explicitly denied"

    allowed = ctx.get("allowed_tools", [])
    if "*" in allowed:
        return True, "Wildcard allows all non-denied tools"

    if tool_name in allowed:
        return True, f"Tool '{tool_name}' is in allowed list"

    return False, f"Tool '{tool_name}' is not in allowed list"


def execute_sandboxed(
    agent_id: str,
    tool_name: str,
    arguments: Optional[Dict[str, Any]] = None,
) -> dict:
    """Run a tool within the agent's security context. Logs everything."""
    ctx = _load_ctx(agent_id)
    if ctx is None:
        _log_audit({"event": "exec_denied", "agent_id": agent_id,
                     "tool": tool_name, "reason": "no_context"})
        return {"error": "No execution context. Call create_execution_context first."}

    allowed, reason = check_tool_access(agent_id, tool_name)
    ctx["total_calls"] = ctx.get("total_calls", 0) + 1
    ctx["last_call_at"] = datetime.now(timezone.utc).isoformat()

    if not allowed:
        ctx["total_denied"] = ctx.get("total_denied", 0) + 1
        _save_ctx(ctx)
        _log_audit({"event": "exec_denied", "agent_id": agent_id,
                     "tool": tool_name, "reason": reason})
        return {"error": f"Access denied: {reason}", "allowed": False}

    # Tool is allowed — record the execution
    _save_ctx(ctx)
    exec_id = uuid.uuid4().hex[:12]
    _log_audit({
        "event": "exec_allowed", "exec_id": exec_id, "agent_id": agent_id,
        "tool": tool_name, "arguments": arguments or {},
        "trust_level": ctx["trust_level"],
    })
    return {
        "exec_id": exec_id,
        "allowed": True,
        "tool": tool_name,
        "agent_id": agent_id,
        "trust_level": ctx["trust_level"],
        "arguments": arguments or {},
    }


def get_trust_level(agent_id: str) -> str:
    """Return current trust level for an agent, or 'unknown'."""
    ctx = _load_ctx(agent_id)
    if ctx is None:
        return "unknown"
    return ctx["trust_level"]


def _change_trust(agent_id: str, new_level: str, direction: str, reason: str) -> dict:
    """Shared logic for promote/demote."""
    if new_level not in TRUST_LEVELS:
        return {"error": f"Invalid trust level: {new_level}"}
    ctx = _load_ctx(agent_id)
    if ctx is None:
        return {"error": f"No context for agent {agent_id}"}
    old_idx, new_idx = TRUST_LEVELS.index(ctx["trust_level"]), TRUST_LEVELS.index(new_level)
    if direction == "up" and new_idx <= old_idx:
        return {"error": f"Cannot promote: {new_level} is not higher than {ctx['trust_level']}"}
    if direction == "down" and new_idx >= old_idx:
        return {"error": f"Cannot demote: {new_level} is not lower than {ctx['trust_level']}"}
    old_level = ctx["trust_level"]
    policy = DEFAULT_POLICIES[new_level]
    ctx.update({"trust_level": new_level, "allowed_tools": list(policy["allowed_tools"]),
                "denied_tools": list(policy["denied_tools"]),
                "max_calls_per_minute": policy["max_calls_per_minute"],
                "can_write": policy["can_write"], "can_execute": policy["can_execute"],
                "updated_at": datetime.now(timezone.utc).isoformat()})
    _save_ctx(ctx)
    event = "trust_promoted" if direction == "up" else "trust_demoted"
    _log_audit({"event": event, "agent_id": agent_id,
                "from": old_level, "to": new_level, "reason": reason})
    return {"agent_id": agent_id, "old_level": old_level, "new_level": new_level}


def promote_trust(agent_id: str, new_level: str, reason: str = "") -> dict:
    """Promote an agent to a higher trust level."""
    return _change_trust(agent_id, new_level, "up", reason)


def demote_trust(agent_id: str, new_level: str, reason: str = "") -> dict:
    """Demote an agent to a lower trust level."""
    return _change_trust(agent_id, new_level, "down", reason)


def get_execution_audit(agent_id: Optional[str] = None, limit: int = 50) -> List[dict]:
    """Audit trail of all sandboxed executions, optionally filtered by agent."""
    entries = _read_jsonl_sync(AUDIT_FILE)
    if agent_id:
        entries = [e for e in entries if e.get("agent_id") == agent_id]
    return entries[-limit:]
