"""
ServiceBridge SDK v2 — New Enterprise Session Protocol

Implements the Channel-based bidi stream session with:
- 8-state FSM lifecycle  
- Adaptive heartbeat (EWMA RTT)
- Credit-based flow control with dynamic window sizing
- Bidirectional selective replay on reconnect
- ConfigPush dynamic transport config
- Epoch fencing / GoawaySignal handling
- Symmetric with Go and Node SDKs
"""

from __future__ import annotations

import asyncio
import logging
import random
from collections import deque
from dataclasses import dataclass, field
from enum import Enum, auto
from typing import Dict, List, Optional

logger = logging.getLogger(__name__)


class SessionStateV2(Enum):
    """Enterprise 8-state FSM states."""
    CONNECTING = auto()
    HANDSHAKING = auto()
    READY = auto()
    ACTIVE = auto()
    SUSPENDED = auto()
    DRAINING = auto()
    FENCED = auto()
    CLOSED = auto()


class TransportMode(str, Enum):
    """Transport mode for worker connections."""
    DIRECT = "direct"
    PROXY = "proxy"


class FallbackPolicy(str, Enum):
    """Fallback policy when primary transport fails."""
    NONE = "none"
    FALLBACK_TO_PROXY = "fallback_to_proxy"
    FALLBACK_TO_DIRECT = "fallback_to_direct"
    FALLBACK_BEST_EFFORT = "fallback_best_effort"


@dataclass
class V2Config:
    """Configuration for the v2 worker session client.

    Symmetric with Go (V2Config) and Node (V2Config) SDKs.
    """
    server_address: str
    instance_id: str
    zone: str = ""
    transport_mode: TransportMode = TransportMode.DIRECT
    max_inflight: int = 128
    min_inflight: int = 1
    sdk_version: str = "python-v2.0.0"
    capabilities: List[str] = field(default_factory=list)
    tls_enabled: bool = False
    cert_pem: str = ""
    key_pem: str = ""
    ca_pem: str = ""

    def validate(self) -> None:
        """Validate the config. Raises ValueError on invalid values."""
        if not self.server_address:
            raise ValueError("servicebridge: server_address is required")
        if not self.instance_id:
            raise ValueError("servicebridge: instance_id is required")


@dataclass
class ResumeState:
    """Resume state for reconnect with bidirectional selective replay.

    Symmetric with Go (ResumeState) and Node (ResumeState) SDKs.
    """
    resume_token: str = ""
    epoch: int = 0
    last_received_seq: int = 0
    last_sent_seq: int = 0
    completed_command_ids: List[str] = field(default_factory=list)


@dataclass
class HelloAckV2:
    """HelloAck received from server on session establishment."""
    session_id: str
    resume_token: str
    epoch: int
    resumed: bool
    resume_from_seq: int
    replayed_commands: int
    reconciled_results: int
    heartbeat_interval_ms: int
    heartbeat_timeout_ms: int
    initial_permits: int
    max_permits: int
    effective_transport_mode: TransportMode
    resume_fail_reason: Optional[str] = None


@dataclass
class FunctionDeclarationV2:
    """Declares a function for registration (v2 Reconcile API)."""
    fn_name: str
    weight: int = 100
    input_schema_json: str = ""
    output_schema_json: str = ""
    allowed_callers: List[str] = field(default_factory=list)


@dataclass
class ConsumerGroupDeclarationV2:
    """Declares a consumer group for registration (v2 Reconcile API)."""
    name: str
    pattern: str
    mode: str = ""
    retry_policy_json: str = ""
    active: bool = True
    filter_expr: str = ""
    weight: int = 100


@dataclass
class HttpRouteDeclarationV2:
    """Declares an HTTP route for registration (v2 Reconcile API)."""
    method: str
    route_pattern: str
    request_schema_json: str = ""
    response_schema_json: str = ""
    allowed_callers: List[str] = field(default_factory=list)


@dataclass
class WorkflowDeclarationV2:
    """Declares a workflow for registration (v2 Reconcile API)."""
    name: str
    definition: str
    opts: str = ""


@dataclass
class JobDeclarationV2:
    """Declares a cron job for registration (v2 Reconcile API)."""
    cron_expr: str
    timezone: str = ""
    misfire_policy: str = ""
    target_type: str = ""
    target_ref: str = ""
    delay_ms: int = 0
    retry_policy_json: str = ""


