clickstream-ch-kafka-superset-demo/generator/src/clickstream_generator/service.py

"""Основной сервисный цикл генератора."""

import hashlib
import json
import logging
import sys
import time
import uuid
from datetime import datetime, timedelta, timezone

from prometheus_client import start_http_server

from clickstream_generator.config import Config
from clickstream_generator.dictionary import EventDictionary
from clickstream_generator.generation import EventGenerator
from clickstream_generator.kafka_io import (
    BatchRecord,
    KafkaBatchHistory,
    KafkaPublisher,
    KafkaStateManager,
    KafkaStartupHistoryManifest,
    ensure_topics,
)
from clickstream_generator.metrics import (
    METRICS_ERRORS_TOTAL,
    METRICS_LAST_SUCCESS,
    METRICS_TICK_DURATION,
)
from clickstream_generator.runtime import TickStreamGenerator


logger = logging.getLogger("generator")


class _TopicManifestStats:
    """Накопительные счётчики одного Kafka-топика для manifest."""

    def __init__(self):
        self.rows = 0
        self.min_event_timestamp: str | None = None
        self.max_event_timestamp: str | None = None
        self._checksum = hashlib.sha256()

    @property
    def checksum(self) -> str:
        return self._checksum.hexdigest()

    def add(self, event: dict, event_timestamp: str | None = None) -> None:
        self.rows += 1
        self._checksum.update(
            json.dumps(event, sort_keys=True, ensure_ascii=True).encode("utf-8")
        )
        timestamp = event_timestamp if event_timestamp is not None else event.get("event_timestamp")
        self.add_timestamp(timestamp)

    def add_timestamp(self, timestamp: str | None) -> None:
        if timestamp is None:
            return
        if self.min_event_timestamp is None or timestamp < self.min_event_timestamp:
            self.min_event_timestamp = timestamp
        if self.max_event_timestamp is None or timestamp > self.max_event_timestamp:
            self.max_event_timestamp = timestamp

    def to_dict(self) -> dict:
        return {
            "rows": self.rows,
            "min_event_timestamp": self.min_event_timestamp,
            "max_event_timestamp": self.max_event_timestamp,
            "checksum_sha256": self.checksum,
        }


class _ManifestCounters:
    """Счётчики стартовой истории для manifest."""

    def __init__(self):
        self.topic_stats = {
            topic: _TopicManifestStats()
            for topic in ("browser_events", "location_events", "device_events", "geo_events")
        }
        self.click_ids: set[str] = set()
        self.user_ids: set[str] = set()

    def add_batch(self, batch: dict[str, list[dict]]) -> None:
        browser_events = batch.get("browser_events", [])
        event_timestamps = {
            event["event_id"]: event.get("event_timestamp")
            for event in browser_events
            if event.get("event_id") and event.get("event_timestamp")
        }
        click_timestamps: dict[str, list[str]] = {}
        for event in browser_events:
            click_id = event.get("click_id")
            timestamp = event.get("event_timestamp")
            if click_id and timestamp:
                click_timestamps.setdefault(click_id, []).append(timestamp)

        for topic, events in batch.items():
            stats = self.topic_stats[topic]
            for event in events:
                event_timestamp = event.get("event_timestamp")
                if event_timestamp is None and topic == "location_events":
                    event_timestamp = event_timestamps.get(event.get("event_id"))
                if event_timestamp is None and topic in ("device_events", "geo_events"):
                    timestamps = click_timestamps.get(event.get("click_id"))
                    if timestamps:
                        event_timestamp = min(timestamps)
                        stats.add_timestamp(max(timestamps))
                stats.add(event, event_timestamp=event_timestamp)
                click_id = event.get("click_id")
                if topic == "browser_events" and click_id:
                    self.click_ids.add(click_id)
                user_id = event.get("user_domain_id")
                if topic == "device_events" and user_id:
                    self.user_ids.add(user_id)

    def to_manifest_topics(self) -> dict:
        return {
            topic: stats.to_dict()
            for topic, stats in self.topic_stats.items()
        }

    def to_manifest_totals(self) -> dict:
        browser_stats = self.topic_stats["browser_events"]
        return {
            "events": browser_stats.rows,
            "visits": len(self.click_ids),
            "users": len(self.user_ids),
            "min_event_timestamp": browser_stats.min_event_timestamp,
            "max_event_timestamp": browser_stats.max_event_timestamp,
        }


