"""Helper functions for state and bookmark management."""

from __future__ import annotations

import copy
import logging
import typing as t

from singer_sdk.exceptions import InvalidStreamSortException
from singer_sdk.helpers._typing import to_json_compatible
from singer_sdk.singerlib.encoding.simple import StateMessage

if t.TYPE_CHECKING:
    import datetime

    from singer_sdk.helpers import types
    from singer_sdk.singerlib.encoding.base import GenericSingerWriter

    _T = t.TypeVar("_T", datetime.datetime, str, int, float)

PROGRESS_MARKERS = "progress_markers"
PROGRESS_MARKER_NOTE = "Note"
SIGNPOST_MARKER = "replication_key_signpost"
STARTING_MARKER = "starting_replication_value"

logger = logging.getLogger("singer_sdk")


def get_state_if_exists(
    tap_state: types.TapState,
    tap_stream_id: str,
    state_partition_context: types.Context | None = None,
    key: str | None = None,
) -> t.Any | None:  # noqa: ANN401
    """Return the stream or partition state, creating a new one if it does not exist.

    Args:
        tap_state: the existing state dict which contains all streams.
        tap_stream_id: the id of the stream
        state_partition_context: keys which identify the partition context,
            by default None (not partitioned)
        key: name of the key searched for, by default None (return entire state if
            found)

    Returns:
        Returns the state if exists, otherwise None

    Raises:
        ValueError: Raised if state is invalid or cannot be parsed.
    """
    if "bookmarks" not in tap_state or tap_stream_id not in tap_state["bookmarks"]:
        return None

    stream_state = tap_state["bookmarks"][tap_stream_id]
    if not state_partition_context:
        return stream_state.get(key, None) if key else stream_state
    if "partitions" not in stream_state:
        return None  # No partitions defined

    matched_partition = _find_in_partitions_list(
        stream_state["partitions"],
        state_partition_context,
    )
    if matched_partition is None:
        return None  # Partition definition not present
    return matched_partition.get(key, None) if key else matched_partition


def get_state_partitions_list(
    tap_state: types.TapState,
    tap_stream_id: str,
) -> list[dict] | None:
    """Return a list of partitions defined in the state, or None if not defined."""
    return (get_state_if_exists(tap_state, tap_stream_id) or {}).get("partitions", None)


def _find_in_partitions_list(
    partitions: list[dict],
    state_partition_context: types.Context,
) -> dict | None:
    found = [
        partition_state
        for partition_state in partitions
        if partition_state["context"] == state_partition_context
    ]
    if len(found) > 1:
        msg = (
            "State file contains duplicate entries for partition: "
            f"{state_partition_context}.\nMatching state values were: {found!s}"
        )
        raise ValueError(msg)
    return found[0] if found else None


def _create_in_partitions_list(
    partitions: list[dict],
    state_partition_context: types.Context,
) -> dict:
    # Existing partition not found. Creating new state entry in partitions list...
    new_partition_state = {"context": state_partition_context}
    partitions.append(new_partition_state)
    return new_partition_state


def get_writeable_state_dict(
    tap_state: types.TapState,
    tap_stream_id: str,
    state_partition_context: types.Context | None = None,
) -> dict:
    """Return the stream or partition state, creating a new one if it does not exist.

    Args:
        tap_state: the existing state dict which contains all streams.
        tap_stream_id: the id of the stream
        state_partition_context: keys which identify the partition context,
            by default None (not partitioned)

    Returns:
        Returns a writeable dict at the stream or partition level.

    Raises:
        ValueError: Raise an error if duplicate entries are found.
    """
    if tap_state is None:
        msg = "Cannot write state to missing state dictionary."  # type: ignore[unreachable]
        raise ValueError(msg)

    tap_state.setdefault("bookmarks", {})
    tap_state["bookmarks"].setdefault(tap_stream_id, {})

    stream_state = tap_state["bookmarks"][tap_stream_id]
    if not state_partition_context:
        return stream_state

    stream_state.setdefault("partitions", [])
    stream_state_partitions: list[dict] = stream_state["partitions"]
    if found := _find_in_partitions_list(
        stream_state_partitions,
        state_partition_context,
    ):
        return found

    return _create_in_partitions_list(stream_state_partitions, state_partition_context)


def write_stream_state(
    tap_state: types.TapState,
    tap_stream_id: str,
    key: str,
    val: t.Any,  # noqa: ANN401
    *,
    state_partition_context: types.Context | None = None,
) -> None:
    """Write stream state."""
    state_dict = get_writeable_state_dict(
        tap_state,
        tap_stream_id,
        state_partition_context=state_partition_context,
    )
    state_dict[key] = val


def reset_state_progress_markers(stream_or_partition_state: dict) -> dict | None:
    """Wipe the state once sync is complete.

    For logging purposes, return the wiped 'progress_markers' object if it existed.
    """
    progress_markers = stream_or_partition_state.pop(PROGRESS_MARKERS, {})
    # Remove auto-generated human-readable note:
    progress_markers.pop(PROGRESS_MARKER_NOTE, None)
    # Return remaining 'progress_markers' if any:
    return progress_markers or None


def write_replication_key_signpost(
    stream_or_partition_state: dict,
    new_signpost_value: t.Any,  # noqa: ANN401
) -> None:
    """Write signpost value."""
    stream_or_partition_state[SIGNPOST_MARKER] = to_json_compatible(new_signpost_value)


