import json
import logging
from typing import Optional

from aioch import Client

from maps_adv.statistics.beekeeper.lib.steps import sqls
from maps_adv.statistics.beekeeper.lib.steps.base import BaseStep
from maps_adv.warden.client.lib import TaskContext

from . import schemas


class MarkProcessedPacket(BaseStep):

    input_schema = schemas.FreeEventsPorterState
    output_schema = schemas.MarkProcessedPacketState

    def __init__(
        self,
        ch_client_params: dict,
        build_revision: Optional[int] = None,
        ch_query_id: str = "",
    ):
        self._ch_client_params = ch_client_params
        self._build_revision = build_revision or 0
        self._ch_query_id = ch_query_id

    async def run(
        self, data: Optional[dict] = None, context: Optional[TaskContext] = None
    ):
        packet_start_ts = int(data["packet_start"].timestamp())
        packet_end_ts = int(data["packet_end"].timestamp())

        ch_client = Client(**self._ch_client_params)
        events_exist = await ch_client.execute(
            sqls.check_packet_is_processed.format(packet_end=packet_end_ts),
            query_id=self._ch_query_id,
            settings={"replace_running_query": 0},
        )
        if events_exist:
            logging.getLogger("beekeeper.mark_processed_packet").error(
                "Ignored duplicate mark event saving for packet (%d, %d)",
                packet_start_ts,
                packet_end_ts,
            )
            return data

        processing_metadata = {}
        if context:
            processing_metadata["warden_executor_id"] = context.client.executor_id

        query = sqls.insert_to_processed_from_subq.format(
            processing_metadata=json.dumps(processing_metadata),
            subq=sqls.select_beekeeper_technical_event_processed_to.format(
                packet_end=packet_end_ts,
                build_revision=self._build_revision,
                normalization_metadata="{}",
            ),
        )
        await ch_client.execute(
            query,
            query_id=self._ch_query_id,
            settings={"replace_running_query": 0},
        )

        return data