class GeneratorService:
    """Основной сервис генератора."""

    def __init__(self, config: Config):
        self.config = config
        self.dictionary = EventDictionary.load(config.data_dir)
        self.generator = EventGenerator(self.dictionary, config)
        self.stream = TickStreamGenerator(self.generator)
        self.publisher: KafkaPublisher | None = None
        self.history: KafkaBatchHistory | None = None
        self.state_manager: KafkaStateManager | None = None
        self.manifest_manager: KafkaStartupHistoryManifest | None = None
        self._running = False
        self._tick = 0
        self._model_time = config.model_t0

    def start(self):
        """Запускает основной цикл."""
        if not self.config.enabled:
            logger.warning("Generator is disabled (GEN_ENABLED=false)")
            return

        logger.info(f"Starting metrics server on port {self.config.metrics_port}")
        start_http_server(self.config.metrics_port)

        logger.info("Starting generator service...")
        logger.info(
            f"Configuration: tick={self.config.tick_seconds}s, "
            f"lambda_base={self.config.lambda_base_per_min}/min, "
            f"jitter={self.config.jitter_pct}%, "
            f"state_enabled={self.config.state_enabled}, "
            f"state_reset={self.config.state_reset}"
        )

        ensure_topics(self.config.kafka_bootstrap_servers)
        self.publisher = KafkaPublisher(self.config.kafka_bootstrap_servers)
        self.history = KafkaBatchHistory(self.config.kafka_bootstrap_servers)

        if self.config.run_mode == "backfill" and not self.config.state_enabled:
            raise ValueError("GEN_STATE_ENABLED must be true for backfill")

        if self.config.state_enabled:
            self.state_manager = KafkaStateManager(self.config.kafka_bootstrap_servers)

            if self.config.run_mode == "backfill":
                self.manifest_manager = KafkaStartupHistoryManifest(
                    self.config.kafka_bootstrap_servers
                )
                logger.info("Backfill mode starts from a fresh generator state")
                self._run_backfill()
                self.stop()
                return

            if not self.config.state_reset:
                restored_state = self.state_manager.load()
                if restored_state:
                    try:
                        self.manifest_manager = KafkaStartupHistoryManifest(
                            self.config.kafka_bootstrap_servers
                        )
                        manifest = self.manifest_manager.load()
                        if self._is_startup_history_state(restored_state, manifest):
                            model_t_end = self._as_aware_utc(
                                datetime.fromisoformat(manifest["model_t_end"])
                            )
                            self.restore_from_startup_history(
                                restored_state,
                                model_t_end=model_t_end,
                            )
                        elif self._is_startup_history_marker(restored_state):
                            raise ValueError(
                                "startup-history state without matching manifest"
                            )
                        else:
                            self._restore_live_state(
                                restored_state,
                                wall_now_utc=datetime.now(timezone.utc),
                            )
                        logger.info(
                            f"Restored state: continuing from tick {self._tick}, "
                            f"model_time={self._model_time.isoformat()}, "
                            f"last_batch_id={restored_state.last_batch_id}"
                        )
                    except Exception as e:
                        logger.warning(
                            f"State data was invalid, starting fresh: {e}"
                        )
                        self.generator = EventGenerator(self.dictionary, self.config)
                        self.stream = TickStreamGenerator(self.generator)
                        self._tick = 0
            else:
                logger.info("State reset requested, starting fresh")
        else:
            logger.info("State management disabled")

        self._running = True

        try:
            self._main_loop()
        except KeyboardInterrupt:
            logger.info("Received shutdown signal")
        finally:
            self.stop()

    def stop(self):
        """Останавливает сервис."""
        logger.info("Stopping generator service...")
        self._running = False
        if self.publisher:
            self.publisher.close()
        if self.history:
            self.history.close()
        if self.state_manager:
            self.state_manager.close()
        if self.manifest_manager:
            self.manifest_manager.close()

    def restore_from_startup_history(
        self,
        state,
        model_t_end: datetime,
    ) -> None:
        """Восстанавливает слепок стартовой истории ровно от T_end."""
        self._restore_state_snapshot(state, resume_model_at=model_t_end)

    def _restore_live_state(self, state, wall_now_utc: datetime) -> None:
        """Восстанавливает live-state с учётом прошедшего настенного времени."""
        self._validate_live_state_config(state)
        resume_model_at = self._calculate_live_resume_model_at(
            state,
            wall_now_utc=wall_now_utc,
        )
        self._restore_state_snapshot(state, resume_model_at=resume_model_at)

    def _restore_state_snapshot(self, state, resume_model_at: datetime) -> None:
        """Применяет state к генератору и тиковому слою."""
        resume_model_at = self._as_aware_utc(resume_model_at)
        self.generator.rng.setstate(state.rng_state)
        self.stream.restore_state(state, resume_model_at=resume_model_at)
        self._tick = state.tick
        self._model_time = resume_model_at

    def _calculate_live_resume_model_at(
        self,
        state,
        wall_now_utc: datetime,
    ) -> datetime:
        """Считает модельную точку live-восстановления по state v2."""
        wall_now_utc = self._as_aware_utc(wall_now_utc)
        wall_saved_at = self._as_aware_utc(state.wall_timestamp)
        idle_seconds = max(0.0, (wall_now_utc - wall_saved_at).total_seconds())
        return self._as_aware_utc(state.model_timestamp) + timedelta(
            seconds=idle_seconds * state.model_time_speed,
        )

    def _validate_live_state_config(self, state) -> None:
        """Проверяет, что state относится к текущей конфигурации live-запуска."""
        mismatches = []
        if state.gen_seed != self.config.seed:
            mismatches.append("gen_seed")
        if self._as_aware_utc(state.model_t0) != self.config.model_t0:
            mismatches.append("model_t0")
        if state.model_timezone != self.config.model_timezone:
            mismatches.append("model_timezone")
        if abs(state.model_time_speed - self.config.model_time_speed) > 1e-9:
            mismatches.append("model_time_speed")

        if mismatches:
            raise ValueError(
                "state config mismatch: " + ", ".join(mismatches)
            )

    def _is_startup_history_state(self, state, manifest: dict | None) -> bool:
        """Проверяет, что state совпадает со слепком стартовой истории."""
        if not manifest or manifest.get("run_mode") != "backfill":
            return False

        expected_state = manifest.get("state") or {}
        if expected_state.get("last_batch_id") != state.last_batch_id:
            return False

        try:
            model_t_end = self._as_aware_utc(
                datetime.fromisoformat(manifest["model_t_end"])
            )
            model_t0 = self._as_aware_utc(
                datetime.fromisoformat(manifest["model_t0"])
            )
        except (KeyError, TypeError, ValueError):
            return False

        if self._as_aware_utc(state.model_timestamp) != model_t_end:
            return False
        if self._as_aware_utc(state.model_t0) != model_t0:
            return False
        if state.gen_seed != manifest.get("gen_seed"):
            return False
        if state.model_timezone != manifest.get("model_timezone"):
            return False
        if (
            self.config.model_t_end is not None
            and self.config.model_t_end != model_t_end
        ):
            return False
        if model_t0 != self.config.model_t0:
            return False
        if manifest.get("gen_seed") != self.config.seed:
            return False
        if manifest.get("model_timezone") != self.config.model_timezone:
            return False

        settings = manifest.get("generation_settings") or {}
        if abs(state.model_time_speed - settings.get("model_time_speed", -1)) > 1e-9:
            return False
        return self._generation_settings() == settings

    def _is_startup_history_marker(self, state) -> bool:
        """Отличает state стартовой истории от обычного live-state."""
        return str(state.last_batch_id).startswith("startup-history-")

    @staticmethod
    def _as_aware_utc(value: datetime) -> datetime:
        if value.tzinfo is None:
            return value.replace(tzinfo=timezone.utc)
        return value.astimezone(timezone.utc)

    def _save_state(self, batch_id: str) -> None:
        """Сохраняет текущее состояние генератора."""
        if not self.state_manager or not self.config.state_enabled:
            return

        try:
            wall_now = datetime.now(timezone.utc)
            state = self.stream.to_state(
                tick=self._tick,
                rng_state=self.generator.rng.getstate(),
                last_batch_id=batch_id,
                last_timestamp=self._model_time,
                model_timestamp=self._model_time,
                wall_timestamp=wall_now,
                model_time_speed=self.config.model_time_speed,
                model_timezone=self.config.model_timezone,
                model_t0=self.config.model_t0,
                gen_seed=self.config.seed,
            )
            self.state_manager.save(state)
            self.state_manager.flush()
            logger.debug(
                f"Saved state: tick={self._tick}, "
                f"model_time={self._model_time.isoformat()}, batch_id={batch_id}"
            )
        except Exception as e:
            logger.warning(f"Failed to save state: {e}")
            METRICS_ERRORS_TOTAL.labels(topic="state").inc()

    def _run_backfill(self) -> None:
        """Проматывает стартовую историю без сна до GEN_MODEL_T_END."""
        if self.config.model_t_end is None:
            raise ValueError("GEN_MODEL_T_END is required for backfill")
        if not self.publisher:
            raise RuntimeError("publisher is not initialized")
        if not self.history:
            raise RuntimeError("history is not initialized")

        logger.info(
            "Running backfill from %s to %s",
            self.config.model_t0.isoformat(),
            self.config.model_t_end.isoformat(),
        )
        counters = _ManifestCounters()

        while self._model_time < self.config.model_t_end:
            self._tick += 1
            batch_id = f"backfill-{self._tick:08d}"
            model_time = self._model_time
            started_at = datetime.now(timezone.utc)

            events_count = self.generator._calculate_events_count(now=model_time)
            batch = self.stream.generate_tick(
                events_count,
                tick_started_at=model_time,
            )
            total_sent, sent_counts, status = self._publish_batch(batch)
            self._raise_on_backfill_publish_error(batch_id, status, sent_counts)
            counters.add_batch(batch)
            self._write_batch_history(
                batch_id=batch_id,
                started_at=started_at,
                sent_counts=sent_counts,
                total_sent=total_sent,
                status=status,
                error_message=None if status == "success" else "Backfill publish error",
            )
            self._advance_model_time()

        final_batch = self.stream.drain_until(
            self.config.model_t_end,
            include_boundary=False,
        )
        final_sent = 0
        final_status = "success"
        if any(final_batch.values()):
            batch_id = f"backfill-{self._tick + 1:08d}-final"
            started_at = datetime.now(timezone.utc)
            final_sent, sent_counts, final_status = self._publish_batch(final_batch)
            self._raise_on_backfill_publish_error(
                batch_id,
                final_status,
                sent_counts,
            )
            counters.add_batch(final_batch)
            self._write_batch_history(
                batch_id=batch_id,
                started_at=started_at,
                sent_counts=sent_counts,
                total_sent=final_sent,
                status=final_status,
                error_message=None if final_status == "success" else "Backfill publish error",
            )

        if self.publisher:
            self.publisher.flush()

        self._model_time = self.config.model_t_end
        state_batch_id = self._startup_state_batch_id(counters)
        state = self.stream.to_state(
            tick=self._tick,
            rng_state=self.generator.rng.getstate(),
            last_batch_id=state_batch_id,
            last_timestamp=self.config.model_t_end,
            model_timestamp=self.config.model_t_end,
            wall_timestamp=datetime.now(timezone.utc),
            model_time_speed=self.config.model_time_speed,
            model_timezone=self.config.model_timezone,
            model_t0=self.config.model_t0,
            gen_seed=self.config.seed,
        )

        manifest = self._build_startup_history_manifest(
            counters=counters,
            state=state,
        )
        if self.manifest_manager:
            self.manifest_manager.save(manifest)
            self.manifest_manager.flush()

        if self.state_manager and self.config.state_enabled:
            self.state_manager.save(state)
            self.state_manager.flush()

        logger.info(
            "Backfill completed: events=%s, visits=%s, users=%s, final_sent=%s",
            manifest["totals"]["events"],
            manifest["totals"]["visits"],
            manifest["totals"]["users"],
            final_sent,
        )

    def _raise_on_backfill_publish_error(
        self,
        batch_id: str,
        status: str,
        sent_counts: dict[str, dict[str, int]],
    ) -> None:
        """Останавливает backfill до записи state/manifest при ошибке Kafka."""
        if status == "success":
            return

        raise RuntimeError(
            f"Backfill publish failed for batch {batch_id}: "
            f"status={status}, sent_counts={sent_counts}"
        )

    def _publish_batch(
        self,
        batch: dict[str, list[dict]],
    ) -> tuple[int, dict[str, dict[str, int]], str]:
        """Публикует batch и возвращает счётчики отправки."""
        total_sent = 0
        total_errors = 0
        sent_counts = {}

        for topic, events in batch.items():
            if events:
                sent, errors = self.publisher.publish(topic, events)
                sent_counts[topic] = {"sent": sent, "errors": errors}
                total_sent += sent
                total_errors += errors

        if total_errors == 0:
            status = "success"
        elif total_sent > 0:
            status = "partial"
        else:
            status = "error"

        return total_sent, sent_counts, status

    def _write_batch_history(
        self,
        batch_id: str,
        started_at: datetime,
        sent_counts: dict[str, dict[str, int]],
        total_sent: int,
        status: str,
        error_message: str | None,
    ) -> None:
        """Пишет служебную историю batch."""
        if not self.history:
            return
        try:
            record = BatchRecord(
                batch_id=batch_id,
                started_at=started_at,
                finished_at=datetime.now(timezone.utc),
                sent_total=total_sent,
                sent_browser=sent_counts.get("browser_events", {}).get("sent", 0),
                sent_location=sent_counts.get("location_events", {}).get("sent", 0),
                sent_device=sent_counts.get("device_events", {}).get("sent", 0),
                sent_geo=sent_counts.get("geo_events", {}).get("sent", 0),
                status=status,
                error_message=error_message,
            )
            self.history.add(record)
            self.history.flush()
        except Exception as hist_err:
            logger.warning(f"Failed to write batch history: {hist_err}")
            METRICS_ERRORS_TOTAL.labels(topic="history").inc()

    def _startup_state_batch_id(self, counters) -> str:
        digest = counters.topic_stats["browser_events"].checksum[:12]
        return f"startup-history-{digest}"

    def _build_startup_history_manifest(self, counters, state) -> dict:
        return {
            "manifest_version": "1.0",
            "generated_at": datetime.now(timezone.utc).isoformat(),
            "gen_seed": self.config.seed,
            "model_t0": self.config.model_t0.isoformat(),
            "model_t_end": self.config.model_t_end.isoformat(),
            "model_timezone": self.config.model_timezone,
            "run_mode": "backfill",
            "generation_settings": self._generation_settings(),
            "state_version": state.version,
            "state": {
                "topic": KafkaStateManager.STATE_TOPIC,
                "key": KafkaStateManager.STATE_KEY,
                "last_batch_id": state.last_batch_id,
                "model_timestamp": state.model_timestamp.isoformat(),
            },
            "topics": counters.to_manifest_topics(),
            "totals": counters.to_manifest_totals(),
        }

    def _generation_settings(self) -> dict:
        return {
            "tick_seconds": self.config.tick_seconds,
            "lambda_base_per_min": self.config.lambda_base_per_min,
            "jitter_pct": self.config.jitter_pct,
            "min_events_per_tick": self.config.min_events_per_tick,
            "max_events_per_tick": self.config.max_events_per_tick,
            "max_session_events": self.config.max_session_events,
            "max_active_sessions": self.config.max_active_sessions,
            "population_max": self.config.population_max,
            "p_new_user": self.config.p_new_user,
            "min_return_minutes": self.config.min_return_minutes,
            "model_time_speed": self.config.model_time_speed,
        }

    def _main_loop(self):
        """Основной цикл тиков."""
        while self._running:
            self._tick += 1
            tick_start = time.time()
            batch_id = str(uuid.uuid4())[:8]
            model_time = self._model_time

            with METRICS_TICK_DURATION.time():
                logger.info(f"=== Tick {self._tick} (batch_id={batch_id}) ===")

                try:
                    events_count = self.generator._calculate_events_count(
                        now=model_time,
                    )
                    logger.info(f"Generating with event budget ~{events_count}")

                    gen_start = time.time()
                    batch = self.stream.generate_tick(
                        events_count,
                        tick_started_at=model_time,
                    )
                    gen_duration = time.time() - gen_start

                    pub_start = time.time()
                    total_sent = 0
                    total_errors = 0

                    sent_counts = {}
                    for topic, events in batch.items():
                        if events:
                            sent, errors = self.publisher.publish(topic, events)
                            sent_counts[topic] = {"sent": sent, "errors": errors}
                            total_sent += sent
                            total_errors += errors

                    if total_errors == 0:
                        status = "success"
                    elif total_sent > 0:
                        status = "partial"
                    else:
                        status = "error"

                    if status in ("success", "partial"):
                        METRICS_LAST_SUCCESS.set_to_current_time()
                        self._advance_model_time()
                        self._save_state(batch_id)

                    self.publisher.flush()
                    pub_duration = time.time() - pub_start

                    try:
                        batch_record = BatchRecord(
                            batch_id=batch_id,
                            started_at=datetime.fromtimestamp(tick_start, tz=timezone.utc),
                            finished_at=datetime.now(timezone.utc),
                            sent_total=total_sent,
                            sent_browser=sent_counts.get("browser_events", {}).get("sent", 0),
                            sent_location=sent_counts.get("location_events", {}).get("sent", 0),
                            sent_device=sent_counts.get("device_events", {}).get("sent", 0),
                            sent_geo=sent_counts.get("geo_events", {}).get("sent", 0),
                            status=status,
                            error_message=None if status == "success" else f"Errors: {total_errors}",
                        )
                        self.history.add(batch_record)
                        self.history.flush()
                    except Exception as hist_err:
                        logger.warning(f"Failed to write batch history: {hist_err}")
                        METRICS_ERRORS_TOTAL.labels(topic="history").inc()

                    tick_duration = time.time() - tick_start
                    logger.info(
                        f"Batch {batch_id} completed: "
                        f"sent={total_sent}, errors={total_errors}, "
                        f"gen_time={gen_duration:.3f}s, pub_time={pub_duration:.3f}s, "
                        f"total_time={tick_duration:.3f}s"
                    )

                    for topic, counts in sent_counts.items():
                        if counts["sent"] > 0:
                            logger.info(f"  {topic}: {counts['sent']} sent")

                except Exception as e:
                    logger.exception(f"Error in tick {self._tick}: {e}")
                    try:
                        self.history.add(
                            BatchRecord(
                                batch_id=batch_id,
                                started_at=datetime.fromtimestamp(tick_start, tz=timezone.utc),
                                finished_at=datetime.now(timezone.utc),
                                sent_total=0,
                                sent_browser=0,
                                sent_location=0,
                                sent_device=0,
                                sent_geo=0,
                                status="error",
                                error_message=str(e),
                            )
                        )
                        self.history.flush()
                    except Exception as hist_err:
                        logger.warning(f"Failed to write error to history: {hist_err}")

            elapsed = time.time() - tick_start
            sleep_time = max(0, self.config.tick_seconds - elapsed)
            if sleep_time > 0:
                logger.debug(f"Sleeping for {sleep_time:.1f}s until next tick")
                time.sleep(sleep_time)

    def _advance_model_time(self) -> None:
        """Сдвигает модельное время после успешного live-тика."""
        self._model_time = self._model_time + timedelta(
            seconds=self.config.tick_seconds * self.config.model_time_speed,
        )


def main():
    """Точка входа сервиса."""
    try:
        config = Config()
        service = GeneratorService(config)
        service.start()
    except Exception as e:
        logger.exception(f"Fatal error: {e}")
        sys.exit(1)