"""
ACT Kernel - Main execution kernel class.
"""

import concurrent.futures
import threading
import time
import uuid
from typing import Any, Callable, Dict, Optional, Tuple

from .protocol import (
    ACTError,
    ACTPermissionError,
    BreakerState,
    CircuitOpenError,
    Envelope,
    ExecutionInfo,
    RetryExhaustedError,
    Status,
    TimeoutError,
    ToolError,
    ToolUnavailableError,
    error_code_from_exception,
)


class TokenBucket:
    """Thread-safe token bucket for rate limiting."""

    def __init__(self, rate: float, capacity: int):
        self.rate = rate
        self.capacity = capacity
        self._tokens = capacity
        self._last_refill = time.monotonic()
        self._lock = threading.Lock()

    def _refill(self):
        """Refill tokens based on elapsed time."""
        now = time.monotonic()
        elapsed = now - self._last_refill
        new_tokens = elapsed * self.rate
        if new_tokens > 0:
            self._tokens = min(self.capacity, self._tokens + new_tokens)
            self._last_refill = now

    def consume(self, tokens: int = 1) -> bool:
        """Try to consume tokens. Returns True if successful."""
        with self._lock:
            self._refill()
            if self._tokens >= tokens:
                self._tokens -= tokens
                return True
            return False

    @property
    def available(self) -> int:
        """Get the number of available tokens."""
        with self._lock:
            self._refill()
            return int(self._tokens)


class QuotaManager:
    """Manages quota/rate limits for tools."""

    def __init__(self, default_rate: float = 10.0, default_capacity: int = 20):
        self._buckets: Dict[Tuple[str, str], TokenBucket] = {}
        self._default_rate = default_rate
        self._default_capacity = default_capacity
        self._lock = threading.Lock()

    def set_limit(self, tool: str, action: str, rate: float, capacity: Optional[int] = None):
        """Set a custom rate limit for a tool/action."""
        if capacity is None:
            capacity = int(rate * 2)
        with self._lock:
            self._buckets[(tool, action)] = TokenBucket(rate, capacity)

    def acquire(self, tool: str, action: str, tokens: int = 1) -> Tuple[bool, int]:
        """Try to acquire quota. Returns (success, remaining_tokens)."""
        key = (tool, action)
        bucket = self._buckets.get(key)
        if bucket is None:
            with self._lock:
                if key not in self._buckets:
                    self._buckets[key] = TokenBucket(self._default_rate, self._default_capacity)
                bucket = self._buckets[key]
        success = bucket.consume(tokens)
        remaining = bucket.available
        return success, remaining


class CircuitBreaker:
    """Circuit breaker to prevent cascading failures."""

    def __init__(self, failure_threshold: int = 5, recovery_timeout: float = 30.0):
        self.failure_threshold = failure_threshold
        self.recovery_timeout = recovery_timeout
        self._state = BreakerState.CLOSED
        self._failure_count = 0
        self._last_failure_time = 0.0
        self._lock = threading.Lock()

    @property
    def state(self) -> BreakerState:
        """Get current circuit breaker state."""
        return self._state

    def call(self, func: Callable, *args, **kwargs) -> Any:
        """Call a function with circuit breaker protection."""
        with self._lock:
            if self._state == BreakerState.OPEN:
                if time.monotonic() - self._last_failure_time >= self.recovery_timeout:
                    self._state = BreakerState.HALF_OPEN
                else:
                    raise CircuitOpenError("Circuit breaker is open")

        try:
            result = func(*args, **kwargs)
        except Exception as e:
            with self._lock:
                self._failure_count += 1
                self._last_failure_time = time.monotonic()
                if self._state == BreakerState.HALF_OPEN or self._failure_count >= self.failure_threshold:
                    self._state = BreakerState.OPEN
            raise e

        with self._lock:
            if self._state == BreakerState.HALF_OPEN:
                self._state = BreakerState.CLOSED
                self._failure_count = 0
            else:
                self._failure_count = max(0, self._failure_count - 1)
        return result


