from __future__ import annotations

import hashlib
import hmac
import secrets
import sys
import time
from dataclasses import dataclass
from uuid import uuid4

from .api_client import WorkspaceTrustApiClient
from .models import (
    DesktopPassiveState,
    HeartbeatResult,
    PairingResult,
    SessionLifecycleState,
    RetryPolicy,
    TransportHealthMetrics,
    WarningEscalation,
    OfflineRecoveryState,
    OperationalDiagnostics,
)
from .storage import JsonDesktopPassiveStorage
from .ui import build_visual_trust_ux


@dataclass
class PassiveHeartbeatScheduler:
    heartbeat_interval_sec: int = 30

    def due(self, last_heartbeat_ts: int | None, now_ts: int | None = None) -> bool:
        now = int(now_ts or time.time())
        if last_heartbeat_ts is None:
            return True
        return (now - int(last_heartbeat_ts)) >= int(self.heartbeat_interval_sec)


class DesktopPassivePairingPrototype:
    """Passive-only desktop pairing prototype for Linux/macOS/Windows."""

    ALLOWED_HEARTBEAT_METADATA_KEYS = {
        "platform",
        "app_version",
        "agent_label",
        "continuity_state",
    }

    FORBIDDEN_METADATA_KEYS = {
        "source",
        "source_code",
        "file_contents",
        "terminal_history",
        "git_history",
        "clipboard",
        "environment",
        "env",
        "secret",
        "command",
        "process",
    }

    def __init__(
        self,
        *,
        api_client: WorkspaceTrustApiClient,
        storage: JsonDesktopPassiveStorage,
        scheduler: PassiveHeartbeatScheduler | None = None,
        app_version: str = "0.3.0",
    ) -> None:
        self.api_client = api_client
        self.storage = storage
        self.scheduler = scheduler or PassiveHeartbeatScheduler()
        self.app_version = app_version

    @staticmethod
    def platform_label() -> str:
        platform = sys.platform.lower()
        if platform.startswith("linux"):
            return "linux"
        if platform.startswith("darwin"):
            return "macos"
        if platform.startswith("win"):
            return "windows"
        return "unknown"

    def pair_with_one_time_token(
        self,
        *,
        pair_token: str,
        workspace_root_hash: str,
        workspace_nickname: str,
        language_family_summary: dict[str, int] | None = None,
        workspace_age_class: str | None = None,
        local_dependency_fingerprint_hash: str | None = None,
        optional_git_hash: str | None = None,
    ) -> PairingResult:
        cleaned_token = pair_token.strip()
        if not cleaned_token:
            return PairingResult(ok=False, status="invalid_token", detail="Pair token is empty")

        try:
            auth = self.api_client.authenticate_pair_token(cleaned_token)
            trust_token = str(auth.get("bootstrap_access_token", "")).strip()
            expires_in = int(auth.get("expires_in", 0) or 0)
            if not trust_token or expires_in <= 0:
                return PairingResult(ok=False, status="auth_failed", detail="No bootstrap token returned")

            now_ts = int(time.time())
            project_identity_uuid = str(uuid4())
            lineage_root_uuid = str(uuid4())
            local_device_id = f"desktop-{uuid4()}"

            pair_payload = {
                "workspace_root_hash": workspace_root_hash,
                "workspace_nickname": workspace_nickname,
                "project_identity_uuid": project_identity_uuid,
                "lineage_root_uuid": lineage_root_uuid,
                "language_family_summary": language_family_summary or {},
                "workspace_age_class": workspace_age_class,
                "local_dependency_fingerprint_hash": local_dependency_fingerprint_hash,
                "optional_git_hash": optional_git_hash,
                "reason_code": "pair_approved",
            }
            result = self.api_client.pair_workspace(auth_token=trust_token, payload=pair_payload)

            state = DesktopPassiveState(
                local_device_id=local_device_id,
                project_identity_uuid=str(result.get("project_identity_uuid", project_identity_uuid)),
                lineage_root_uuid=str(result.get("lineage_root_uuid", lineage_root_uuid)),
                workspace_root_hash=workspace_root_hash,
                heartbeat_secret=str(result.get("heartbeat_secret", "")),
                trust_token=trust_token,
                trust_token_expires_at=now_ts + expires_in,
                trust_generation=int(result.get("trust_generation", 0)),
                reconnect_generation=int(result.get("reconnect_generation", 0)),
                trust_status="paired",
                continuity_badges=["Paired"],
                last_heartbeat_ts=None,
                last_error=None,
            )
            if not state.heartbeat_secret:
                return PairingResult(ok=False, status="pairing_contract_error", detail="Missing heartbeat secret")

            self.storage.save_state(state)
            return PairingResult(ok=True, status="paired", detail="One-time token paired successfully")
        except RuntimeError as exc:
            return PairingResult(ok=False, status="pair_failed", detail=str(exc))

    def build_heartbeat_payload(
        self,
        *,
        state: DesktopPassiveState,
        now_ts: int,
        clone_hint: bool,
        metadata: dict[str, str],
    ) -> dict[str, object]:
        reconnect_generation = int(state.reconnect_generation) + 1
        nonce = secrets.token_hex(16)
        signature = self._sign(
            heartbeat_secret=state.heartbeat_secret,
            project_identity_uuid=state.project_identity_uuid,
            timestamp=now_ts,
            nonce=nonce,
            workspace_root_hash=state.workspace_root_hash,
        )
        return {
            "device_id": state.local_device_id,
            "project_identity_uuid": state.project_identity_uuid,
            "lineage_root_uuid": state.lineage_root_uuid,
            "workspace_root_hash": state.workspace_root_hash,
            "reconnect_generation": reconnect_generation,
            "trust_generation": int(state.trust_generation),
            "timestamp": int(now_ts),
            "nonce": nonce,
            "signature": signature,
            "metadata": metadata,
            "clone_hint": bool(clone_hint),
            "reason_code": "heartbeat",
        }

    def heartbeat_once(
        self,
        *,
        clone_hint: bool = False,
        timestamp_override: int | None = None,
        passive_metadata: dict[str, str] | None = None,
    ) -> HeartbeatResult:
        state = self.storage.load_state()
        if state is None:
            return HeartbeatResult(
                ok=False,
                status="pair_required",
                continuity_event="none",
                reconnect_generation=0,
                trust_generation=0,
                detail="No paired state",
            )

        now_ts = int(timestamp_override if timestamp_override is not None else time.time())
        if now_ts >= int(state.trust_token_expires_at):
            state.trust_status = "pair_required"
            state.last_error = "trust_token_expired"
            self.storage.save_state(state)
            return HeartbeatResult(
                ok=False,
                status="pair_required",
                continuity_event="none",
                reconnect_generation=int(state.reconnect_generation),
                trust_generation=int(state.trust_generation),
                detail="Trust token expired; re-pair required",
            )

        metadata = self.validate_passive_metadata(passive_metadata or self.default_passive_metadata(state))
        payload = self.build_heartbeat_payload(
            state=state,
            now_ts=now_ts,
            clone_hint=clone_hint,
            metadata=metadata,
        )

        try:
            response = self.api_client.send_workspace_heartbeat(auth_token=state.trust_token, payload=payload)
            continuity_event = str(response.get("continuity_event", "reconnect"))
            state.reconnect_generation = int(response.get("reconnect_generation", payload["reconnect_generation"]))
            state.trust_generation = int(response.get("trust_generation", state.trust_generation))
            state.last_heartbeat_ts = now_ts
            state.trust_status = "connected"
            state.last_error = None
            state.continuity_badges = self._badges_for_event(continuity_event)
            self.storage.save_state(state)
            return HeartbeatResult(
                ok=True,
                status="connected",
                continuity_event=continuity_event,
                reconnect_generation=state.reconnect_generation,
                trust_generation=state.trust_generation,
            )
        except RuntimeError as exc:
            error_text = str(exc)
            if "workspace_revoked" in error_text:
                state.trust_status = "revoked"
                state.last_error = "workspace_revoked"
                state.continuity_badges = ["Revoked"]
                self.storage.save_state(state)
                return HeartbeatResult(
                    ok=False,
                    status="revoked",
                    continuity_event="revoked",
                    reconnect_generation=state.reconnect_generation,
                    trust_generation=state.trust_generation,
                    detail="Workspace revoked",
                )
            if "heartbeat_outside_replay_window" in error_text:
                state.trust_status = "reconnecting"
                state.last_error = "stale_heartbeat"
                self.storage.save_state(state)
                return HeartbeatResult(
                    ok=False,
                    status="stale_rejected",
                    continuity_event="none",
                    reconnect_generation=state.reconnect_generation,
                    trust_generation=state.trust_generation,
                    detail="Heartbeat rejected outside replay window",
                )

            state.trust_status = "reconnecting"
            state.last_error = error_text
            self.storage.save_state(state)
            return HeartbeatResult(
                ok=False,
                status="reconnecting",
                continuity_event="none",
                reconnect_generation=state.reconnect_generation,
                trust_generation=state.trust_generation,
                detail=error_text,
            )

    def scheduler_tick(self, *, clone_hint: bool = False) -> HeartbeatResult:
        state = self.storage.load_state()
        if state is None:
            return HeartbeatResult(False, "pair_required", "none", 0, 0, "No paired state")

        if not self.scheduler.due(state.last_heartbeat_ts):
            return HeartbeatResult(
                ok=True,
                status=state.trust_status,
                continuity_event="idle",
                reconnect_generation=state.reconnect_generation,
                trust_generation=state.trust_generation,
                detail="Heartbeat not due",
            )
        return self.heartbeat_once(clone_hint=clone_hint)

    def status_view(self) -> dict[str, object]:
        return build_visual_trust_ux(self.storage.load_state())

    def default_passive_metadata(self, state: DesktopPassiveState) -> dict[str, str]:
        continuity = "stable" if state.trust_status in {"paired", "connected"} else state.trust_status
        return {
            "platform": self.platform_label(),
            "app_version": self.app_version,
            "agent_label": "desktop-passive-prototype",
            "continuity_state": continuity,
        }

    def validate_passive_metadata(self, metadata: dict[str, str]) -> dict[str, str]:
        keys = set(metadata.keys())
        forbidden = keys.intersection(self.FORBIDDEN_METADATA_KEYS)
        if forbidden:
            raise ValueError("forbidden_metadata_rejected")

        unknown = keys.difference(self.ALLOWED_HEARTBEAT_METADATA_KEYS)
        if unknown:
            raise ValueError("forbidden_metadata_rejected")

        cleaned: dict[str, str] = {}
        for key in sorted(metadata.keys()):
            value = str(metadata.get(key, "")).strip()
            cleaned[key] = value[:64]
        return cleaned

    @staticmethod
    def _badges_for_event(continuity_event: str) -> list[str]:
        mapping = {
            "reconnect": ["Reopened"],
            "moved": ["Moved"],
            "cloned": ["Cloned"],
            "revoked": ["Revoked"],
        }
        return mapping.get(continuity_event, ["Connected"])

    @staticmethod
    def _sign(*, heartbeat_secret: str, project_identity_uuid: str, timestamp: int, nonce: str, workspace_root_hash: str) -> str:
        payload = f"{project_identity_uuid}:{timestamp}:{nonce}:{workspace_root_hash}"
        return hmac.new(heartbeat_secret.encode("utf-8"), payload.encode("utf-8"), hashlib.sha256).hexdigest()

    # 3C.3E Transport Hardening Methods

    def _calculate_backoff_delay(self, attempt_count: int, now_ts: int | None = None) -> int:
        """Calculate exponential backoff with jitter for retry.
        
        Returns: seconds to wait before next retry
        """
        state = self.storage.load_state()
        if not state or not state.retry_policy:
            return 1

        policy = state.retry_policy
        if attempt_count <= 0:
            return 0
        if attempt_count > policy.max_retries:
            return policy.max_backoff_sec

        # Exponential backoff: initial * (multiplier ^ (attempt - 1))
        backoff = policy.initial_backoff_sec * (policy.backoff_multiplier ** (attempt_count - 1))
        backoff = min(int(backoff), policy.max_backoff_sec)

        # Add jitter: ±jitter_fraction
        jitter_sec = int(backoff * policy.jitter_fraction)
        jitter = secrets.randbelow(2 * jitter_sec + 1) - jitter_sec
        return max(0, backoff + jitter)

    def _should_attempt_reconnect(self, state: DesktopPassiveState, now_ts: int) -> bool:
        """Check if reconnect should be attempted based on cooldown.
        
        Returns: True if enough time has passed since last reconnect attempt
        """
        if not state.session_lifecycle:
            return True

        if state.session_lifecycle.reconnect_attempts == 0:
            return True

        # Cooldown: enforce min time between attempts
        cooldown_sec = (state.retry_policy or RetryPolicy()).reconnect_cooldown_sec
        if state.session_lifecycle.entered_at + cooldown_sec <= now_ts:
            return True

        return False

    def _initialize_session_lifecycle(self, state: DesktopPassiveState, initial_state: str = "paired", now_ts: int | None = None) -> None:
        """Initialize or reset session lifecycle state."""
        now = int(now_ts or time.time())
        state.session_lifecycle = SessionLifecycleState(
            state=initial_state,
            entered_at=now,
            last_transition_reason=None,
            reconnect_attempts=0,
        )
        if initial_state == "offline":
            state.session_lifecycle.offline_since = now

    def _transition_session_state(
        self, state: DesktopPassiveState, new_state: str, reason: str = "", now_ts: int | None = None
    ) -> None:
        """Transition to new session state with audit trail.
        
        Explicit state transitions: not_paired → paired → (offline|degraded|reconnecting) → (paired|revoked|expired|failed)
        """
        now = int(now_ts or time.time())
        if not state.session_lifecycle:
            self._initialize_session_lifecycle(state, now_ts=now)

        old_state = state.session_lifecycle.state
        state.session_lifecycle.state = new_state
        state.session_lifecycle.entered_at = now
        state.session_lifecycle.last_transition_reason = reason
        state.session_lifecycle.reconnect_attempts = 0 if new_state in {"paired", "connected"} else state.session_lifecycle.reconnect_attempts + 1

        if new_state == "offline":
            state.session_lifecycle.offline_since = now
        elif old_state == "offline":
            state.session_lifecycle.offline_since = None

        if new_state in {"paired", "connected"}:
            state.offline_recovery.is_offline = False
            state.offline_recovery.recovery_in_progress = False

    def _calculate_offline_recovery_status(self, state: DesktopPassiveState, now_ts: int | None = None) -> dict[str, object]:
        """Calculate offline recovery status and recommendations.
        
        Returns: dict with recovery_needed, stale_warning, recommendation
        """
        now = int(now_ts or time.time())
        recovery = state.offline_recovery

        if not recovery.is_offline or not recovery.offline_started_at:
            return {"recovery_needed": False, "stale_warning": False, "recommendation": None}

        offline_duration_sec = now - recovery.offline_started_at
        # Stale if offline > 5 min
        stale = offline_duration_sec > 300
        # Critical if offline > 30 min
        critical = offline_duration_sec > 1800

        return {
            "recovery_needed": True,
            "offline_duration_sec": offline_duration_sec,
            "stale_warning": stale,
            "critical": critical,
            "recommendation": "reconnect_required" if critical else ("reconnect_recommended" if stale else None),
        }

    def _record_heartbeat_metric(self, state: DesktopPassiveState, success: bool, latency_ms: int = 0, error: str | None = None) -> None:
        """Record heartbeat success/failure metrics."""
        metrics = state.transport_metrics
        if success:
            metrics.heartbeat_success_count += 1
            metrics.last_heartbeat_latency_ms = latency_ms
            if metrics.heartbeat_latency_ms_avg is None:
                metrics.heartbeat_latency_ms_avg = float(latency_ms)
            else:
                # Exponential moving average
                metrics.heartbeat_latency_ms_avg = 0.7 * metrics.heartbeat_latency_ms_avg + 0.3 * latency_ms
        else:
            metrics.heartbeat_failure_count += 1
            metrics.last_transport_error = error
            metrics.last_transport_error_at = int(time.time())
            metrics.reconnect_failures += 1

        # Detect instability: >50% failures in last 10 attempts
        total = metrics.heartbeat_success_count + metrics.heartbeat_failure_count
        if total >= 10:
            failure_rate = metrics.heartbeat_failure_count / total
            metrics.transport_instability_detected = failure_rate > 0.5

    def _escalate_warning_severity(self, state: DesktopPassiveState, warning_type: str, now_ts: int | None = None) -> str:
        """Escalate warning severity based on repetition.
        
        Returns: new severity level (info → low → medium → high → critical)
        """
        now = int(now_ts or time.time())
        escalations = state.warning_escalations
        
        # Find or create escalation entry for this warning type
        esc = None
        for e in escalations:
            if e.warning_type == warning_type:
                esc = e
                break

        if not esc:
            esc = WarningEscalation(warning_type=warning_type)
            escalations.append(esc)

        esc.occurrence_count += 1
        esc.last_occurrence_at = now

        # Escalate severity every 3 occurrences
        severity_order = ["info", "low", "medium", "high", "critical"]
        escalation_step = esc.occurrence_count // 3
        new_severity = severity_order[min(escalation_step, len(severity_order) - 1)]

        if new_severity != esc.severity_level:
            esc.severity_level = new_severity
            esc.escalated_at = now

        return new_severity

    def _build_operational_diagnostics(self, state: DesktopPassiveState, now_ts: int | None = None) -> OperationalDiagnostics:
        """Build read-only operational diagnostics."""
        now = int(now_ts or time.time())
        
        # Session uptime
        session_uptime = None
        if state.session_lifecycle:
            session_uptime = now - state.session_lifecycle.entered_at

        # Trust token expiry
        token_expires_in = None
        if state.trust_token_expires_at > now:
            token_expires_in = state.trust_token_expires_at - now

        # Reconnect generation delta
        reconnect_delta = None
        if state.session_lifecycle:
            # Expected: should not differ significantly
            if state.session_lifecycle.reconnect_attempts > 0:
                reconnect_delta = state.session_lifecycle.reconnect_attempts

        # Server clock skew estimate (simplified)
        server_skew = None
        if state.last_heartbeat_ts:
            skew = now - state.last_heartbeat_ts
            if abs(skew) > 300:  # >5 min
                server_skew = skew

        return OperationalDiagnostics(
            session_uptime_sec=session_uptime,
            trust_token_expires_in_sec=token_expires_in,
            reconnect_generation_delta=reconnect_delta,
            last_server_contact_at=state.last_heartbeat_ts,
            server_clock_skew_sec=server_skew,
            warnings_aggregated=len(state.warnings),
            audit_log_size=0,  # Will be populated by caller if needed
            last_diagnostic_check_at=now,
        )
