"""qRemoteX — main entry point.

Wires together configuration, connection, protocol routing, execution
engine, metrics collection, and the job queue.
"""

from __future__ import annotations

import asyncio
import logging
import multiprocessing as mp
import sys
from datetime import datetime, timezone
from typing import Any

from qremotex.config import ExecutorConfig
from qremotex.connection.auth import AuthManager
from qremotex.connection.key_rotation import KeyRotationHandler
from qremotex.connection.manager import ConnectionManager
from qremotex.execution.job_queue import Job, JobQueue
from qremotex.execution.script_executor import ScriptExecutor
from qremotex.execution.tool_executor import RemoteToolExecutor
from qremotex.executor_group import DrainController
from qremotex.logging.logger import setup_logging
from qremotex.metrics.collector import MetricsCollector
from qremotex.protocol.messages import (
    ExecutionStatus,
    MessageType,
    ack_message,
    error_message,
    log_message,
    result_message,
    status_message,
)
from qremotex.protocol.router import MessageRouter
from qremotex.transfer.downloader import FileDownloader
from qremotex.transfer.uploader import FileUploader

logger = logging.getLogger("qremotex")


class QRemoteXAgent:
    """Top-level agent that orchestrates all components."""

    def __init__(self, config: ExecutorConfig) -> None:
        self._config = config
        self._metrics = MetricsCollector()
        self._script_executor = ScriptExecutor()
        self._tool_executor = RemoteToolExecutor(config)
        self._router = MessageRouter()
        self._conn: ConnectionManager | None = None
        self._job_queue: JobQueue | None = None
        self._downloader = FileDownloader()
        self._uploader: FileUploader | None = None  # needs send_fn, set after conn init
        self._key_rotation: KeyRotationHandler | None = None
        self._drain_controller: DrainController | None = None

    async def run(self) -> None:
        """Start the agent — connects, registers, and processes jobs."""
        # Wire up the job queue
        self._job_queue = JobQueue(
            concurrency=self._config.concurrency,
            worker=self._execute_job,
        )
        self._job_queue.start()

        # Register message handlers
        self._router.register(MessageType.EXECUTE_SCRIPT, self._on_execute_script)
        self._router.register(MessageType.EXECUTE_TOOL, self._on_execute_tool)
        self._router.register(MessageType.CANCEL, self._on_cancel)
        self._router.register(MessageType.CONFIG_UPDATE, self._on_config_update)
        # File transfer handlers (stubs for Phase A)
        self._router.register(MessageType.DOWNLOAD_FILE_START, self._on_download_start)
        self._router.register(MessageType.DOWNLOAD_FILE_CHUNK, self._on_download_chunk)
        self._router.register(MessageType.DOWNLOAD_FILE_COMPLETE, self._on_download_complete)
        self._router.register(MessageType.KEY_ROTATION_NOTICE, self._on_key_rotation_notice)
        self._router.register(MessageType.DRAIN, self._on_drain)

        # Wire up the connection manager
        self._conn = ConnectionManager(
            config=self._config,
            on_message=self._router.dispatch,
            metrics_fn=self._metrics_snapshot,
        )

        # Wire up key rotation handler
        self._key_rotation = KeyRotationHandler(
            config=self._config,
            auth_manager=self._conn._auth,
            on_key_rotated=self._on_key_rotated,
        )
        self._key_rotation.register_signal_handler()

        # Wire up drain controller (HA groups)
        self._drain_controller = DrainController(
            active_jobs_fn=lambda: (
                self._job_queue.running_count + self._job_queue.queued_count
                if self._job_queue else 0
            ),
            on_drained=self._on_fully_drained,
        )

        # Wire up the file uploader (needs send fn from connection manager)
        self._uploader = FileUploader(
            send_fn=self._send,
            executor_id=self._config.executor_id,
        )

        logger.info(
            "qRemoteX agent starting — executor_id=%s name=%s concurrency=%d group=%s",
            self._config.executor_id,
            self._config.executor_name,
            self._config.concurrency,
            self._config.group or "(none)",
        )

        await self._conn.run()

    # ------------------------------------------------------------------ #
    # Metrics
    # ------------------------------------------------------------------ #
    def _metrics_snapshot(self) -> dict[str, Any]:
        if self._job_queue:
            self._metrics.set_queued_count(self._job_queue.queued_count)
        return self._metrics.snapshot()

    # ------------------------------------------------------------------ #
    # Message handlers
    # ------------------------------------------------------------------ #
    async def _on_execute_script(self, msg: dict[str, Any]) -> None:
        """Handle execute_script message — enqueue as a job."""
        request_id = msg.get("request_id", "")
        if not request_id:
            logger.warning("execute_script missing request_id")
            return

        # Reject if draining
        if self._drain_controller and self._drain_controller.is_draining:
            await self._send(error_message(
                request_id=request_id,
                executor_id=self._config.executor_id,
                error="Executor is draining — not accepting new tasks",
            ))
            return

        # ACK immediately
        await self._send(ack_message(request_id, self._config.executor_id))

        # Enqueue
        job = Job(
            request_id=request_id,
            job_type="script",
            payload=msg,
            timeout_sec=int(msg.get("timeout_sec", self._config.exec_timeout_sec)),
        )
        position = await self._job_queue.enqueue(job)  # type: ignore[union-attr]
        await self._send(status_message(
            request_id, self._config.executor_id, ExecutionStatus.QUEUED, position=position,
        ))
        logger.info("Script job queued: %s (position %d)", request_id, position)

    async def _on_execute_tool(self, msg: dict[str, Any]) -> None:
        """Handle execute_tool message — enqueue as a job."""
        request_id = msg.get("request_id", "")
        if not request_id:
            logger.warning("execute_tool missing request_id")
            return

        # Reject if draining
        if self._drain_controller and self._drain_controller.is_draining:
            await self._send(error_message(
                request_id=request_id,
                executor_id=self._config.executor_id,
                error="Executor is draining — not accepting new tasks",
            ))
            return

        await self._send(ack_message(request_id, self._config.executor_id))

        job = Job(
            request_id=request_id,
            job_type="tool",
            payload=msg,
            timeout_sec=int(msg.get("timeout_sec", self._config.exec_timeout_sec)),
        )
        position = await self._job_queue.enqueue(job)  # type: ignore[union-attr]
        await self._send(status_message(
            request_id, self._config.executor_id, ExecutionStatus.QUEUED, position=position,
        ))
        logger.info("Tool job queued: %s (position %d)", request_id, position)

    async def _on_cancel(self, msg: dict[str, Any]) -> None:
        """Handle cancel message."""
        request_id = msg.get("request_id", "")
        if not request_id:
            return
        cancelled = await self._job_queue.cancel(request_id)  # type: ignore[union-attr]
        if cancelled:
            await self._send(status_message(
                request_id, self._config.executor_id, ExecutionStatus.CANCELLED,
            ))
            logger.info("Cancelled job: %s", request_id)

    async def _on_config_update(self, msg: dict[str, Any]) -> None:
        """Handle config_update — placeholder for future dynamic config."""
        logger.info("Received config_update (not yet implemented)")

    async def _on_key_rotation_notice(self, msg: dict[str, Any]) -> None:
        """Handle KEY_ROTATION_NOTICE from gateway."""
        if self._key_rotation:
            result = self._key_rotation.handle_rotation_notice(msg)
            if result.get("action") == "auto_reloaded":
                logger.info("Key auto-reloaded from key file after rotation notice")
            else:
                logger.warning(
                    "Key rotation notice received — manual update required. "
                    "Set QRAPTOR_API_KEY_FILE and send SIGHUP, or restart with new key."
                )

    def _on_key_rotated(self, new_auth: AuthManager) -> None:
        """Callback when key rotation completes — update connection auth."""
        if self._conn:
            self._conn._auth = new_auth
            logger.info("Connection manager auth updated with rotated key")
    async def _on_drain(self, msg: dict[str, Any]) -> None:
        """Handle DRAIN command from gateway — stop accepting new work."""
        if self._drain_controller:
            if self._conn:
                self._conn.start_drain()
            self._drain_controller.handle_drain_message(msg)

    async def _on_fully_drained(self) -> None:
        """Called when all in-flight jobs complete after drain."""
        logger.info("Executor fully drained — stopping connection")
        if self._conn:
            await self._conn.stop()
    async def _on_download_start(self, msg: dict[str, Any]) -> None:
        """Handle download_file_start — begin receiving a file from DMS."""
        result = self._downloader.handle_start(msg)
        if not result.get("ok"):
            logger.error("Download start failed: %s", result.get("error"))
            # Notify qRunX of failure
            request_id = msg.get("request_id", "")
            if request_id:
                await self._send(error_message(
                    request_id=request_id,
                    executor_id=self._config.executor_id,
                    error=f"File download start failed: {result.get('error')}",
                ))

    async def _on_download_chunk(self, msg: dict[str, Any]) -> None:
        """Handle download_file_chunk — append data to in-progress download."""
        result = self._downloader.handle_chunk(msg)
        if not result.get("ok"):
            logger.error("Download chunk failed: %s", result.get("error"))
            request_id = msg.get("request_id", "")
            if request_id:
                await self._send(error_message(
                    request_id=request_id,
                    executor_id=self._config.executor_id,
                    error=f"File download chunk failed: {result.get('error')}",
                ))

    async def _on_download_complete(self, msg: dict[str, Any]) -> None:
        """Handle download_file_complete — verify checksum and finalize file."""
        result = self._downloader.handle_complete(msg)
        if result.get("ok"):
            logger.info("File downloaded: %s (%d bytes)", result.get("local_path"), result.get("size_bytes", 0))
        else:
            logger.error("Download complete failed: %s", result.get("error"))
            request_id = msg.get("request_id", "")
            if request_id:
                await self._send(error_message(
                    request_id=request_id,
                    executor_id=self._config.executor_id,
                    error=f"File download failed: {result.get('error')}",
                ))

    # ------------------------------------------------------------------ #
    # Job executor (callback for the JobQueue)
    # ------------------------------------------------------------------ #
    async def _execute_job(self, job: Job) -> None:
        """Execute a single job (script or tool)."""
        request_id = job.request_id
        context = job.payload.get("context", {})
        run_id = context.get("run_id", "") or ""
        self._metrics.record_execution_start()

        logger.info("Executing %s job: %s (run_id=%s)", job.job_type, request_id, run_id or "N/A")

        # Send started status
        await self._send(status_message(
            request_id, self._config.executor_id, ExecutionStatus.STARTED,
        ))

        try:
            if job.job_type == "script":
                result = await self._run_script(job, run_id=run_id)
            elif job.job_type == "tool":
                result = await self._run_tool(job)
            else:
                result = {"success": False, "error": f"Unknown job type: {job.job_type}"}

            success = result.get("success", False)
            self._metrics.record_execution_end(success)

            logger.info(
                "Job %s completed: success=%s duration_ms=%s",
                request_id, success, result.get("duration_ms", "?"),
            )

            if success:
                await self._send(result_message(
                    request_id=request_id,
                    executor_id=self._config.executor_id,
                    success=True,
                    result=result.get("result"),
                    duration_ms=result.get("duration_ms", 0),
                    peak_memory_mb=result.get("peak_memory_mb"),
                    run_id=run_id,
                    console_logs=result.get("console_logs"),
                ))
            else:
                await self._send(error_message(
                    request_id=request_id,
                    executor_id=self._config.executor_id,
                    error=result.get("error", "Unknown error"),
                    trace=result.get("trace"),
                    run_id=run_id,
                ))

        except asyncio.CancelledError:
            self._metrics.record_execution_end(False)
            await self._send(status_message(
                request_id, self._config.executor_id, ExecutionStatus.CANCELLED,
            ))
            raise
        except Exception as e:
            self._metrics.record_execution_end(False)
            await self._send(error_message(
                request_id=request_id,
                executor_id=self._config.executor_id,
                error=str(e),
                run_id=run_id,
            ))

    async def _run_script(self, job: Job, *, run_id: str = "") -> dict[str, Any]:
        """Execute a Python script job."""
        payload = job.payload
        script_b64 = payload.get("script_b64", "")
        input_vars = payload.get("input_vars", {})
        config = payload.get("config", {})
        output_variable = config.get("output_variable", "result")

        async def on_log(entry: dict[str, Any]) -> None:
            """Forward log entries from child process to qRunX."""
            await self._send(log_message(
                request_id=job.request_id,
                executor_id=self._config.executor_id,
                level=entry.get("level", "info"),
                stream=entry.get("stream", "stdout"),
                message=entry.get("message", ""),
                timestamp=entry.get("timestamp", datetime.now(timezone.utc).isoformat()),
                run_id=run_id,
            ))

        return await self._script_executor.execute(
            request_id=job.request_id,
            script_b64=script_b64,
            timeout_sec=job.timeout_sec,
            input_vars=input_vars,
            on_log=on_log,
            output_variable=output_variable,
        )

    async def _run_tool(self, job: Job) -> dict[str, Any]:
        """Execute a tool job via RemoteToolExecutor (Phase C).

        The QRP execute_tool message contains:
          - tool_key:            e.g. "system.db_query", "project.crm_api"
          - tool_params:         Fully-resolved input parameters
          - tool_definition:     Metadata incl. execution_target_kind, http_config
          - context:             run_id, node_id, subscription_id, trace_id
        """
        payload = job.payload
        tool_key = payload.get("tool_key", "")
        tool_params = payload.get("tool_params", {})
        tool_definition = payload.get("tool_definition", {})
        context = payload.get("context", {})

        if not tool_key:
            return {"success": False, "error": "Missing tool_key in execute_tool message"}

        return await self._tool_executor.execute(
            tool_key=tool_key,
            tool_params=tool_params,
            tool_definition=tool_definition,
            context=context,
        )

    # ------------------------------------------------------------------ #
    # Send helper
    # ------------------------------------------------------------------ #
    async def _send(self, msg: dict[str, Any]) -> None:
        """Send a message through the connection manager."""
        if self._conn:
            await self._conn.send(msg)


# ──────────────────────────────────────────────── entrypoint ──

def main() -> None:
    """CLI entry point for the qRemoteX agent."""
    mp.set_start_method("spawn", force=True)

    config = ExecutorConfig()
    setup_logging(config.log_level)

    errors = config.validate()
    if errors:
        for err in errors:
            print(f"[ERROR] {err}", file=sys.stderr)
        sys.exit(1)

    logger.info("qRemoteX v%s", config.version)
    logger.info("Gateway: %s", config.gateway_url)
    logger.info("Executor ID: %s", config.executor_id)
    logger.info("Concurrency: %d", config.concurrency)

    agent = QRemoteXAgent(config)

    try:
        asyncio.run(agent.run())
    except KeyboardInterrupt:
        logger.info("Shutting down (KeyboardInterrupt)")
    except Exception:
        logger.exception("Fatal error")
        sys.exit(1)


if __name__ == "__main__":
    main()
