"""Ingestion worker: consumes from the Kafka ``data-ingestion`` topic and
writes run data into FoundationDB via the storage layer.
"""

from __future__ import annotations

import asyncio
import time
from collections.abc import Callable
from datetime import datetime
from typing import TYPE_CHECKING

from aiokafka import AIOKafkaConsumer
from loguru import logger
from matyan_api_models.kafka import IngestionMessage
from pydantic import ValidationError

from matyan_backend.config import SETTINGS
from matyan_backend.fdb_types import FDBError
from matyan_backend.storage import entities, runs, sequences
from matyan_backend.storage.context_utils import context_to_id
from matyan_backend.storage.fdb_client import ensure_directories, init_fdb

if TYPE_CHECKING:
    from matyan_backend.fdb_types import Database


class IngestionWorker:
    def __init__(self) -> None:
        self._consumer: AIOKafkaConsumer | None = None
        self._db: Database | None = None
        self._running = False

    async def start(self) -> None:
        self._db = init_fdb()
        ensure_directories(self._db)
        logger.info("FDB initialized")

        self._consumer = AIOKafkaConsumer(
            SETTINGS.kafka_data_ingestion_topic,
            bootstrap_servers=SETTINGS.kafka_bootstrap_servers,
            group_id="ingestion-workers",
            auto_offset_reset="earliest",
            enable_auto_commit=False,
            value_deserializer=lambda v: v.decode(),
        )
        await self._consumer.start()
        logger.info(
            "Ingestion worker started (topic={}, group=ingestion-workers)",
            SETTINGS.kafka_data_ingestion_topic,
        )

        self._running = True
        try:
            await self._consume_loop()
        finally:
            await self.stop()

    async def stop(self) -> None:
        self._running = False
        if self._consumer is not None:
            await self._consumer.stop()
            logger.info("Ingestion worker stopped")

    async def _consume_loop(self) -> None:
        assert self._consumer is not None
        async for record in self._consumer:
            try:
                msg = IngestionMessage.model_validate_json(record.value)
                await asyncio.to_thread(self._handle_message, msg)
                await self._consumer.commit()
            except ValidationError:
                logger.warning("Invalid ingestion message, skipping", offset=record.offset, partition=record.partition)
            except FDBError:
                logger.exception("FDB error processing message", offset=record.offset, partition=record.partition)
            except Exception:  # noqa: BLE001
                logger.exception(
                    "Unexpected error processing message",
                    offset=record.offset,
                    partition=record.partition,
                )

    def _handle_message(self, msg: IngestionMessage) -> None:
        db = self._db
        assert db is not None
        handler = _HANDLERS.get(msg.type)
        if handler is None:
            logger.warning("Unknown message type: {}", msg.type)
            return
        handler(db, msg)

    # ------------------------------------------------------------------
    # Per-type handlers (all run synchronously inside asyncio.to_thread)
    # ------------------------------------------------------------------


def _parse_client_ts(payload: dict) -> float | None:
    raw = payload.get("client_datetime")
    if not raw:
        return None
    try:
        return datetime.fromisoformat(raw).timestamp()
    except (ValueError, TypeError):
        return None


def _handle_create_run(db: Database, msg: IngestionMessage) -> None:
    force_resume = msg.payload.get("force_resume", False)
    client_ts = _parse_client_ts(msg.payload)
    if force_resume:
        existing = runs.resume_run(db, msg.run_id)
        if existing is not None:
            if client_ts is not None:
                runs.update_run_meta(db, msg.run_id, client_start_ts=client_ts)
            logger.debug("Resumed run {}", msg.run_id)
            return
    runs.create_run(db, msg.run_id)
    if client_ts is not None:
        runs.update_run_meta(db, msg.run_id, client_start_ts=client_ts)
    logger.debug("Created run {}", msg.run_id)