class Registry:
    """Registry for tool handlers and metadata."""

    def __init__(self):
        self._handlers: Dict[Tuple[str, str], Callable] = {}
        self._metadata: Dict[Tuple[str, str], Dict] = {}

    def register(self, tool: str, action: str, handler: Callable, metadata: Optional[Dict] = None):
        """Register a tool handler with optional metadata."""
        self._handlers[(tool, action)] = handler
        self._metadata[(tool, action)] = metadata or {}

    def get(self, tool: str, action: str) -> Callable:
        """Get a registered tool handler."""
        return self._handlers[(tool, action)]

    def get_metadata(self, tool: str, action: str) -> Dict:
        """Get metadata for a registered tool."""
        return self._metadata.get((tool, action), {})


class Router:
    """Prefix-based intent router."""

    def __init__(self):
        self._rules: List[Tuple[str, str, str]] = []

    def add_rule(self, pattern: str, tool: str, action: str):
        """Add a routing rule."""
        self._rules.append((pattern, tool, action))

    def route(self, intent: str) -> Optional[Tuple[str, str]]:
        """Route an intent to a tool/action pair."""
        for pattern, tool, action in self._rules:
            if intent.startswith(pattern):
                return tool, action
        return None


class TimeoutExecutor:
    """Shared thread pool for timeout execution."""

    def __init__(self, max_workers: int = 10):
        self.executor = concurrent.futures.ThreadPoolExecutor(max_workers=max_workers)
        self._lock = threading.Lock()

    def call(self, func: Callable, timeout_sec: float) -> Any:
        """Execute a function with timeout."""
        future = self.executor.submit(func)
        try:
            return future.result(timeout=timeout_sec)
        except concurrent.futures.TimeoutError:
            raise TimeoutError(f"Execution exceeded timeout {timeout_sec}s")

    def shutdown(self, wait: bool = True):
        """Shutdown the executor."""
        self.executor.shutdown(wait=wait)

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.shutdown(wait=False)


def retry_with_count(
    func: Callable,
    retries: int = 2,
    delay: float = 0.0,
    backoff: float = 1.0,
    exceptions: Tuple[type, ...] = (Exception,),
) -> Tuple[Any, int]:
    """Retry function with count. Returns (result, actual_retries)."""
    last_exception = None
    current_delay = delay
    attempt = 0
    for attempt in range(retries + 1):
        try:
            result = func()
            return result, attempt
        except exceptions as e:
            last_exception = e
            if attempt < retries:
                time.sleep(current_delay)
                current_delay *= backoff
    raise RetryExhaustedError(f"Retry exhausted after {retries} attempts") from last_exception