@dataclass
class ReconcileRequestV2:
    """Declarative worker registration request (v2 Reconcile API).

    Symmetric with Go (ReconcileRequestV2) and Node (ReconcileRequestV2) SDKs.
    """
    instance_id: str
    endpoint: str = ""
    transport: str = ""
    zone: str = ""
    transport_mode: TransportMode = TransportMode.DIRECT
    functions: List[FunctionDeclarationV2] = field(default_factory=list)
    consumer_groups: List[ConsumerGroupDeclarationV2] = field(default_factory=list)
    http_routes: List[HttpRouteDeclarationV2] = field(default_factory=list)
    workflows: List[WorkflowDeclarationV2] = field(default_factory=list)
    jobs: List[JobDeclarationV2] = field(default_factory=list)
    cpu_percent: float = 0.0
    ram_mb: int = 0


@dataclass
class CircuitBreakerConfigV2:
    """Circuit breaker configuration."""
    failure_threshold: int = 5
    success_threshold: int = 3
    open_duration_ms: int = 30000
    half_open_max_requests: int = 2


@dataclass
class ZoneConfigV2:
    """Zone-aware routing configuration."""
    local_zone: str = ""
    prefer_zone: bool = False
    zone_weight_boost: int = 3


@dataclass
class ServiceTransportOverride:
    """Per-service transport config override."""
    mode: Optional[TransportMode] = None
    fallback_policy: Optional[FallbackPolicy] = None
    direct_timeout_ms: int = 0
    lb_policy: str = ""
    circuit_breaker: Optional[CircuitBreakerConfigV2] = None


@dataclass
class FunctionTransportOverride:
    """Per-function transport config override."""
    mode: Optional[TransportMode] = None
    fallback_policy: Optional[FallbackPolicy] = None
    timeout_ms: int = 0


@dataclass
class TransportConfigV2:
    """Dynamic transport configuration pushed from server via ConfigPush."""
    default_mode: Optional[TransportMode] = None
    service_overrides: Dict[str, ServiceTransportOverride] = field(default_factory=dict)
    function_overrides: Dict[str, FunctionTransportOverride] = field(default_factory=dict)
    default_circuit_breaker: Optional[CircuitBreakerConfigV2] = None
    zone: Optional[ZoneConfigV2] = None


@dataclass
class IssueCertificateRequestV2:
    """Certificate issuance request (auto mTLS).

    Symmetric with Go (IssueCertificateRequestV2) and Node (IssueCertificateRequestV2) SDKs.
    """
    instance_id: str
    csr_pem: str
    advertise_address: str = ""
    extra_sans: List[str] = field(default_factory=list)


@dataclass
class IssueCertificateResponseV2:
    """Certificate issuance response."""
    cert_pem: str
    ca_pem: str
    expires_at: int  # Unix ms


class AdaptiveHeartbeatV2:
    """Adaptive heartbeat controller (EWMA RTT-based interval adjustment).

    Symmetric with Go (AdaptiveHeartbeatV2) and Node (AdaptiveHeartbeatV2) SDKs.
    """

    def __init__(self, interval_ms: int, timeout_ms: int) -> None:
        self._interval_ms = interval_ms
        self._timeout_ms = timeout_ms
        self._miss_count = 0
        self._rtt_ewma: float = 0.0
        self._alpha = 0.3

    def next_interval_ms(self) -> int:
        base = self._interval_ms / 3
        if self._miss_count > 0:
            shift = min(self._miss_count, 4)
            base = base / (2 ** shift)
        elif self._rtt_ewma > 0 and self._rtt_ewma < 50:
            base = min(base * 2, 30_000)
        return int(max(2_000, min(30_000, base)))

    def on_pong(self, rtt_ms: float) -> None:
        self._rtt_ewma = self._alpha * rtt_ms + (1 - self._alpha) * self._rtt_ewma
        self._miss_count = 0

    def on_miss(self) -> int:
        self._miss_count += 1
        return self._miss_count

    def reset_miss(self) -> None:
        self._miss_count = 0


class FlowControlStateV2:
    """Credit-based flow control state with dynamic window sizing.

    Symmetric with Go (FlowControlState) and Node (FlowControlStateV2) SDKs.
    """

    def __init__(self, initial: int, min_w: int, max_w: int) -> None:
        self._permits = initial
        self._window = initial
        self._min_w = min_w
        self._max_w = max_w
        self._lock = asyncio.Lock()

    async def try_consume(self) -> bool:
        async with self._lock:
            if self._permits <= 0:
                return False
            self._permits -= 1
            return True

    async def release(self, n: int) -> None:
        async with self._lock:
            self._permits += n

    async def available(self) -> int:
        async with self._lock:
            return self._permits

    async def set_window(self, n: int) -> None:
        async with self._lock:
            self._window = max(self._min_w, min(self._max_w, n))
            self._permits = self._window


