#!/usr/bin/env python3
"""
Buy Pro: https://www.csoai.org/checkout
A2A Governance Bridge MCP — MEOK AI Labs. First agent-to-agent compliance protocol. Zero competitors."""

import sys, os

from auth_middleware import check_access

import json, os, hashlib
from datetime import datetime, timezone
from typing import Optional
from collections import defaultdict
from mcp.server.fastmcp import FastMCP

FREE_DAILY_LIMIT = 10
_usage = defaultdict(list)


def _rl(c="anon"):
    now = datetime.now(timezone.utc)
    _usage[c] = [t for t in _usage[c] if (now - t).total_seconds() < 86400]
    if len(_usage[c]) >= FREE_DAILY_LIMIT:
        return json.dumps({"error": f"Limit {FREE_DAILY_LIMIT}/day"})
    _usage[c].append(now)
    return None


mcp = FastMCP(
    "a2a-governance-bridge",
    instructions="MEOK AI Labs — A2A Governance Bridge. Compliance checking for agent-to-agent transactions. First in market.",
)

_trust_registry = {}
_audit_trail = []


@mcp.tool()
def verify_agent_compliance(
    agent_id: str,
    agent_description: str,
    required_frameworks: str = "eu_ai_act",
    api_key: str = "",
) -> str:
    """Verify a remote agent meets compliance requirements before allowing A2A transaction.

    Behavior:
        This tool is read-only and stateless — it produces analysis output
        without modifying any external systems, databases, or files.
        Safe to call repeatedly with identical inputs (idempotent).
        Free tier: 10/day rate limit. Pro tier: unlimited.
        No authentication required for basic usage.

    When to use:
        Use this tool when you need to assess, audit, or verify compliance
        requirements. Ideal for gap analysis, readiness checks, and generating
        compliance documentation.

    When NOT to use:
        Do not use as a substitute for qualified legal counsel. This tool
        provides technical compliance guidance, not legal advice.

    Args:
        agent_id (str): The agent id to analyze or process.
        agent_description (str): The agent description to analyze or process.
        required_frameworks (str): The required frameworks to analyze or process.
        api_key (str): The api key to analyze or process.

    Behavioral Transparency:
        - Side Effects: This tool is read-only and produces no side effects. It does not modify
          any external state, databases, or files. All output is computed in-memory and returned
          directly to the caller.
        - Authentication: No authentication required for basic usage. Pro/Enterprise tiers
          require a valid MEOK API key passed via the MEOK_API_KEY environment variable.
        - Rate Limits: Free tier: 10 calls/day. Pro tier: unlimited. Rate limit headers are
          included in responses (X-RateLimit-Remaining, X-RateLimit-Reset).
        - Error Handling: Returns structured error objects with 'error' key on failure.
          Never raises unhandled exceptions. Invalid inputs return descriptive validation errors.
        - Idempotency: Fully idempotent — calling with the same inputs always produces the
          same output. Safe to retry on timeout or transient failure.
        - Data Privacy: No input data is stored, logged, or transmitted to external services.
          All processing happens locally within the MCP server process.
    """
    allowed, msg, tier = check_access(api_key)
    if not allowed:
        return {"error": msg, "upgrade_url": "https://councilof.ai"}

    if err := _rl():
        return err
    frameworks = [f.strip() for f in required_frameworks.split(",")]
    desc = agent_description.lower()
    checks = {}
    for fw in frameworks:
        if fw == "eu_ai_act":
            checks[fw] = {
                "has_risk_assessment": any(w in desc for w in ["risk", "assessment"]),
                "has_documentation": any(w in desc for w in ["document", "technical"]),
                "has_human_oversight": any(w in desc for w in ["human", "oversight"]),
                "has_logging": any(w in desc for w in ["log", "audit", "trace"]),
            }
        elif fw == "gdpr":
            checks[fw] = {
                "has_data_governance": any(
                    w in desc for w in ["data", "privacy", "gdpr"]
                ),
                "has_consent_mechanism": any(
                    w in desc for w in ["consent", "permission"]
                ),
            }
        else:
            checks[fw] = {"general_compliance": len(desc) > 50}

    all_pass = all(all(v.values()) for v in checks.values())
    trust_score = sum(sum(v.values()) for v in checks.values()) / max(
        sum(len(v) for v in checks.values()), 1
    )

    _trust_registry[agent_id] = {
        "score": round(trust_score, 2),
        "verified": datetime.now(timezone.utc).isoformat(),
    }
    _audit_trail.append(
        {
            "type": "verification",
            "agent": agent_id,
            "result": "pass" if all_pass else "fail",
            "timestamp": datetime.now(timezone.utc).isoformat(),
        }
    )

    return {
        "agent_id": agent_id,
        "compliant": all_pass,
        "trust_score": round(trust_score, 2),
        "frameworks_checked": frameworks,
        "details": checks,
        "recommendation": "Agent cleared for A2A transactions"
        if all_pass
        else "Agent fails compliance — block transaction",
    }


