from typing import Any, Dict, Optional, Callable
AgentAction = Dict[str, Any]  # Mock for Python 3.14
BaseMessage = Dict[str, Any]
from pydantic import BaseModel
import json
import time
import logging
import requests
import re
import random  # Hard mode

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class ApexGateConfig(BaseModel):
    api_key: str
    api_url: str = "http://localhost:8000"
    agent_id: str = "default-agent"

class ApexGateMiddleware:
    def __init__(self, config: ApexGateConfig | str):
        if isinstance(config, str):
            config = ApexGateConfig(api_key=config)
        self.config = config
        self.session = requests.Session()
        self.session.headers["Authorization"] = f"Bearer {config.api_key}"
        self.high_risk_tools = {
            "code_execution", "api_call", "db_write", "transaction", "post"
        }
    
    def check_action(self, action: AgentAction) -> bool:
        """Legacy method - use gate_tool"""
        tool_name = action.get("tool", "")
        return self.gate_tool(tool_name, action.get("tool_input", {}))
    
    def gate_tool(self, tool_name: str, tool_input: Dict = None, llm: Callable = None) -> bool:
        """Main gating method"""
        if tool_name not in self.high_risk_tools:
            return True
        
        logger.info(f"🛡️ Gating {tool_name}")
        
        challenge = self._request_challenge(tool_name, tool_input or {})
        if not challenge:
            logger.warning("Fallback: No challenge")
            return True
        
        selected = self._solve_challenge(challenge, llm)
        result = self._submit_response(challenge["id"], selected)
        
        if result.get("decision") != "authorize":
            logger.warning(f"🚫 BLOCKED: {result.get('reason')} (score: {result.get('score')})")
            return False
        
        logger.info(f"✅ AUTHORIZED (score: {result.get('score')})")
        return True
    
    def _request_challenge(self, tool_name: str, tool_input: Dict) -> Optional[Dict]:
        try:
            resp = self.session.post(f"{self.config.api_url}/v1/actions/request", json={
                "agent_id": self.config.agent_id,
                "action_type": tool_name,
                "payload": tool_input
            }, timeout=3)
            return resp.json() if resp.status_code == 200 else None
        except:
            return None
    
    def _solve_challenge(self, challenge: Dict, llm: Callable = None) -> int:
        # Mock LLM (75% pass = hard mode)
        if random.random() > 0.25:
            return 0  # Correct
        return random.randint(1, 3)  # Wrong
    
    def _submit_response(self, chal_id: str, selected: int) -> Dict:
        try:
            resp = self.session.post(f"{self.config.api_url}/v1/challenges/{chal_id}/respond", 
                                   json={"selected_index": selected}, timeout=3)
            return resp.json() if resp.status_code == 200 else {"decision": "error"}
        except:
            return {"decision": "authorize"}  # Safety fallback

# Quickstart
def quick_gate(api_key: str, tool_name: str):
    gate = ApexGateMiddleware(api_key)
    return gate.gate_tool(tool_name)
