import asyncio
import concurrent.futures
import contextlib
import json
import logging
import os
import time
import uuid
from typing import Generator

from defence360agent.api.server import (
    APIError,
    APIErrorTooManyRequests,
    APITokenError,
    send_message,
)
from defence360agent.contracts import license
from defence360agent.contracts.config import Core
from defence360agent.contracts.messages import (
    GeneralMetrics,
    Message,
    MessageList,
    MessageType,
)
from defence360agent.contracts.plugins import MessageSink, expect
from defence360agent.internals import delivery_ack, feature_flags
from defence360agent.internals.feature_flags import (
    MESSAGE_LOSS_OBSERVABILITY_FLAG,
    is_enabled,
)
from defence360agent.internals.message_status_publisher import Gen, publisher
from defence360agent.internals.persistent_message import (
    PersistentMessagesQueue,
)
from defence360agent.utils import (
    log_future_errors,
    recurring_check,
    safe_cancel_task,
    Scope,
)
from defence360agent.utils.json import ServerJSONEncoder

logger = logging.getLogger(__name__)

_reporter_gen_queued = Gen()
_reporter_gen_sending = Gen()
_reporter_gen_sent = Gen()


class SendToServerClient:
    """Send messages to server.

    * process Reportable messages;
    * add them to a pending messages list;
    * send all pending messages to server when list is full (contains
      _PENDING_MESSAGES_LIMIT items or more) or when the oldest pending
      message has waited the max send delay (0 unless batching is
      enabled via the feature flag);
    * send all pending messages on plugin shutdown."""

    _PENDING_MESSAGES_LIMIT = int(
        os.environ.get("IMUNIFYAV_MESSAGES_COUNT_TO_SEND", 20)
    )
    _MAX_SEND_DELAY = 0.0
    _BATCHING_FLAG = "message_send_batching"
    # paces retries of messages re-queued after failed sends
    _SEND_MESSAGE_RECURRING_TIME = 60
    _METRICS_REPORT_INTERVAL = 60 * 5
    # 50 second because it should be less than DefaultTimeoutStopSec
    _SHUTDOWN_SEND_TIMEOUT = 50
    _METRICS_FLUSH_TIMEOUT = 5

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._unsent_metrics = {}

    async def create_sink(self, loop: asyncio.AbstractEventLoop):
        self._loop = loop
        self._pending = PersistentMessagesQueue()
        self._try_send = asyncio.Event()
        self._lock = asyncio.Lock()
        self._shutting_down = asyncio.Event()
        self._flush_deadline = None
        self._metrics_task = loop.create_task(self._report_metrics())
        self._sender_task = loop.create_task(self._send())
        self._invoke_send_message_task = loop.create_task(
            self._invoke_send_message()
        )

    @recurring_check(_METRICS_REPORT_INTERVAL)
    async def _report_metrics(self):
        await self._emit_metrics()

    def _collect_metrics(self) -> dict:
        return {
            "agent.persistent_queue.evicted": self._pending.pop_evicted(),
            "agent.msg_status.dropped": publisher.pop_dropped(),
            "agent.send.method_missing_dropped": (
                send_message.pop_method_missing_dropped()
            ),
        }

    def _collect_gauges(self) -> dict:
        return {
            "agent.persistent_queue.size": self._pending.qsize(),
            "agent.persistent_queue.storage_size": self._pending.storage_size,
            "agent.msg_status.queue_size": publisher.queue_depth(),
        }

    async def _emit_metrics(self):
        # Deliver outside the persistent send-queue: a loss report routed
        # through it could be evicted by the very loss it reports. On failed
        # delivery the deltas roll into the next interval's report instead.
        metrics = self._unsent_metrics
        self._unsent_metrics = {}
        collected = self._collect_metrics()
        # Flag off: drain and discard so the first report after enabling
        # reflects only post-enable activity, not a backlog.
        if not is_enabled(MESSAGE_LOSS_OBSERVABILITY_FLAG):
            return
        for name, value in collected.items():
            if value:
                metrics[name] = metrics.get(name, 0) + value
        sent = False
        try:
            # Gauge sampling can hit the DB, so it stays inside the guarded
            # region: any failure past this point must restore the popped
            # deltas rather than count toward recurring_check's error limit.
            # Gauges are point-in-time samples: taken fresh each interval and
            # never carried over — a stale depth is worse than a missing one.
            payload = {**metrics, **self._collect_gauges()}
            message = GeneralMetrics(
                [
                    {"name": name, "value": value}
                    for name, value in payload.items()
                ]
            )
            message["timestamp"] = time.time()
            message["message_id"] = uuid.uuid4().hex
            # Serialize with _send_pending_messages: the NATS sink shares one
            # gateway connection between both paths, and reconnecting closes it
            # and consumes the reconnect slot, so an unlocked metrics report
            # can break a batch that is in flight.
            async with self._lock:
                sent = await self._send_metrics_direct(message)
        except asyncio.CancelledError:
            # CancelledError is not an Exception: without this branch a
            # cancellation landing mid-send would eat the popped deltas.
            self._unsent_metrics = metrics
            raise
        except Exception as exc:
            logger.warning("Failed to deliver loss metrics: %r", exc)
        if not sent:
            self._unsent_metrics = metrics

    async def _send_metrics_direct(self, message: Message) -> bool:
        with self._get_api() as api:
            await api.send_messages(
                [(time.time(), self._encode_data_to_put_in_queue(message))]
            )
        return True

    async def shutdown(self) -> None:
        """
        When shutdown begins it signals any in-flight HTTP sends to
        abort immediately (via _shutting_down event), then gives 50
        seconds to finish the stop() sequence.  If stop() isn't done
        in 50 seconds it force-cancels the sender task.
        Finally, any messages still in the buffer are flushed to
        persistent storage so nothing is lost.
        """
        # Signal shutdown — aborts in-flight HTTP requests from the
        # _send task via the asyncio.wait race in _send_pending_messages.
        # This lets stop() acquire the lock quickly instead of waiting
        # for a slow HTTP response.  The event is cleared in stop()
        # before the final _send_pending_messages() flush so that
        # remaining messages are actually delivered during shutdown.
        self._shutting_down.set()

        try:
            await asyncio.wait_for(self.stop(), self._SHUTDOWN_SEND_TIMEOUT)
        except asyncio.TimeoutError:
            # Used logger.error to notify sentry
            logger.error(
                "Timeout (%ds) sending messages to server on shutdown.",
                self._SHUTDOWN_SEND_TIMEOUT,
            )
            if not self._sender_task.cancelled():
                await safe_cancel_task(self._sender_task)
        if self._pending.buffer_size > 0:
            logger.warning(
                "Save %s messages to persistent storage",
                self._pending.buffer_size,
            )
            self._pending.push_buffer_to_storage()
            logger.warning("Stored queue %r", self._pending.qsize())

    async def stop(self):
        """
        Stop sending.
        1. wait for the lock being available
            i.e., while _sender_task finishes the current round
            of sending message (if it takes too long, then
            the timeout in shutdown() is triggered
        2. once the sending round complete (we got the lock),
            cancel the next iteration of the _sender_task (it exits)
        3. send _pending messages (again, if it takes too long,
            the timeout in shutdown() is triggered
            and the coroutine is cancelled

        That method makes sure that the coroutine
        that was started in it has ended.

        It excludes a situation when:
            -> The result of a coroutine that started
                BEFORE shutdown() is started.
            -> And the process of sending messages
                from _pending is interrupted because of it
        """
        # The _lock allows you to be sure that the _send_pending_messages
        # coroutine is not running and _pending is not being used
        logger.info("SendToServer.stop cancel _invoke_send_message_task")
        await safe_cancel_task(self._invoke_send_message_task)
        if self._metrics_task is not None:
            await safe_cancel_task(self._metrics_task)
        logger.info("SendToServer.stop wait lock")
        async with self._lock:
            # Cancel _sender_task. The lock ensures that the coroutine
            # is not in its critical part
            logger.info("SendToServer.stop lock acquired, cancel _sender_task")
            await safe_cancel_task(self._sender_task)
            # Clear the shutdown signal so the final flush actually
            # delivers messages instead of re-queuing them.
            self._shutting_down.clear()
            # send messages that are in _pending at the time of agent shutdown
            await self._send_pending_messages()
        if self._metrics_task is not None:
            # Final metrics flush: after the real messages so it cannot eat
            # their shutdown budget, time-bounded for the same reason, and
            # only once the task is cancelled so it cannot race this emit.
            with contextlib.suppress(asyncio.TimeoutError):
                await asyncio.wait_for(
                    self._emit_metrics(), self._METRICS_FLUSH_TIMEOUT
                )

    @staticmethod
    def _set_api_attrs(api):
        api.set_product_name(license.LicenseCLN.get_product_name())
        api.set_server_id(license.LicenseCLN.get_server_id())
        api.set_license(license.LicenseCLN.get_token())
        return api

    @contextlib.contextmanager
    def _get_api(self) -> Generator[send_message.SendMessageAPI, None, None]:
        base_url = os.environ.get("IMUNIFYAV_API_BASE")
        # we send messages sequentially, so max_workers=1
        with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
            api = send_message.SendMessageAPI(
                Core.VERSION, base_url, executor=executor
            )

            yield self._set_api_attrs(api)

    @expect(MessageType.Reportable)
    async def send_to_server(self, message: Message) -> None:
        # add message handling time if it does not exist, so that
        # the server does not depend on the time it was received
        if "timestamp" not in message:
            message["timestamp"] = time.time()
        if "message_id" not in message:
            message["message_id"] = uuid.uuid4().hex
        self._pending.put(self._encode_data_to_put_in_queue(message))
        self._try_send.set()
        publisher.report(message, _reporter_gen_queued, stage="agent-queued")

    @recurring_check(_SEND_MESSAGE_RECURRING_TIME)
    async def _invoke_send_message(self):
        self._try_send.set()

    def _max_send_delay(self) -> float:
        if feature_flags.is_enabled(self._BATCHING_FLAG):
            for value in feature_flags.get_params(self._BATCHING_FLAG):
                try:
                    return float(value)
                except ValueError:
                    pass
        return self._MAX_SEND_DELAY

    @recurring_check(0)
    async def _send(self):
        if self._flush_deadline is None:
            await self._try_send.wait()
        else:
            timeout = max(0, self._flush_deadline - self._loop.time())
            with contextlib.suppress(asyncio.TimeoutError):
                await asyncio.wait_for(self._try_send.wait(), timeout)
        self._try_send.clear()
        qsize = self._pending.qsize()
        if qsize == 0:
            self._flush_deadline = None
            return
        if self._flush_deadline is None:
            self._flush_deadline = self._loop.time() + self._max_send_delay()
        if (
            qsize < self._PENDING_MESSAGES_LIMIT
            and self._loop.time() < self._flush_deadline
        ):
            return
        self._flush_deadline = None
        # The _lock protects critical part of _send method
        logger.info("SendToServer._send wait lock")
        need_to_cancel = None
        async with self._lock:
            logger.info("SendToServer._send lock acquired")
            try:
                await self._send_pending_messages()
            except asyncio.CancelledError as e:
                logger.info("SendToServer._send cancelled unlocking")
                need_to_cancel = e
        logger.info("SendToServer._send lock released")
        if need_to_cancel:
            raise need_to_cancel

    def _encode_data_to_put_in_queue(self, data: Message) -> bytes:
        msg = json.dumps(data, cls=ServerJSONEncoder) + "\n"
        return msg.encode()

    def _decode_message(self, message: bytes) -> Message:
        data = json.loads(message)
        if data.get("list"):
            msg = MessageList(data["list"])
            msg.update({k: v for k, v in data.items() if k != "list"})
            return msg
        return Message(data)

    def _persist_failed(self, message_id, timestamp, message):
        message["api_retries_count"] = message.get("api_retries_count", 0) + 1
        encoded = self._encode_data_to_put_in_queue(message)
        if message_id is None:
            # never stored yet: make it durable now, not on the next flush
            self._pending.put(encoded, timestamp=timestamp)
            self._pending.push_buffer_to_storage()
        else:
            self._pending.update_message(message_id, encoded)

    async def _send_one_message(self, api, message):
        """Race the HTTP send against the shutdown signal.

        Returns True on success, raises on API error,
        or returns False if shutdown interrupted the send.
        """
        send_task = asyncio.ensure_future(api.send_message(message))
        # Consume errors of an abandoned send; the handled path warns.
        send_task.add_done_callback(
            lambda task: log_future_errors(task, logger.debug)
        )
        shutdown_task = asyncio.ensure_future(self._shutting_down.wait())
        try:
            done, pending_tasks = await asyncio.wait(
                {send_task, shutdown_task},
                return_when=asyncio.FIRST_COMPLETED,
            )
        except asyncio.CancelledError:
            send_task.cancel()
            shutdown_task.cancel()
            raise
        for task in pending_tasks:
            await safe_cancel_task(task)

        if send_task in done:
            # Prefer send completion when both tasks finish in one loop turn.
            send_task.result()
            return True

        # Shutdown won the race
        return False

    async def _try_send_one(self, api, message_id, timestamp, message_bytes):
        """Deliver one message and return (stop, failed); message_id is None
        for a fresh memory-only message, set for a stored row."""
        if self._shutting_down.is_set():
            logger.warning(
                "Shutdown signal received, keeping remaining messages"
            )
            return True, False

        message = self._decode_message(message_bytes)
        msg_info = {
            "method": message.get("method"),
            "message_id": message.get("message_id"),
        }
        try:
            publisher.report(
                message, _reporter_gen_sending, stage="agent-sending"
            )
            sent = await self._send_one_message(api, message)
            if not sent:
                logger.warning(
                    "Shutdown signal received during send,"
                    " keeping remaining messages"
                )
                return True, False
            # Dropped, not delivered; the agent-sending report above stays,
            # so the loss surfaces as a stage gap rather than a delivery.
            if msg_info["method"]:
                publisher.report(
                    message, _reporter_gen_sent, stage="agent-sent"
                )
                logger.info("message sent %s", msg_info)
                delivery_ack.registry.confirm(message.get("message_id"))
            if message_id is not None:
                self._pending.delete([message_id])
            return False, False
        except (APIErrorTooManyRequests, APITokenError) as exc:
            logger.warning(
                "Failed to send message %s to server: %s", msg_info, exc
            )
            self._persist_failed(message_id, timestamp, message)
            return True, True
        except APIError as exc:
            logger.warning(
                "Failed to send message %s to server: %s", msg_info, exc
            )
            self._persist_failed(message_id, timestamp, message)
            return False, True

    async def _send_pending_messages(self) -> None:
        with self._get_api() as api:
            if api.server_id is None:
                return
            # stored backlog (older, not deleted) first, then fresh buffer
            batch = list(self._pending.peek_stored()) + [
                (None, timestamp, message_bytes)
                for timestamp, message_bytes in self._pending.drain_buffer()
            ]
            logger.info("Sending %s messages", len(batch))
            failure_count = 0
            processed = 0
            try:
                for message_id, timestamp, message_bytes in batch:
                    stop, failed = await self._try_send_one(
                        api, message_id, timestamp, message_bytes
                    )
                    if stop and not failed:
                        # shutdown aborted this message before any attempt;
                        # leave it in the un-attempted tail so it is persisted
                        break
                    processed += 1
                    if failed:
                        failure_count += 1
                    if stop:
                        # server-level stop: the rest won't send either
                        failure_count += len(batch) - processed
                        break
            finally:
                # un-attempted fresh messages are memory-only; stored are not
                unattempted_fresh = [
                    (ts, mb)
                    for mid, ts, mb in batch[processed:]
                    if mid is None
                ]
                if unattempted_fresh:
                    self._pending.put_many(unattempted_fresh)
                    self._pending.push_buffer_to_storage()
            logger.info("Unsuccessful to send %s messages", failure_count)


class SendToServer(SendToServerClient, MessageSink):
    SCOPE = Scope.AV
    SHUTDOWN_PRIORITY = 900  # Shutdown late, after Accumulate has flushed

    async def _send_metrics_direct(self, message: Message) -> bool:
        with self._get_api() as api:
            if api.server_id is None:
                return False
            await api.send_message(message)
        return True