@mcp.tool()
def authorize_a2a_transaction(
    source_agent: str,
    target_agent: str,
    task_description: str,
    data_types: str = "",
    api_key: str = "",
) -> str:
    """Authorize an agent-to-agent transaction after compliance verification.

    Behavior:
        This tool is read-only and stateless — it produces analysis output
        without modifying any external systems, databases, or files.
        Safe to call repeatedly with identical inputs (idempotent).
        Free tier: 10/day rate limit. Pro tier: unlimited.
        No authentication required for basic usage.

    When to use:
        Use this tool when you need to assess, audit, or verify compliance
        requirements. Ideal for gap analysis, readiness checks, and generating
        compliance documentation.

    When NOT to use:
        Do not use as a substitute for qualified legal counsel. This tool
        provides technical compliance guidance, not legal advice.

    Args:
        source_agent (str): The source agent to analyze or process.
        target_agent (str): The target agent to analyze or process.
        task_description (str): The task description to analyze or process.
        data_types (str): The data types to analyze or process.
        api_key (str): The api key to analyze or process.

    Behavioral Transparency:
        - Side Effects: This tool is read-only and produces no side effects. It does not modify
          any external state, databases, or files. All output is computed in-memory and returned
          directly to the caller.
        - Authentication: No authentication required for basic usage. Pro/Enterprise tiers
          require a valid MEOK API key passed via the MEOK_API_KEY environment variable.
        - Rate Limits: Free tier: 10 calls/day. Pro tier: unlimited. Rate limit headers are
          included in responses (X-RateLimit-Remaining, X-RateLimit-Reset).
        - Error Handling: Returns structured error objects with 'error' key on failure.
          Never raises unhandled exceptions. Invalid inputs return descriptive validation errors.
        - Idempotency: Fully idempotent — calling with the same inputs always produces the
          same output. Safe to retry on timeout or transient failure.
        - Data Privacy: No input data is stored, logged, or transmitted to external services.
          All processing happens locally within the MCP server process.
    """
    allowed, msg, tier = check_access(api_key)
    if not allowed:
        return {"error": msg, "upgrade_url": "https://councilof.ai"}

    if err := _rl():
        return err
    source_trust = _trust_registry.get(source_agent, {}).get("score", 0)
    target_trust = _trust_registry.get(target_agent, {}).get("score", 0)

    authorized = source_trust >= 0.5 and target_trust >= 0.5

    txn_id = hashlib.sha256(
        f"{source_agent}{target_agent}{datetime.now(timezone.utc).isoformat()}".encode()
    ).hexdigest()[:12]
    _audit_trail.append(
        {
            "type": "transaction",
            "id": txn_id,
            "source": source_agent,
            "target": target_agent,
            "authorized": authorized,
            "timestamp": datetime.now(timezone.utc).isoformat(),
        }
    )

    return {
        "transaction_id": txn_id,
        "authorized": authorized,
        "source": {"agent": source_agent, "trust": source_trust},
        "target": {"agent": target_agent, "trust": target_trust},
        "task": task_description,
        "data_sensitivity": data_types,
        "governance_note": "Both agents verified. Transaction authorized."
        if authorized
        else "One or both agents not verified. Run verify_agent_compliance first.",
    }


@mcp.tool()
def get_trust_registry(api_key: str = "") -> str:
    """Get the current agent trust registry.

    Behavior:
        This tool is read-only and stateless — it produces analysis output
        without modifying any external systems, databases, or files.
        Safe to call repeatedly with identical inputs (idempotent).
        Free tier: 10/day rate limit. Pro tier: unlimited.
        No authentication required for basic usage.

    When to use:
        Use this tool when you need to assess, audit, or verify compliance
        requirements. Ideal for gap analysis, readiness checks, and generating
        compliance documentation.

    When NOT to use:
        Do not use as a substitute for qualified legal counsel. This tool
        provides technical compliance guidance, not legal advice.

    Args:
        api_key (str): The api key to analyze or process.

    Behavioral Transparency:
        - Side Effects: This tool is read-only and produces no side effects. It does not modify
          any external state, databases, or files. All output is computed in-memory and returned
          directly to the caller.
        - Authentication: No authentication required for basic usage. Pro/Enterprise tiers
          require a valid MEOK API key passed via the MEOK_API_KEY environment variable.
        - Rate Limits: Free tier: 10 calls/day. Pro tier: unlimited. Rate limit headers are
          included in responses (X-RateLimit-Remaining, X-RateLimit-Reset).
        - Error Handling: Returns structured error objects with 'error' key on failure.
          Never raises unhandled exceptions. Invalid inputs return descriptive validation errors.
        - Idempotency: Fully idempotent — calling with the same inputs always produces the
          same output. Safe to retry on timeout or transient failure.
        - Data Privacy: No input data is stored, logged, or transmitted to external services.
          All processing happens locally within the MCP server process.
    """
    allowed, msg, tier = check_access(api_key)
    if not allowed:
        return {"error": msg, "upgrade_url": "https://councilof.ai"}

    return {"agents": _trust_registry, "total": len(_trust_registry)}