def _handle_log_metric(db: Database, msg: IngestionMessage) -> None:
    p = msg.payload
    ctx = p.get("context") or {}
    ctx_id = context_to_id(ctx)
    name: str = p["name"]
    value = p["value"]
    step: int | None = p.get("step")
    epoch: int | None = p.get("epoch")
    dtype: str = p.get("dtype") or "float"

    if step is None:
        step = _next_step(db, msg.run_id, ctx_id, name)

    runs.set_context(db, msg.run_id, ctx_id, ctx)

    sequences.write_sequence_step(
        db,
        msg.run_id,
        ctx_id,
        name,
        step,
        value,
        epoch=epoch,
        timestamp=time.time(),
    )

    runs.set_trace_info(
        db,
        msg.run_id,
        ctx_id,
        name,
        dtype=dtype,
        last=value,
        last_step=step,
    )

    logger.debug("Logged metric {}.{} step={} for run {}", name, ctx_id, step, msg.run_id)


def _handle_log_hparams(db: Database, msg: IngestionMessage) -> None:
    value = msg.payload.get("value", {})
    if not value or not isinstance(value, dict):
        return
    path: list[str] = []
    current = value
    while isinstance(current, dict) and len(current) == 1:
        key = next(iter(current))
        path.append(key)
        current = current[key]
    runs.set_run_attrs(db, msg.run_id, tuple(path), current)
    logger.debug("Set attrs at {} for run {}", path, msg.run_id)


def _handle_finish_run(db: Database, msg: IngestionMessage) -> None:
    client_finish = _parse_client_ts(msg.payload)
    meta = runs.get_run_meta(db, msg.run_id)
    client_start = meta.get("client_start_ts")

    duration: float | None = None
    if client_finish is not None and client_start is not None:
        duration = max(0.0, client_finish - client_start)

    created_at = meta.get("created_at", 0.0)
    finalized_at = (created_at + duration) if duration is not None else time.time()

    runs.update_run_meta(
        db,
        msg.run_id,
        active=False,
        finalized_at=finalized_at,
        duration=duration,
    )
    logger.debug("Finished run {} (duration={}s)", msg.run_id, duration)


def _handle_log_custom_object(db: Database, msg: IngestionMessage) -> None:
    p = msg.payload
    ctx = p.get("context") or {}
    ctx_id = context_to_id(ctx)
    name: str = p["name"]
    value: dict = p["value"]
    step: int | None = p.get("step")
    epoch: int | None = p.get("epoch")
    dtype: str = p.get("dtype") or "custom"

    if step is None:
        step = _next_step(db, msg.run_id, ctx_id, name)

    runs.set_context(db, msg.run_id, ctx_id, ctx)

    sequences.write_sequence_step(
        db,
        msg.run_id,
        ctx_id,
        name,
        step,
        value,
        epoch=epoch,
        timestamp=time.time(),
    )

    runs.set_trace_info(
        db,
        msg.run_id,
        ctx_id,
        name,
        dtype=dtype,
        last=0.0,
        last_step=step,
    )

    logger.debug("Logged custom object {}.{} step={} for run {}", name, ctx_id, step, msg.run_id)


def _handle_blob_ref(db: Database, msg: IngestionMessage) -> None:
    p = msg.payload
    artifact_path = p.get("artifact_path", "unknown")
    runs.set_run_attrs(
        db,
        msg.run_id,
        ("__blobs__", artifact_path),
        {
            "s3_key": p.get("s3_key", ""),
            "content_type": p.get("content_type", "application/octet-stream"),
        },
    )
    logger.debug("Stored blob ref {} for run {}", artifact_path, msg.run_id)