class BackoffV2:
    """Exponential backoff with full jitter for reconnect.

    Symmetric with Go (BackoffV2) and Node (BackoffV2) SDKs.
    """

    def __init__(self) -> None:
        self._attempt = 0
        self._base_ms = 100
        self._max_ms = 30_000
        self._circuit_fails = 0

    def next_ms(self) -> float:
        self._attempt += 1
        calculated = self._base_ms * (2 ** min(self._attempt, 10))
        capped = min(calculated, self._max_ms)
        jittered = random.random() * capped
        return max(self._base_ms, jittered)

    def reset(self) -> None:
        self._attempt = 0
        self._circuit_fails = 0

    def record_fail(self) -> None:
        self._circuit_fails += 1

    def is_circuit_open(self) -> bool:
        return self._circuit_fails >= 10

    def next(self) -> float:
        """Alias for next_ms() for symmetry with Go and Node SDKs."""
        return self.next_ms()


class PositionTrackerV2:
    """Tracks sequence numbers and completed command IDs for resume.

    Symmetric with Go (PositionTracker) and Node (PositionTrackerV2) SDKs.
    """

    def __init__(self) -> None:
        self.last_received_seq: int = 0
        self.last_sent_seq: int = 0
        self._completed_ids: deque[str] = deque(maxlen=512)

    def record_received(self, seq: int) -> None:
        if seq > self.last_received_seq:
            self.last_received_seq = seq

    def record_sent(self, seq: int) -> None:
        if seq > self.last_sent_seq:
            self.last_sent_seq = seq

    def mark_completed(self, command_id: str) -> None:
        self._completed_ids.append(command_id)

    def get_resume_state(self, token: str, epoch: int) -> ResumeState:
        return ResumeState(
            resume_token=token,
            epoch=epoch,
            last_received_seq=self.last_received_seq,
            last_sent_seq=self.last_sent_seq,
            completed_command_ids=list(self._completed_ids),
        )


class ConfigPushStateV2:
    """Manages dynamically pushed transport configuration.

    Symmetric with Go (ConfigPushState) and Node (ConfigPushStateV2) SDKs.
    """

    def __init__(self) -> None:
        self._default_mode: TransportMode = TransportMode.DIRECT
        self._service_overrides: Dict[str, ServiceTransportOverride] = {}
        self._function_overrides: Dict[str, FunctionTransportOverride] = {}

    def apply(self, config: TransportConfigV2) -> None:
        if config.default_mode is not None:
            self._default_mode = config.default_mode
        self._service_overrides.update(config.service_overrides)
        self._function_overrides.update(config.function_overrides)

    def resolve_mode(self, fn_name: str) -> TransportMode:
        if fn_name in self._function_overrides:
            override = self._function_overrides[fn_name]
            if override.mode is not None:
                return override.mode
        svc_name = fn_name.split("/")[0] if "/" in fn_name else fn_name
        if svc_name in self._service_overrides:
            override = self._service_overrides[svc_name]
            if override.mode is not None:
                return override.mode
        return self._default_mode