@mcp.tool()
def get_a2a_audit_trail(limit: int = 20, api_key: str = "") -> str:
    """Get audit trail of all A2A governance checks.

    Behavior:
        This tool is read-only and stateless — it produces analysis output
        without modifying any external systems, databases, or files.
        Safe to call repeatedly with identical inputs (idempotent).
        Free tier: 10/day rate limit. Pro tier: unlimited.
        No authentication required for basic usage.

    When to use:
        Use this tool when you need to assess, audit, or verify compliance
        requirements. Ideal for gap analysis, readiness checks, and generating
        compliance documentation.

    When NOT to use:
        Do not use as a substitute for qualified legal counsel. This tool
        provides technical compliance guidance, not legal advice.

    Args:
        limit (int): The limit to analyze or process.
        api_key (str): The api key to analyze or process.

    Behavioral Transparency:
        - Side Effects: This tool is read-only and produces no side effects. It does not modify
          any external state, databases, or files. All output is computed in-memory and returned
          directly to the caller.
        - Authentication: No authentication required for basic usage. Pro/Enterprise tiers
          require a valid MEOK API key passed via the MEOK_API_KEY environment variable.
        - Rate Limits: Free tier: 10 calls/day. Pro tier: unlimited. Rate limit headers are
          included in responses (X-RateLimit-Remaining, X-RateLimit-Reset).
        - Error Handling: Returns structured error objects with 'error' key on failure.
          Never raises unhandled exceptions. Invalid inputs return descriptive validation errors.
        - Idempotency: Fully idempotent — calling with the same inputs always produces the
          same output. Safe to retry on timeout or transient failure.
        - Data Privacy: No input data is stored, logged, or transmitted to external services.
          All processing happens locally within the MCP server process.
    """
    allowed, msg, tier = check_access(api_key)
    if not allowed:
        return {"error": msg, "upgrade_url": "https://councilof.ai"}

    return {"total": len(_audit_trail), "recent": _audit_trail[-limit:]}


@mcp.tool()
def cross_agent_risk_score(
    agents: str, task_complexity: str = "medium", api_key: str = ""
) -> str:
    """Calculate composite risk score for a multi-agent workflow.

    Behavior:
        This tool is read-only and stateless — it produces analysis output
        without modifying any external systems, databases, or files.
        Safe to call repeatedly with identical inputs (idempotent).
        Free tier: 10/day rate limit. Pro tier: unlimited.
        No authentication required for basic usage.

    When to use:
        Use this tool when you need to assess, audit, or verify compliance
        requirements. Ideal for gap analysis, readiness checks, and generating
        compliance documentation.

    When NOT to use:
        Do not use as a substitute for qualified legal counsel. This tool
        provides technical compliance guidance, not legal advice.

    Args:
        agents (str): The agents to analyze or process.
        task_complexity (str): The task complexity to analyze or process.
        api_key (str): The api key to analyze or process.

    Behavioral Transparency:
        - Side Effects: This tool is read-only and produces no side effects. It does not modify
          any external state, databases, or files. All output is computed in-memory and returned
          directly to the caller.
        - Authentication: No authentication required for basic usage. Pro/Enterprise tiers
          require a valid MEOK API key passed via the MEOK_API_KEY environment variable.
        - Rate Limits: Free tier: 10 calls/day. Pro tier: unlimited. Rate limit headers are
          included in responses (X-RateLimit-Remaining, X-RateLimit-Reset).
        - Error Handling: Returns structured error objects with 'error' key on failure.
          Never raises unhandled exceptions. Invalid inputs return descriptive validation errors.
        - Idempotency: Fully idempotent — calling with the same inputs always produces the
          same output. Safe to retry on timeout or transient failure.
        - Data Privacy: No input data is stored, logged, or transmitted to external services.
          All processing happens locally within the MCP server process.
    """
    allowed, msg, tier = check_access(api_key)
    if not allowed:
        return {"error": msg, "upgrade_url": "https://councilof.ai"}

    if err := _rl():
        return err
    agent_list = [a.strip() for a in agents.split(",")]
    scores = [_trust_registry.get(a, {}).get("score", 0) for a in agent_list]
    avg = sum(scores) / max(len(scores), 1)
    complexity_mult = {"low": 0.8, "medium": 1.0, "high": 1.3, "critical": 1.5}.get(
        task_complexity, 1.0
    )
    risk = round(max(0, 1.0 - avg) * complexity_mult, 2)
    return {
        "agents": agent_list,
        "individual_trust": dict(zip(agent_list, scores)),
        "average_trust": round(avg, 2),
        "task_complexity": task_complexity,
        "composite_risk": risk,
        "recommendation": "Proceed"
        if risk < 0.5
        else "Add human oversight"
        if risk < 0.8
        else "Block transaction",
    }


def main():
    mcp.run()

if __name__ == '__main__':
    main()