def write_starting_replication_value(
    stream_or_partition_state: dict,
    initial_value: t.Any,  # noqa: ANN401
) -> None:
    """Write initial replication value to state."""
    stream_or_partition_state[STARTING_MARKER] = to_json_compatible(initial_value)


def get_starting_replication_value(stream_or_partition_state: dict) -> t.Any | None:  # noqa: ANN401
    """Retrieve initial replication marker value from state."""
    if not stream_or_partition_state:
        return None
    return stream_or_partition_state.get(STARTING_MARKER)


def increment_state(
    stream_or_partition_state: dict,
    *,
    latest_record: dict,
    replication_key: str,
    is_sorted: bool,
    check_sorted: bool,
) -> None:
    """Update the state using data from the latest record.

    Raises InvalidStreamSortException if is_sorted=True, check_sorted=True and unsorted
    data is detected in the stream.
    """
    progress_dict = stream_or_partition_state
    if not is_sorted:
        if PROGRESS_MARKERS not in stream_or_partition_state:
            stream_or_partition_state[PROGRESS_MARKERS] = {
                PROGRESS_MARKER_NOTE: "Progress is not resumable if interrupted.",
            }
            logger.warning(
                "Stream is assumed to be unsorted, progress is not resumable if "
                "interrupted",
                extra={"replication_key": replication_key},
            )
        progress_dict = stream_or_partition_state[PROGRESS_MARKERS]
    # TODO: Instead of forcing all values to be JSON-compatible strings and hope
    # we catch all cases, we should allow the stream to define how to
    # the values from the state and the record should be pre-processed.
    # https://github.com/meltano/sdk/issues/2753
    old_rk_value = to_json_compatible(progress_dict.get("replication_key_value"))
    new_rk_value = to_json_compatible(latest_record[replication_key])

    if new_rk_value is None:
        logger.warning("New replication value is null")
        return

    if old_rk_value is None or not check_sorted or new_rk_value >= old_rk_value:
        progress_dict["replication_key"] = replication_key
        progress_dict["replication_key_value"] = new_rk_value
        return

    if is_sorted:
        msg = (
            f"Unsorted data detected in stream. Latest value '{new_rk_value}' is "
            f"smaller than previous max '{old_rk_value}'."
        )
        raise InvalidStreamSortException(msg)


def _greater_than_signpost(
    signpost: _T,
    new_value: _T,
) -> bool:
    """Compare and return True if new_value is greater than signpost."""
    # fails if signpost and bookmark are incompatible types
    return new_value > signpost


def is_state_non_resumable(stream_or_partition_state: dict) -> bool:
    """Return True when state is non-resumable.

    This is determined by checking for a "progress marker" tag in the state artifact.
    """
    return PROGRESS_MARKERS in stream_or_partition_state


def finalize_state_progress_markers(stream_or_partition_state: dict) -> dict | None:
    """Promote or wipe progress markers once sync is complete.

    This marks any non-resumable progress markers as finalized. If there are
    valid bookmarks present, they will be promoted to be resumable.
    """
    signpost_value = stream_or_partition_state.pop(SIGNPOST_MARKER, None)
    stream_or_partition_state.pop(STARTING_MARKER, None)
    if (
        is_state_non_resumable(stream_or_partition_state)
        and "replication_key" in stream_or_partition_state[PROGRESS_MARKERS]
    ):
        # Replication keys valid (only) after sync is complete
        progress_markers = stream_or_partition_state[PROGRESS_MARKERS]
        stream_or_partition_state["replication_key"] = progress_markers.pop(
            "replication_key",
        )
        new_rk_value = progress_markers.pop("replication_key_value")
        if signpost_value and _greater_than_signpost(signpost_value, new_rk_value):
            new_rk_value = signpost_value
        stream_or_partition_state["replication_key_value"] = new_rk_value
    # Wipe and return any markers that have not been promoted
    return reset_state_progress_markers(stream_or_partition_state)


def log_sort_error(
    *,
    ex: Exception,
    log_fn: t.Callable,
    stream_name: str,
    current_context: types.Context | None,
    state_partition_context: types.Context | None,
    record_count: int,
    partition_record_count: int,
) -> None:
    """Log a sort error."""
    msg = f"Sorting error detected in '{stream_name}' on record #{record_count}. "
    if partition_record_count != record_count:
        msg += (
            f"Record was partition record "
            f"#{partition_record_count} with"
            f" state partition context {state_partition_context}. "
        )
    if current_context:
        msg += f"Context was {current_context!s}. "
    msg += str(ex)
    log_fn(msg)


class StateWriter:
    """Centralized state message writer that prevents duplicate state emissions.

    This class manages the writing of STATE messages to ensure that duplicate
    state messages are not emitted across multiple streams or tap-level operations.
    It tracks the last emitted state and only writes new messages when the state
    has actually changed.
    """

    def __init__(self, message_writer: GenericSingerWriter) -> None:
        """Initialize the StateWriter.

        Args:
            message_writer: The message writer instance (typically from tap)
                           that has a write_message method.
        """
        self._message_writer = message_writer
        self._last_emitted_state: types.TapState | None = None

    def write_state(self, state: types.TapState) -> None:
        """Write a state message if the state has changed.

        This method checks if the provided state is different from the last
        emitted state and only writes a STATE message if there are changes.

        Args:
            state: The current tap state to potentially emit.
        """
        # Check if state has changed since last emission
        if self._last_emitted_state is None or state != self._last_emitted_state:
            self._message_writer.write_message(StateMessage(value=state))
            self._last_emitted_state = copy.deepcopy(state)
