"""Main event loop orchestrator for the Curadise agent."""

from __future__ import annotations

import asyncio
import contextlib
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any

import structlog

from curadise_agent import __version__
from curadise_agent.auth.jwt import TokenManager
from curadise_agent.auth.registration import RegistrationService
from curadise_agent.collectors.manager import CollectionManager
from curadise_agent.collectors.registry import (
    CollectorRegistry,
    create_collectors_from_config,
    register_builtin_collectors,
)
from curadise_agent.collectors.self_metrics import SelfMetricsCollector
from curadise_agent.config.watcher import ConfigWatcher, can_hot_reload
from curadise_agent.executor.runner import CommandRunner, create_command_runner
from curadise_agent.heartbeat.health import HealthChecker
from curadise_agent.heartbeat.service import HeartbeatService
from curadise_agent.models.api.command import CommandRequest
from curadise_agent.state.manager import AgentState, StateManager, create_state_manager
from curadise_agent.transport.batch import BatchBuilder, BatchConfig
from curadise_agent.transport.buffer import RingBuffer
from curadise_agent.transport.client import TransportClient, create_transport_client
from curadise_agent.utils.shutdown import ShutdownManager, get_shutdown_manager
from curadise_agent.utils.system_info import SystemInfo, collect_system_info

if TYPE_CHECKING:
    from pathlib import Path

    from curadise_agent.config.schema import AgentConfig
    from curadise_agent.models.domain.metric import DomainMetric

log = structlog.get_logger(__name__)


