"""Main RemediationService orchestrator."""

from __future__ import annotations

import asyncio
import base64
import contextlib
from dataclasses import dataclass, field
from datetime import UTC, datetime, timedelta
from typing import TYPE_CHECKING, Any

import structlog

from curadise_agent.auth.crypto import SignatureInfo, SignatureVerifier
from curadise_agent.errors import RemediationError
from curadise_agent.remediation.detector import Detector, create_detector
from curadise_agent.remediation.models import (
    DetectionResult,
    DetectionRule,
    Playbook,
    PlaybookBundle,
    RemediationResult,
    RuleState,
    StepResult,
)
from curadise_agent.remediation.security import SecurityError, execute_command_safe
from curadise_agent.remediation.stats import RemediationStats

if TYPE_CHECKING:
    from curadise_agent.remediation.config import RemediationConfig
    from curadise_agent.state.manager import StateManager
    from curadise_agent.transport.client import TransportClient

log = structlog.get_logger(__name__)

# Nonce expiration time (24 hours)
NONCE_EXPIRATION_HOURS = 24


@dataclass
class RemediationService:
    """
    Main orchestrator for auto-remediation.

    Coordinates detection, remediation execution, verification, and reporting.
    """

    config: RemediationConfig
    transport: TransportClient
    state_manager: StateManager
    verifier: SignatureVerifier | None = None

    # Internal components
    _detector: Detector | None = field(default=None, init=False)
    _stats: RemediationStats = field(default_factory=RemediationStats, init=False)

    # Rule and playbook storage
    _rules: dict[str, DetectionRule] = field(default_factory=dict, init=False)
    _playbooks: dict[str, Playbook] = field(default_factory=dict, init=False)
    _rule_to_playbook: dict[str, str] = field(default_factory=dict, init=False)
    _rule_states: dict[str, RuleState] = field(default_factory=dict, init=False)

    # Nonce tracking for replay prevention (nonce -> timestamp)
    _used_nonces: dict[str, datetime] = field(default_factory=dict, init=False)

    # Runtime state
    _running: bool = field(default=False, init=False)
    _check_task: asyncio.Task[None] | None = field(default=None, init=False)
    _sync_task: asyncio.Task[None] | None = field(default=None, init=False)
    _shutdown_event: asyncio.Event | None = field(default=None, init=False)
    _semaphore: asyncio.Semaphore | None = field(default=None, init=False)
    _last_remediation_time: datetime | None = field(default=None, init=False)

    def __post_init__(self) -> None:
        """Initialize detector, semaphore, and load public keys."""
        self._detector = create_detector(self.config)
        self._semaphore = asyncio.Semaphore(self.config.max_concurrent)

        # Load public keys from configuration
        self._load_public_keys()

    def _load_public_keys(self) -> None:
        """Load public keys from configuration into the verifier."""
        if self.verifier is None:
            return

        sig_config = self.config.signature

        # Load keys from configuration dictionary
        for key_id, pem_key in sig_config.public_keys.items():
            try:
                self.verifier.add_public_key(key_id, pem_key)
                log.info("loaded_public_key_from_config", key_id=key_id)
            except Exception as e:
                log.error(
                    "failed_to_load_public_key",
                    key_id=key_id,
                    error=str(e),
                )

        # Load keys from directory if specified
        if sig_config.public_key_directory:
            try:
                from curadise_agent.remediation.config import load_public_keys_from_directory

                dir_keys = load_public_keys_from_directory(
                    sig_config.public_key_directory
                )
                for key_id, pem_key in dir_keys.items():
                    try:
                        self.verifier.add_public_key(key_id, pem_key)
                    except Exception as e:
                        log.error(
                            "failed_to_add_public_key",
                            key_id=key_id,
                            error=str(e),
                        )
            except Exception as e:
                log.error(
                    "failed_to_load_keys_from_directory",
                    directory=str(sig_config.public_key_directory),
                    error=str(e),
                )

    @property
    def stats(self) -> RemediationStats:
        """Get remediation statistics."""
        return self._stats

    @property
    def is_running(self) -> bool:
        """Check if service is running."""
        return self._running

    @property
    def active_rules_count(self) -> int:
        """Get count of active rules."""
        return len([r for r in self._rules.values() if r.enabled])

    async def start(self, shutdown_event: asyncio.Event) -> None:
        """
        Start the remediation service.

        Args:
            shutdown_event: Event to signal shutdown
        """
        if not self.config.enabled:
            log.info("remediation_service_disabled")
            return

        self._shutdown_event = shutdown_event
        self._running = True

        # Load cached bundle from state if available
        await self._load_cached_bundle()

        # Start sync task
        self._sync_task = asyncio.create_task(
            self._sync_loop(),
            name="remediation-sync",
        )

        # Start check task
        self._check_task = asyncio.create_task(
            self._check_loop(),
            name="remediation-check",
        )

        log.info(
            "remediation_service_started",
            sync_interval=self.config.sync_interval,
            check_interval=self.config.check_interval,
            dry_run=self.config.dry_run,
        )

    async def stop(self) -> None:
        """Stop the remediation service."""
        self._running = False

        # Cancel tasks
        for task in [self._sync_task, self._check_task]:
            if task:
                task.cancel()
                with contextlib.suppress(asyncio.CancelledError):
                    await task

        log.info("remediation_service_stopped")

    async def _sync_loop(self) -> None:
        """Periodically sync rules and playbooks from server."""
        assert self._shutdown_event is not None

        while self._running and not self._shutdown_event.is_set():
            try:
                await self._sync_bundle()
            except Exception as e:
                log.error("bundle_sync_error", error=str(e))
                self._stats.record_bundle_sync(success=False)

            # Wait for next sync interval
            try:
                await asyncio.wait_for(
                    self._shutdown_event.wait(),
                    timeout=self.config.sync_interval,
                )
                break  # Shutdown requested
            except TimeoutError:
                continue

    async def _check_loop(self) -> None:
        """Periodically run detection checks."""
        assert self._shutdown_event is not None

        while self._running and not self._shutdown_event.is_set():
            try:
                await self._run_all_checks()
            except Exception as e:
                log.error("check_loop_error", error=str(e))

            # Wait for next check interval
            try:
                await asyncio.wait_for(
                    self._shutdown_event.wait(),
                    timeout=self.config.check_interval,
                )
                break  # Shutdown requested
            except TimeoutError:
                continue

    async def _sync_bundle(self) -> None:
        """Fetch and apply a new bundle from the server."""
        log.debug("syncing_remediation_bundle")

        try:
            bundle_data = await self.transport.fetch_remediation_bundle()

            if not bundle_data:
                log.debug("no_bundle_available")
                return

            bundle = PlaybookBundle.model_validate(bundle_data)

            # Verify signature if required
            if self.config.require_signature:
                if not self._verify_bundle_signature(bundle):
                    log.error("bundle_signature_invalid")
                    self._stats.record_signature_validation(valid=False)
                    return
                self._stats.record_signature_validation(valid=True)

            # Check nonce for replay prevention with timestamp-based expiration
            if not self._check_and_add_nonce(bundle.nonce):
                log.warning("bundle_nonce_reused", nonce=bundle.nonce)
                return

            # Apply bundle
            self._apply_bundle(bundle)

            # Cache bundle to state for offline operation
            await self._cache_bundle(bundle)

            self._stats.record_bundle_sync(success=True)
            self._stats.last_bundle_sync_time = datetime.now(UTC)

            log.info(
                "bundle_synced",
                version=bundle.version,
                rules_count=len(bundle.rules),
                playbooks_count=len(bundle.playbooks),
            )

        except Exception as e:
            self._stats.record_bundle_sync(success=False)
            raise RemediationError(f"Bundle sync failed: {e}", cause=e) from e

    def _verify_bundle_signature(self, bundle: PlaybookBundle) -> bool:
        """Verify the signature on a playbook bundle."""
        if self.verifier is None:
            log.warning("no_signature_verifier_configured")
            return False

        try:
            signature_bytes = base64.b64decode(bundle.signature)
            sig_info = SignatureInfo(
                algorithm=bundle.signature_algorithm,
                key_id=bundle.key_id,
                signature=signature_bytes,
            )

            data = bundle.get_signature_data()
            return self.verifier.verify(data, sig_info)

        except Exception as e:
            log.error("signature_verification_error", error=str(e))
            return False

    def _apply_bundle(self, bundle: PlaybookBundle) -> None:
        """Apply a bundle's rules and playbooks."""
        # Clear existing
        self._rules.clear()
        self._playbooks.clear()
        self._rule_to_playbook.clear()

        # Add rules
        for rule in bundle.rules:
            self._rules[rule.id] = rule
            # Initialize rule state if not exists
            if rule.id not in self._rule_states:
                self._rule_states[rule.id] = RuleState(rule_id=rule.id)

        # Add playbooks and map to rules
        for playbook in bundle.playbooks:
            self._playbooks[playbook.id] = playbook
            self._rule_to_playbook[playbook.rule_id] = playbook.id

        # Update stats
        self._stats.active_rules = len([r for r in self._rules.values() if r.enabled])
        self._stats.active_playbooks = len(self._playbooks)

    def _check_and_add_nonce(self, nonce: str) -> bool:
        """
        Check if nonce has been used and add it with timestamp if not.

        Uses timestamp-based expiration instead of size-based pruning.
        This ensures proper FIFO behavior and prevents replay attacks
        within the expiration window.

        Args:
            nonce: The nonce to check

        Returns:
            True if nonce is fresh (not used), False if already used
        """
        now = datetime.now(UTC)

        # Clean expired nonces first (older than NONCE_EXPIRATION_HOURS)
        cutoff = now - timedelta(hours=NONCE_EXPIRATION_HOURS)
        self._used_nonces = {
            n: t for n, t in self._used_nonces.items() if t > cutoff
        }

        # Check if nonce was already used
        if nonce in self._used_nonces:
            return False

        # Add nonce with current timestamp
        self._used_nonces[nonce] = now
        return True

    async def _load_cached_bundle(self) -> None:
        """Load cached bundle from state storage."""
        try:
            cached = self.state_manager.storage.get("remediation_bundle")
            if cached:
                bundle = PlaybookBundle.model_validate(cached)
                self._apply_bundle(bundle)
                log.info(
                    "cached_bundle_loaded",
                    version=bundle.version,
                    rules_count=len(bundle.rules),
                )
        except Exception as e:
            log.warning("failed_to_load_cached_bundle", error=str(e))

    async def _cache_bundle(self, bundle: PlaybookBundle) -> None:
        """Cache bundle to state storage."""
        try:
            self.state_manager.storage.set(
                "remediation_bundle",
                bundle.model_dump(mode="json"),
                secure=False,
            )
        except Exception as e:
            log.warning("failed_to_cache_bundle", error=str(e))

    async def _run_all_checks(self) -> None:
        """Run all enabled detection checks."""
        assert self._detector is not None

        now = datetime.now(UTC)
        self._stats.last_check_time = now

        # Reset hourly counter if needed
        if (
            self._stats.hour_start_time is None
            or (now - self._stats.hour_start_time).total_seconds() >= 3600
        ):
            self._stats.reset_hourly_counter(now)

        for rule in self._rules.values():
            if not rule.enabled:
                continue

            # Check if enough time has passed since last check
            state = self._rule_states.get(rule.id)
            if state and state.last_check_time:
                elapsed = (now - state.last_check_time).total_seconds()
                if elapsed < rule.interval:
                    continue

            # Run the check
            result = await self._detector.run_check(rule)

            # Update state
            if rule.id not in self._rule_states:
                self._rule_states[rule.id] = RuleState(rule_id=rule.id)
            self._rule_states[rule.id].last_check_time = now

            # Record stats
            self._stats.record_check(result.passed, result.duration_ms)

            if result.passed:
                # Reset consecutive failures on success
                self._rule_states[rule.id].consecutive_failures = 0
                log.debug(
                    "detection_check_passed",
                    rule_id=rule.id,
                    rule_name=rule.name,
                )
            else:
                # Increment consecutive failures
                self._rule_states[rule.id].consecutive_failures += 1
                failures = self._rule_states[rule.id].consecutive_failures

                log.warning(
                    "detection_check_failed",
                    rule_id=rule.id,
                    rule_name=rule.name,
                    message=result.message,
                    consecutive_failures=failures,
                )

                # Check if threshold reached for remediation
                if failures >= rule.failure_threshold:
                    await self._trigger_remediation(rule, result)

    async def _trigger_remediation(
        self, rule: DetectionRule, detection_result: DetectionResult
    ) -> None:
        """Trigger remediation for a failed rule."""
        # Find playbook for this rule
        playbook_id = self._rule_to_playbook.get(rule.id)
        if not playbook_id:
            log.warning("no_playbook_for_rule", rule_id=rule.id)
            return

        playbook = self._playbooks.get(playbook_id)
        if not playbook:
            log.warning("playbook_not_found", playbook_id=playbook_id)
            return

        state = self._rule_states[rule.id]

        # Check cooldown
        if state.last_remediation_time:
            elapsed = (datetime.now(UTC) - state.last_remediation_time).total_seconds()
            if elapsed < playbook.cooldown:
                log.debug(
                    "remediation_in_cooldown",
                    rule_id=rule.id,
                    cooldown_remaining=playbook.cooldown - elapsed,
                )
                self._stats.record_remediation(
                    success=False, duration_ms=0, skipped_cooldown=True
                )
                return

        # Check global cooldown
        if self._last_remediation_time:
            elapsed = (datetime.now(UTC) - self._last_remediation_time).total_seconds()
            if elapsed < self.config.global_cooldown:
                log.debug(
                    "global_cooldown_active",
                    cooldown_remaining=self.config.global_cooldown - elapsed,
                )
                self._stats.record_remediation(
                    success=False, duration_ms=0, skipped_cooldown=True
                )
                return

        # Check rate limit
        if self._stats.remediations_this_hour >= self.config.max_remediations_per_hour:
            log.warning(
                "rate_limit_reached",
                limit=self.config.max_remediations_per_hour,
            )
            self._stats.record_remediation(
                success=False, duration_ms=0, skipped_rate_limit=True
            )
            return

        # Check retry limit
        if state.remediation_count >= playbook.max_retries:
            log.warning(
                "max_retries_reached",
                rule_id=rule.id,
                max_retries=playbook.max_retries,
            )
            return

        # Execute remediation with semaphore (concurrency limit)
        assert self._semaphore is not None
        async with self._semaphore:
            await self._execute_remediation(rule, playbook, detection_result)

    async def _execute_remediation(
        self,
        rule: DetectionRule,
        playbook: Playbook,
        detection_result: DetectionResult,
    ) -> None:
        """Execute a remediation playbook."""
        assert self._detector is not None

        started_at = datetime.now(UTC)
        state = self._rule_states[rule.id]
        step_results: list[StepResult] = []
        success = True
        error_msg: str | None = None

        log.info(
            "remediation_started",
            rule_id=rule.id,
            playbook_id=playbook.id,
            playbook_name=playbook.name,
            dry_run=self.config.dry_run,
        )

        if self.config.dry_run:
            log.info(
                "remediation_dry_run",
                rule_id=rule.id,
                playbook_id=playbook.id,
                steps=[s.name for s in playbook.steps],
            )
            self._stats.record_remediation(
                success=True, duration_ms=0, dry_run=True
            )
            return

        # Execute each step
        for step in playbook.steps:
            step_result = await self._execute_step(step)
            step_results.append(step_result)

            if not step_result.success:
                if step.continue_on_error:
                    log.warning(
                        "step_failed_continuing",
                        step_name=step.name,
                        exit_code=step_result.exit_code,
                    )
                else:
                    success = False
                    error_msg = f"Step '{step.name}' failed: {step_result.stderr}"
                    break

        completed_at = datetime.now(UTC)
        duration_ms = (completed_at - started_at).total_seconds() * 1000

        # Verification
        verified = False
        verification_passed: bool | None = None

        if success and playbook.verify_after:
            verification_result = await self._detector.run_check(rule)
            verified = True
            verification_passed = verification_result.passed
            self._stats.record_verification(verification_passed)

            if not verification_passed:
                log.warning(
                    "verification_failed",
                    rule_id=rule.id,
                    message=verification_result.message,
                )

        # Update state
        state.last_remediation_time = datetime.now(UTC)
        state.remediation_count += 1
        self._last_remediation_time = state.last_remediation_time

        if verification_passed:
            state.consecutive_failures = 0  # Reset on successful remediation

        # Record stats
        self._stats.record_remediation(success=success, duration_ms=duration_ms)
        self._stats.last_remediation_time = completed_at

        # Create result
        result = RemediationResult(
            rule_id=rule.id,
            playbook_id=playbook.id,
            triggered_at=started_at,
            completed_at=completed_at,
            success=success,
            steps_executed=step_results,
            verified=verified,
            verification_passed=verification_passed,
            error=error_msg,
            dry_run=False,
        )

        # Report to server
        try:
            await self.transport.report_remediation(result.model_dump(mode="json"))
        except Exception as e:
            log.error("failed_to_report_remediation", error=str(e))

        log.info(
            "remediation_completed",
            rule_id=rule.id,
            playbook_id=playbook.id,
            success=success,
            verified=verified,
            verification_passed=verification_passed,
            duration_ms=duration_ms,
        )

    async def _execute_step(self, step: Any) -> StepResult:
        """Execute a single playbook step using secure command execution."""
        assert self._detector is not None

        started_at = datetime.now(UTC)
        command = step.command

        # Validate command against remediation allowlist
        if not self._detector.validate_command(command, is_remediation=True):
            return StepResult(
                step_name=step.name,
                command=command,
                exit_code=-1,
                stdout="",
                stderr=f"Command not in remediation allowlist: {command}",
                duration_ms=0,
                success=False,
            )

        try:
            # Use secure command execution (subprocess_exec, not shell)
            exit_code, stdout_str, stderr_str = await execute_command_safe(
                command=command,
                timeout=step.timeout,
            )

            completed_at = datetime.now(UTC)
            duration_ms = (completed_at - started_at).total_seconds() * 1000

            return StepResult(
                step_name=step.name,
                command=command,
                exit_code=exit_code,
                stdout=stdout_str[:10000],  # Truncate
                stderr=stderr_str[:5000],
                duration_ms=duration_ms,
                success=(exit_code == 0),
            )

        except SecurityError as e:
            return StepResult(
                step_name=step.name,
                command=command,
                exit_code=-1,
                stdout="",
                stderr=f"Command rejected by security validation: {e}",
                duration_ms=0,
                success=False,
            )

        except TimeoutError:
            return StepResult(
                step_name=step.name,
                command=command,
                exit_code=-1,
                stdout="",
                stderr=f"Command timed out after {step.timeout}s",
                duration_ms=step.timeout * 1000,
                success=False,
            )

        except Exception as e:
            return StepResult(
                step_name=step.name,
                command=command,
                exit_code=-1,
                stdout="",
                stderr=str(e),
                duration_ms=(datetime.now(UTC) - started_at).total_seconds() * 1000,
                success=False,
            )

    def get_status(self) -> dict[str, Any]:
        """Get service status for health checks."""
        return {
            "enabled": self.config.enabled,
            "running": self._running,
            "dry_run": self.config.dry_run,
            "active_rules": self.active_rules_count,
            "active_playbooks": len(self._playbooks),
            "stats": self._stats.to_dict(),
        }


def create_remediation_service(
    config: RemediationConfig,
    transport: TransportClient,
    state_manager: StateManager,
    verifier: SignatureVerifier | None = None,
) -> RemediationService:
    """
    Factory function to create a remediation service.

    Args:
        config: Remediation configuration
        transport: Transport client for server communication
        state_manager: State manager for persistence
        verifier: Optional signature verifier

    Returns:
        Configured RemediationService
    """
    return RemediationService(
        config=config,
        transport=transport,
        state_manager=state_manager,
        verifier=verifier,
    )