class ACTKernel:
    """Main ACT Kernel class."""

    def __init__(
        self,
        default_timeout_ms: int = 1000,
        default_retries: int = 1,
        executor_max_workers: int = 10,
    ):
        self.registry = Registry()
        self.router = Router()
        self.quota_manager = QuotaManager()
        self.circuit_breakers: Dict[Tuple[str, str], CircuitBreaker] = {}
        self.default_timeout_ms = default_timeout_ms
        self.default_retries = default_retries
        self._lock = threading.Lock()
        self._executor = TimeoutExecutor(max_workers=executor_max_workers)

    def register_tool(
        self,
        tool: str,
        action: str,
        handler: Callable,
        metadata: Optional[Dict] = None,
    ):
        """Register a tool with the kernel."""
        self.registry.register(tool, action, handler, metadata)
        with self._lock:
            self.circuit_breakers[(tool, action)] = CircuitBreaker(
                failure_threshold=metadata.get("failure_threshold", 5) if metadata else 5,
                recovery_timeout=metadata.get("recovery_timeout", 30.0) if metadata else 30.0,
            )
        if metadata and "rate_limit" in metadata:
            rate = metadata["rate_limit"]
            capacity = metadata.get("capacity", int(rate * 2))
            self.quota_manager.set_limit(tool, action, rate, capacity)

    def add_route(self, pattern: str, tool: str, action: str):
        """Add a routing rule."""
        self.router.add_rule(pattern, tool, action)

    def dispatch(
        self,
        intent: str,
        params: Optional[Dict] = None,
        context: Optional[Dict] = None,
    ) -> Dict[str, Any]:
        """Dispatch an intent to execute a tool."""
        t_start = time.monotonic()
        request_id = (context or {}).get("request_id", str(uuid.uuid4()))

        envelope = Envelope(
            intent={"raw_intent": intent, "params": params, "request_id": request_id}
        )
        envelope.execution.queue_ms = 0

        route_result = self.router.route(intent)
        if route_result is None:
            envelope.execution.status = Status.ERROR
            envelope.error = {
                "code": "routing.not_found",
                "message": f"No route for intent: {intent}",
            }
            return envelope.to_dict()

        tool, action = route_result
        envelope.intent["tool"] = tool
        envelope.intent["action"] = action

        metadata = self.registry.get_metadata(tool, action)
        timeout_ms = metadata.get("timeout_ms", self.default_timeout_ms)
        retries = metadata.get("retries", self.default_retries)

        quota_success, quota_remaining = self.quota_manager.acquire(tool, action, tokens=1)
        envelope.execution.quota_requested = 1
        envelope.execution.quota_remaining = quota_remaining
        if not quota_success:
            envelope.execution.status = Status.REJECTED
            envelope.error = {"code": "quota.exceeded", "message": "Rate limit exceeded"}
            envelope.execution.quota_consumed = 0
            envelope.execution.latency_ms = int((time.monotonic() - t_start) * 1000)
            return envelope.to_dict()
        envelope.execution.quota_consumed = 1

        try:
            handler = self.registry.get(tool, action)
        except KeyError:
            envelope.execution.status = Status.ERROR
            envelope.error = {
                "code": "tool.not_found",
                "message": f"Tool {tool}/{action} not registered",
            }
            envelope.execution.latency_ms = int((time.monotonic() - t_start) * 1000)
            return envelope.to_dict()

        def task():
            return handler(params or {})

        breaker = self.circuit_breakers.get((tool, action))
        if not breaker:
            with self._lock:
                breaker = self.circuit_breakers.setdefault((tool, action), CircuitBreaker())

        actual_retries = 0
        try:

            def call_with_controls():
                nonlocal actual_retries

                def timeout_task():
                    return self._executor.call(task, timeout_ms / 1000.0)

                if retries > 0:
                    retry_exceptions = (TimeoutError, ToolUnavailableError)
                    result, actual_retries = retry_with_count(
                        timeout_task, retries=retries, exceptions=retry_exceptions
                    )
                    return result
                else:
                    actual_retries = 0
                    return timeout_task()

            result = breaker.call(call_with_controls)

            envelope.execution.status = Status.OK
            envelope.payload = result
        except TimeoutError as e:
            envelope.execution.status = Status.TIMEOUT
            envelope.error = {"code": error_code_from_exception(e), "message": str(e)}
        except CircuitOpenError as e:
            envelope.execution.status = Status.REJECTED
            envelope.error = {"code": error_code_from_exception(e), "message": str(e)}
        except RetryExhaustedError as e:
            envelope.execution.status = Status.ERROR
            envelope.error = {"code": error_code_from_exception(e), "message": str(e)}
        except ACTPermissionError as e:
            envelope.execution.status = Status.REJECTED
            envelope.error = {"code": error_code_from_exception(e), "message": str(e)}
        except ToolError as e:
            envelope.execution.status = Status.ERROR
            envelope.error = {"code": error_code_from_exception(e), "message": str(e)}
        except (ValueError, TypeError) as e:
            envelope.execution.status = Status.ERROR
            envelope.error = {"code": error_code_from_exception(e), "message": str(e)}
        except Exception as e:
            envelope.execution.status = Status.ERROR
            envelope.error = {"code": error_code_from_exception(e), "message": str(e)}
        finally:
            envelope.execution.latency_ms = int((time.monotonic() - t_start) * 1000)
            envelope.execution.breaker_state = breaker.state
            envelope.execution.timeout_ms = timeout_ms
            envelope.execution.retries = actual_retries

        return envelope.to_dict()

    def shutdown(self):
        """Shutdown the kernel."""
        self._executor.shutdown(wait=False)

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.shutdown()


class Middleware:
    """Middleware base class (v1.2 placeholder)."""

    def __call__(self, ctx: Dict, next: Callable) -> Any:
        result = next(ctx)
        return result