@dataclass
class Agent:
    """
    Main agent orchestrator.

    Coordinates all agent components and manages the main event loop.
    """

    config: AgentConfig
    config_path: Path | None = None

    # Core components
    _state_manager: StateManager | None = field(default=None, init=False)
    _transport: TransportClient | None = field(default=None, init=False)
    _buffer: RingBuffer[DomainMetric] | None = field(default=None, init=False)
    _batch_builder: BatchBuilder | None = field(default=None, init=False)

    # Services
    _registry: CollectorRegistry | None = field(default=None, init=False)
    _collection_manager: CollectionManager | None = field(default=None, init=False)
    _heartbeat_service: HeartbeatService | None = field(default=None, init=False)
    _command_runner: CommandRunner | None = field(default=None, init=False)
    _config_watcher: ConfigWatcher | None = field(default=None, init=False)

    # Helpers
    _shutdown_manager: ShutdownManager | None = field(default=None, init=False)
    _health_checker: HealthChecker | None = field(default=None, init=False)
    _token_manager: TokenManager | None = field(default=None, init=False)
    _system_info: SystemInfo | None = field(default=None, init=False)
    _self_collector: SelfMetricsCollector | None = field(default=None, init=False)

    # Tasks
    _flush_task: asyncio.Task[None] | None = field(default=None, init=False)
    _command_poll_task: asyncio.Task[None] | None = field(default=None, init=False)

    async def initialize(self) -> None:
        """Initialize all agent components."""
        log.info("agent_initializing", version=__version__)

        # Collect system info
        self._system_info = collect_system_info(extra={"agent_version": __version__})

        # Initialize state manager
        self._state_manager = create_state_manager(self.config.state)
        await self._state_manager.load()

        # Initialize transport
        self._transport = create_transport_client(
            self.config.server,
            self.config.auth,
            self.config.retry,
        )
        await self._transport.start()

        # Initialize buffer
        persistence_path = None
        if self.config.buffer.persistence_path:
            persistence_path = self.config.buffer.persistence_path
        elif self.config.state.directory:
            persistence_path = self.config.state.directory / self.config.state.buffer_file

        self._buffer = RingBuffer(
            max_size=self.config.buffer.max_size,
            persistence_path=persistence_path,
        )
        await self._buffer.restore()

        # Initialize batch builder
        self._batch_builder = BatchBuilder(
            config=BatchConfig(
                max_size=self.config.buffer.batch_size,
                max_age_seconds=self.config.buffer.flush_interval,
                compression_enabled=self.config.buffer.compression_enabled,
                compression_level=self.config.buffer.compression_level,
            )
        )

        # Initialize token manager
        self._token_manager = TokenManager()
        if self._state_manager.token:
            self._token_manager.set_tokens(self._state_manager.token)
            self._transport.set_auth_token(self._state_manager.token)

        # Initialize collectors
        self._registry = CollectorRegistry()
        register_builtin_collectors(self._registry)
        create_collectors_from_config(self.config.collectors, self._registry, self.config.tags)

        # Initialize collection manager
        self._collection_manager = CollectionManager(
            registry=self._registry,
            on_metrics=self._handle_metrics,
        )

        # Initialize command runner
        self._command_runner = create_command_runner(self.config.executor)

        # Initialize health checker
        self._health_checker = HealthChecker(
            state_manager=self._state_manager,
            transport=self._transport,
            buffer=self._buffer,
            collection_manager=self._collection_manager,
            command_runner=self._command_runner,
        )

        # Initialize heartbeat service
        self._heartbeat_service = HeartbeatService(
            transport=self._transport,
            state_manager=self._state_manager,
            health_checker=self._health_checker,
            config=self.config.heartbeat,
            system_info=self._system_info,
            tags=self.config.tags,
            on_commands_pending=self._fetch_and_execute_commands,
        )

        # Initialize self-metrics collector
        self._self_collector = SelfMetricsCollector(
            buffer=self._buffer,
            transport=self._transport,
            state_manager=self._state_manager,
        )
        self._registry._collectors["self"] = self._self_collector

        # Initialize config watcher
        if self.config_path:
            self._config_watcher = ConfigWatcher(
                config_path=self.config_path,
                on_change=self._handle_config_change,
            )

        # Initialize shutdown manager
        self._shutdown_manager = get_shutdown_manager()
        self._shutdown_manager.register_callback("buffer_flush", self._flush_buffer)
        self._shutdown_manager.register_callback("transport_stop", self._stop_transport)
        self._shutdown_manager.register_callback("state_save", self._save_state)

        log.info("agent_initialized")

    async def run(self) -> None:
        """Run the main agent loop."""
        assert self._state_manager is not None
        assert self._shutdown_manager is not None

        async with self._shutdown_manager.managed_lifecycle():
            # Check registration
            if not self._state_manager.is_registered:
                await self._register()

            # Transition to running state
            self._state_manager.transition_to(AgentState.RUNNING)

            # Start services
            await self._start_services()

            log.info("agent_running", agent_id=self._state_manager.agent_id)

            # Wait for shutdown
            await self._shutdown_manager.wait_for_shutdown()

            log.info("agent_shutting_down")
            self._state_manager.transition_to(AgentState.SHUTTING_DOWN)

            # Stop services
            await self._stop_services()

            self._state_manager.transition_to(AgentState.STOPPED)
            log.info("agent_stopped")

    async def _register(self) -> None:
        """Register the agent with the server."""
        assert self._transport is not None
        assert self._state_manager is not None
        assert self._system_info is not None

        enrollment_token = None
        if self.config.auth.enrollment_token:
            enrollment_token = self.config.auth.enrollment_token.get_secret_value()

        if not enrollment_token:
            log.error("no_enrollment_token")
            raise RuntimeError("No enrollment token configured")

        registration = RegistrationService(self._transport, self._state_manager)
        result = await registration.register(
            enrollment_token=enrollment_token,
            system_info=self._system_info.to_dict(),
            tags=self.config.tags,
        )

        # Update token manager
        if self._token_manager:
            self._token_manager.set_tokens(result.access_token, result.refresh_token)

    async def _start_services(self) -> None:
        """Start all agent services."""
        assert self._shutdown_manager is not None

        # Start collection manager
        if self._collection_manager:
            await self._collection_manager.start()

        # Start heartbeat service
        if self._heartbeat_service:
            await self._heartbeat_service.start(self._shutdown_manager.shutdown_event)

        # Start config watcher
        if self._config_watcher:
            await self._config_watcher.start(self.config)

        # Start buffer flush task
        self._flush_task = asyncio.create_task(
            self._flush_loop(),
            name="buffer-flush",
        )

        # Start command poll task
        self._command_poll_task = asyncio.create_task(
            self._command_poll_loop(),
            name="command-poll",
        )

    async def _stop_services(self) -> None:
        """Stop all agent services."""
        # Cancel tasks
        for task in [self._flush_task, self._command_poll_task]:
            if task:
                task.cancel()
                with contextlib.suppress(asyncio.CancelledError):
                    await task

        # Stop services
        if self._config_watcher:
            await self._config_watcher.stop()

        if self._heartbeat_service:
            await self._heartbeat_service.stop()

        if self._collection_manager:
            await self._collection_manager.stop()

    async def _handle_metrics(self, metrics: list[DomainMetric]) -> None:
        """Handle collected metrics."""
        assert self._buffer is not None
        assert self._batch_builder is not None

        # Add to buffer
        await self._buffer.add_many(metrics)

        # Try to batch
        for metric in metrics:
            batch = await self._batch_builder.add(metric)
            if batch:
                await self._send_batch(batch)

    async def _flush_loop(self) -> None:
        """Periodic buffer flush loop."""
        assert self._shutdown_manager is not None
        assert self._batch_builder is not None

        while not self._shutdown_manager.is_shutting_down:
            try:
                await asyncio.sleep(self.config.buffer.flush_interval)

                # Check for aged batches
                batch = await self._batch_builder.check_age()
                if batch:
                    await self._send_batch(batch)

            except asyncio.CancelledError:
                break
            except Exception as e:
                log.error("flush_loop_error", error=str(e))

    async def _send_batch(self, batch: Any) -> None:
        """Send a batch of metrics to the server."""
        assert self._transport is not None

        try:
            await self._transport.send_batch(batch)
        except Exception as e:
            log.error("batch_send_failed", batch_id=batch.id, error=str(e))
            # Batch stays in buffer for retry

    async def _command_poll_loop(self) -> None:
        """Poll for commands periodically."""
        assert self._shutdown_manager is not None

        # Poll every 60 seconds as a backup to heartbeat-triggered fetches
        poll_interval = 60.0

        while not self._shutdown_manager.is_shutting_down:
            try:
                await asyncio.sleep(poll_interval)
                await self._fetch_and_execute_commands(0)  # 0 = check regardless
            except asyncio.CancelledError:
                break
            except Exception as e:
                log.error("command_poll_error", error=str(e))

    async def _fetch_and_execute_commands(self, pending_count: int) -> None:
        """Fetch and execute pending commands."""
        assert self._transport is not None
        assert self._command_runner is not None

        try:
            commands = await self._transport.fetch_commands()

            for cmd_data in commands:
                command = CommandRequest.from_dict(cmd_data)
                response = await self._command_runner.execute(command)

                # Report result
                await self._transport.report_command_result(command.id, response.to_dict())

                # Record in self-metrics
                if self._self_collector:
                    self._self_collector.record_command_execution(response.success)

        except Exception as e:
            log.error("command_execution_error", error=str(e))

    async def _handle_config_change(
        self,
        old_config: AgentConfig,
        new_config: AgentConfig,
    ) -> None:
        """Handle configuration changes."""
        can_reload, reasons = can_hot_reload(old_config, new_config)

        if not can_reload:
            log.warning(
                "config_change_requires_restart",
                reasons=reasons,
            )
            return

        # Apply hot-reloadable changes
        self.config = new_config

        # Update heartbeat interval
        if self._heartbeat_service:
            self._heartbeat_service.config = new_config.heartbeat

        # Update collector configs
        # (More sophisticated hot-reload could be implemented here)

        log.info("config_hot_reloaded")

    async def _flush_buffer(self) -> None:
        """Flush remaining buffer contents on shutdown."""
        assert self._buffer is not None
        assert self._batch_builder is not None

        # Flush any pending batch
        batch = await self._batch_builder.flush()
        if batch:
            try:
                await self._send_batch(batch)
            except Exception as e:
                log.error("final_batch_send_failed", error=str(e))

        # Persist buffer
        await self._buffer.persist()

    async def _stop_transport(self) -> None:
        """Stop transport on shutdown."""
        if self._transport:
            await self._transport.stop()

    async def _save_state(self) -> None:
        """Save state on shutdown."""
        if self._state_manager:
            await self._state_manager.save()


async def run_agent(config: AgentConfig, config_path: Path | None = None) -> None:
    """
    Run the Curadise agent.

    Args:
        config: Agent configuration
        config_path: Optional path to config file for hot-reload
    """
    agent = Agent(config=config, config_path=config_path)
    await agent.initialize()
    await agent.run()