class V2SessionClient:
    """Worker session client implementing the new Enterprise protocol.

    Symmetric with Go (V2SessionClient) and Node (V2SessionClient) SDKs.
    """

    def __init__(self, config: V2Config) -> None:
        config.validate()
        self.config = config
        self._state = SessionStateV2.CONNECTING
        self._resume_token = ""
        self._epoch = 0
        self._lock = asyncio.Lock()

        self.position = PositionTrackerV2()
        self.flow_ctrl = FlowControlStateV2(
            config.max_inflight,
            config.min_inflight,
            config.max_inflight,
        )
        self.heartbeat = AdaptiveHeartbeatV2(10_000, 30_000)
        self.backoff = BackoffV2()
        self.config_push = ConfigPushStateV2()

        self._draining = False
        self._stopped = False

    @property
    def state(self) -> SessionStateV2:
        return self._state

    @property
    def is_draining(self) -> bool:
        return self._draining

    @property
    def is_stopped(self) -> bool:
        return self._stopped

    def get_hello_fields(self) -> dict:
        """Returns Hello fields for initial connection or resume.

        Symmetric with Go (GetResumeState) and Node (getHelloFields) SDKs.
        """
        rs = self.position.get_resume_state(self._resume_token, self._epoch)
        return {
            "identity": {
                "instance_id": self.config.instance_id,
                "transport": self.config.transport_mode.value,
            },
            "max_inflight": self.config.max_inflight,
            "min_inflight": self.config.min_inflight,
            "resume_token": rs.resume_token,
            "epoch": rs.epoch,
            "last_received_seq": rs.last_received_seq,
            "last_sent_seq": rs.last_sent_seq,
            "completed_command_ids": rs.completed_command_ids,
            "sdk_version": self.config.sdk_version,
            "capabilities": self.config.capabilities,
            "zone": self.config.zone,
            "transport_mode": self.config.transport_mode.value,
        }

    async def on_hello_ack(self, ack: HelloAckV2) -> None:
        """Process HelloAck from server."""
        async with self._lock:
            self._resume_token = ack.resume_token
            self._epoch = ack.epoch
            self._state = SessionStateV2.READY
        await self.flow_ctrl.set_window(ack.initial_permits)
        self.heartbeat = AdaptiveHeartbeatV2(ack.heartbeat_interval_ms, ack.heartbeat_timeout_ms)
        self.backoff.reset()
        logger.info(
            "session-v2: hello_ack session=%s resumed=%s epoch=%d",
            ack.session_id, ack.resumed, ack.epoch,
        )

    async def on_command_received(self, seq: int, command_id: str) -> bool:
        """Process an incoming command. Returns False if backpressure."""
        self.position.record_received(seq)
        if not await self.flow_ctrl.try_consume():
            logger.warning("session-v2: no permits for command %s, backpressure", command_id)
            return False
        async with self._lock:
            if self._state == SessionStateV2.READY:
                self._state = SessionStateV2.ACTIVE
        return True

    async def on_command_completed(self, seq: int, command_id: str) -> None:
        """Process command completion — releases 1 permit."""
        self.position.mark_completed(command_id)
        self.position.record_sent(seq)
        await self.flow_ctrl.release(1)

    async def on_permit_grant(self, additional: int) -> None:
        """Handle PermitGrant from server."""
        await self.flow_ctrl.release(additional)

    async def on_flow_control_update(self, new_window_size: int, reason: str) -> None:
        """Handle FlowControlUpdate from server."""
        await self.flow_ctrl.set_window(new_window_size)
        logger.info("session-v2: flow_control_update window=%d reason=%s", new_window_size, reason)

    def on_pong(self, rtt_ms: float) -> None:
        """Handle pong response."""
        self.heartbeat.on_pong(rtt_ms)
        if self._state == SessionStateV2.SUSPENDED:
            self._state = SessionStateV2.READY
            logger.info("session-v2: heartbeat resumed, state→READY")

    def on_heartbeat_miss(self) -> bool:
        """Handle missed pong. Returns True if state→SUSPENDED."""
        count = self.heartbeat.on_miss()
        if count >= 2 and self._state in (SessionStateV2.READY, SessionStateV2.ACTIVE):
            self._state = SessionStateV2.SUSPENDED
            logger.warning("session-v2: heartbeat miss %d, state→SUSPENDED", count)
            return True
        return False

    def on_drain(self, reason: str, deadline_ms: int) -> None:
        """Initiate graceful drain (worker-initiated shutdown)."""
        if not self._draining:
            self._draining = True
            self._state = SessionStateV2.DRAINING
            logger.info("session-v2: drain initiated reason=%s deadline=%dms", reason, deadline_ms)

    def on_goaway(self, code: str, reason: str) -> None:
        """Handle GoawaySignal from server."""
        logger.info("session-v2: goaway code=%s reason=%s", code, reason)
        if code == "GOAWAY_FENCED":
            self._state = SessionStateV2.FENCED
            self.stop()
        elif code == "GOAWAY_GRACEFUL_SHUTDOWN":
            self._state = SessionStateV2.CLOSED
        else:
            self._state = SessionStateV2.CLOSED

    def on_config_push(self, config: TransportConfigV2) -> None:
        """Apply a ConfigPush from server."""
        self.config_push.apply(config)
        logger.info("session-v2: config push applied")

    def resolve_transport_mode(self, fn_name: str) -> TransportMode:
        """Resolve the effective transport mode for a function."""
        return self.config_push.resolve_mode(fn_name)

    def stop(self) -> None:
        """Stop the session."""
        self._stopped = True
        self._state = SessionStateV2.CLOSED


__all__ = [
    "V2Config",
    "V2SessionClient",
    "SessionStateV2",
    "TransportMode",
    "FallbackPolicy",
    "ResumeState",
    "HelloAckV2",
    "ReconcileRequestV2",
    "FunctionDeclarationV2",
    "ConsumerGroupDeclarationV2",
    "HttpRouteDeclarationV2",
    "WorkflowDeclarationV2",
    "JobDeclarationV2",
    "TransportConfigV2",
    "ServiceTransportOverride",
    "FunctionTransportOverride",
    "CircuitBreakerConfigV2",
    "ZoneConfigV2",
    "IssueCertificateRequestV2",
    "IssueCertificateResponseV2",
    "AdaptiveHeartbeatV2",
    "FlowControlStateV2",
    "BackoffV2",
    "PositionTrackerV2",
    "ConfigPushStateV2",
]
