import time
from collections import defaultdict, deque
from threading import RLock

from hiveos.intelligence.contracts import validate_tool_request


def _normalize_caps(capabilities):
    if isinstance(capabilities, list):
        return capabilities
    if isinstance(capabilities, str):
        return [capabilities]
    return []


class ToolCapabilityRegistry:
    """
    Deny-by-default tool authorization and execution.
    """

    def __init__(self):
        self._lock = RLock()
        self._tools = {}
        self._call_windows = defaultdict(deque)  # (agent_id, tool_id) -> timestamps
        self._cost_windows = defaultdict(deque)  # (agent_id, tool_id) -> (ts, cost)
        self._decisions = {"allow": 0, "deny": 0}

    def register_tool(self, tool_id, capability, handler, request_schema=None, limits=None):
        with self._lock:
            self._tools[tool_id] = {
                "capability": capability,
                "handler": handler,
                "request_schema": request_schema or {},
                "limits": limits or {},
            }

    def authorize(self, agent, tool_request):
        request = validate_tool_request(tool_request)
        tool_id = request["tool_id"]

        with self._lock:
            tool = self._tools.get(tool_id)
            if not tool:
                self._decisions["deny"] += 1
                return {"allowed": False, "reason": "tool_not_registered"}

            required_cap = tool["capability"]
            agent_caps = _normalize_caps(agent.report_capabilities())
            if required_cap not in agent_caps:
                self._decisions["deny"] += 1
                return {"allowed": False, "reason": "capability_mismatch"}

            if not self._validate_schema(tool["request_schema"], request.get("args", {})):
                self._decisions["deny"] += 1
                return {"allowed": False, "reason": "schema_validation_failed"}

            if not self._within_limits(tool["limits"], request):
                self._decisions["deny"] += 1
                return {"allowed": False, "reason": "rate_or_budget_exceeded"}

            self._record_usage(request)
            self._decisions["allow"] += 1
            return {"allowed": True, "reason": "authorized"}

    def execute(self, tool_request, context=None):
        request = validate_tool_request(tool_request)
        with self._lock:
            tool = self._tools.get(request["tool_id"])
            if not tool:
                raise ValueError("Tool not registered")
            handler = tool["handler"]
        return handler(request.get("args", {}), context or {})

    def get_stats(self):
        with self._lock:
            return {
                "registered_tools": len(self._tools),
                "decisions": dict(self._decisions),
            }

    def _validate_schema(self, schema, args):
        required_args = schema.get("required_args", [])
        if not isinstance(args, dict):
            return False
        for key in required_args:
            if key not in args:
                return False
        return True

    def _within_limits(self, limits, request):
        agent_id = request["agent_id"]
        tool_id = request["tool_id"]
        now = time.time()

        # Calls per minute window
        max_calls = limits.get("max_calls_per_minute")
        call_key = (agent_id, tool_id)
        calls = self._call_windows[call_key]
        while calls and now - calls[0] > 60:
            calls.popleft()
        if max_calls is not None and len(calls) >= int(max_calls):
            return False

        # Cost per hour window
        max_cost = limits.get("max_cost_per_hour")
        cost_key = (agent_id, tool_id)
        costs = self._cost_windows[cost_key]
        while costs and now - costs[0][0] > 3600:
            costs.popleft()
        cost_now = sum(c for _, c in costs)
        req_cost = float(request.get("budget_cost", 0))
        if max_cost is not None and (cost_now + req_cost) > float(max_cost):
            return False

        return True

    def _record_usage(self, request):
        agent_id = request["agent_id"]
        tool_id = request["tool_id"]
        now = time.time()
        self._call_windows[(agent_id, tool_id)].append(now)
        self._cost_windows[(agent_id, tool_id)].append((now, float(request.get("budget_cost", 0))))