def _handle_set_run_property(db: Database, msg: IngestionMessage) -> None:
    p = msg.payload
    run = runs.get_run(db, msg.run_id)
    if run is None:
        runs.create_run(db, msg.run_id)

    meta_updates: dict[str, object] = {}
    if "name" in p:
        meta_updates["name"] = p["name"]
    if "description" in p:
        meta_updates["description"] = p["description"]
    if "archived" in p:
        meta_updates["is_archived"] = p["archived"]
    if meta_updates:
        runs.update_run_meta(db, msg.run_id, **meta_updates)

    if "experiment" in p:
        exp_name = p["experiment"]
        exp = entities.get_experiment_by_name(db, exp_name)
        if exp is None:
            exp = entities.create_experiment(db, exp_name)
        entities.set_run_experiment(db, msg.run_id, exp["id"])
        runs.set_run_experiment(db, msg.run_id, exp["id"])

    logger.debug("Set properties {} for run {}", list(p.keys()), msg.run_id)


def _handle_add_tag(db: Database, msg: IngestionMessage) -> None:
    tag_name = msg.payload["tag_name"]
    run = runs.get_run(db, msg.run_id)
    if run is None:
        runs.create_run(db, msg.run_id)

    tag = entities.get_tag_by_name(db, tag_name)
    if tag is None:
        tag = entities.create_tag(db, tag_name)
    entities.add_tag_to_run(db, msg.run_id, tag["id"])
    runs.add_tag_to_run(db, msg.run_id, tag["id"])
    logger.debug("Added tag {!r} to run {}", tag_name, msg.run_id)


def _handle_remove_tag(db: Database, msg: IngestionMessage) -> None:
    tag_name = msg.payload["tag_name"]
    tag = entities.get_tag_by_name(db, tag_name)
    if tag is None:
        logger.warning("Tag {!r} not found, skipping remove for run {}", tag_name, msg.run_id)
        return
    entities.remove_tag_from_run(db, msg.run_id, tag["id"])
    runs.remove_tag_from_run(db, msg.run_id, tag["id"])
    logger.debug("Removed tag {!r} from run {}", tag_name, msg.run_id)


def _handle_log_terminal_line(db: Database, msg: IngestionMessage) -> None:
    p = msg.payload
    line: str = p["line"]
    step: int = p["step"]

    runs.set_context(db, msg.run_id, 0, {})

    sequences.write_sequence_step(
        db,
        msg.run_id,
        0,
        "logs",
        step,
        line,
        timestamp=time.time(),
    )

    runs.set_trace_info(db, msg.run_id, 0, "logs", dtype="logs", last=0.0, last_step=step)
    logger.debug("Logged terminal line step={} for run {}", step, msg.run_id)


def _handle_log_record(db: Database, msg: IngestionMessage) -> None:
    p = msg.payload
    record_dict = {
        "message": p["message"],
        "log_level": p["level"],
        "timestamp": p["timestamp"],
        "args": p.get("extra_args"),
        "logger_info": p.get("logger_info"),
    }

    step = _next_step(db, msg.run_id, 0, "__log_records")

    runs.set_context(db, msg.run_id, 0, {})

    sequences.write_sequence_step(
        db,
        msg.run_id,
        0,
        "__log_records",
        step,
        record_dict,
        timestamp=time.time(),
    )

    runs.set_trace_info(db, msg.run_id, 0, "__log_records", dtype="log_records", last=0.0, last_step=step)
    logger.debug("Logged record step={} for run {}", step, msg.run_id)


def _next_step(db: Database, run_hash: str, ctx_id: int, name: str) -> int:
    """Auto-assign the next step if the client didn't provide one."""
    return sequences.get_sequence_length(db, run_hash, ctx_id, name)


_HandlerFn = Callable[["Database", IngestionMessage], None]

_HANDLERS: dict[str, _HandlerFn] = {
    "create_run": _handle_create_run,
    "log_metric": _handle_log_metric,
    "log_custom_object": _handle_log_custom_object,
    "log_hparams": _handle_log_hparams,
    "finish_run": _handle_finish_run,
    "blob_ref": _handle_blob_ref,
    "set_run_property": _handle_set_run_property,
    "add_tag": _handle_add_tag,
    "remove_tag": _handle_remove_tag,
    "log_terminal_line": _handle_log_terminal_line,
    "log_record": _handle_log_record,
}
