#  Copyright 2026 Synnax Labs, Inc.
#
#  Use of this software is governed by the Business Source License included in the file
#  licenses/BSL.txt.
#
#  As of the Change Date specified in that file, in accordance with the Business Source
#  License, use of this software will be governed by the Apache License, Version 2.0,
#  included in the file licenses/APL.txt.

from __future__ import annotations

from enum import Enum
from typing import Literal, TypeAlias, cast, overload
from uuid import uuid4

from freighter import (
    EOF,
    ExceptionPayload,
    Stream,
    WebsocketClient,
    decode_exception,
)
from freighter.transport import P
from freighter.websocket import Message
from pydantic import BaseModel

import synnax.channel.payload as channel
from synnax.exceptions import UnexpectedError
from synnax.framer.adapter import WriteFrameAdapter
from synnax.framer.codec import (
    HIGH_PERF_SPECIAL_CHAR,
    LOW_PERF_SPECIAL_CHAR,
    WSFramerCodec,
)
from synnax.framer.frame import CrudeFrame, FramePayload
from synnax.telem import CrudeSeries, CrudeTimeStamp, TimeSpan, TimeStamp
from synnax.telem.control import Authority, CrudeAuthority, Subject
from synnax.util.normalize import normalize


class WriterCommand(int, Enum):
    OPEN = 0
    WRITE = 1
    COMMIT = 2
    SET_AUTHORITY = 3


class WriterMode(int, Enum):
    PERSIST_STREAM = 1
    PERSIST = 2
    STREAM = 3


CrudeWriterMode: TypeAlias = (
    WriterMode
    | Literal["persist_stream"]
    | Literal["persist"]
    | Literal["stream"]
    | int
)


class WriterConfig(BaseModel):
    authorities: list[int] = [Authority.ABSOLUTE]
    control_subject: Subject = Subject(name="", key=str(uuid4()))
    start: TimeStamp | None = None
    keys: list[channel.Key]
    mode: WriterMode = WriterMode.PERSIST_STREAM
    err_on_unauthorized: bool = False
    enable_auto_commit: bool = True
    auto_index_persist_interval: TimeSpan = 1 * TimeSpan.SECOND


class WriterRequest(BaseModel):
    config: WriterConfig | None = None
    command: WriterCommand
    frame: FramePayload | None = None


class WriterResponse(BaseModel):
    command: WriterCommand
    end: TimeStamp | None
    err: ExceptionPayload


class WSWriterCodec(WSFramerCodec):
    def encode(self, data: BaseModel) -> bytes:
        if not isinstance(data, Message):
            raise TypeError(f"expected Message, got {type(data)}")
        if (
            data.type == "close"
            or data.payload is None
            or data.payload.command != WriterCommand.WRITE
        ):
            return bytes([LOW_PERF_SPECIAL_CHAR]) + self.lower_perf_codec.encode(data)
        encoded = self.codec.encode(data.payload.frame, 1)
        buf = bytearray(encoded)
        buf[0] = HIGH_PERF_SPECIAL_CHAR
        return bytes(buf)

    def decode(self, data: bytes, pld_t: type[P]) -> P:
        if data[0] == LOW_PERF_SPECIAL_CHAR:
            return self.lower_perf_codec.decode(data[1:], pld_t)
        frame = self.codec.decode(data, 1)
        msg: Message[WriterRequest] = Message(type="data")
        msg.payload = WriterRequest(command=WriterCommand.WRITE, frame=frame)
        return cast(P, msg)


def parse_writer_mode(mode: CrudeWriterMode) -> WriterMode:
    if mode == "persist_stream":
        return WriterMode.PERSIST_STREAM
    if mode == "persist":
        return WriterMode.PERSIST
    if mode == "stream":
        return WriterMode.STREAM
    if isinstance(mode, WriterMode):
        return mode
    if isinstance(mode, int):
        try:
            return WriterMode(mode)
        except:
            ...
    raise ValueError(f"invalid writer mode {mode}")


ALWAYS_INDEX_PERSIST_ON_AUTO_COMMIT: TimeSpan = TimeSpan(-1)


class WriterClosed(Exception): ...


class Writer:
    """A writer is used to write telemetry to a set of channels in time order. It
    should not be constructed directly, and should instead be created using the Synnax
    client's open_writer method.

    The writer is a streaming protocol that is heavily optimized for performance. This
    comes at the cost of increased complexity, and should only be used directly when
    writing large volumes of data (such as recording telemetry from a sensor or
    ingesting data from a file). Simpler methods (such as client.write()) should be
    used in most cases.

    For a detailed guide on writing data to Synnax, see
    https://docs.synnaxlabs.com/reference/concepts/writes. A rough summary of the write
    process is detailed below:

    1. The writer is opened with a starting timestamp and a list of channel keys (or
    names). The writer will fail to open if the starting timestamp overlaps with any
    existing telemetry for any of the channels specified. If the writer is opened
    successfully, the caller is then free to write frames to the writer.

    2. To write a frame, the caller can use the write method and follow the validation
    rules described in the method's documentation. This process is asynchronous,
    meaning that write will return before the frame has been written to the cluster. This
    also means that the writer can accumulate an error after write is called. If the
    writer accumulates an error, subsequent calls to `write` will repeatedly throw that
    error. If write throws an error, the writer can no longer be used, and it must
    be closed and re-opened to continue writing.

    3. To commit the written frames to the database, the caller can call the commit
    method. Unlike write, commit is synchronous, meaning that it will not return until
    all frames have been committed to the database. If commit fails or the writer
    has accumulated an error through previous calls to commit or write, then commit
    will throw that error. If commit throws an error, the writer can no longer be used,
    and it must be closed and re-opened to continue writing. Commit can be called
    several times throughout a writer's lifetime, and will commit all frames written
    since the last commit.

    4. A writer MUST be closed after use in order to prevent resource leaks. Close
    should typically be called in a 'finally' block. If the writer has accumulated an
    error, close will raise the accumulated error.
    """

    _stream: Stream[WriterRequest, WriterResponse]
    _adapter: WriteFrameAdapter
    _close_exc: Exception | None = None

    start: CrudeTimeStamp

    def __init__(
        self,
        start: CrudeTimeStamp,
        client: WebsocketClient,
        adapter: WriteFrameAdapter,
        name: str = "",
        authorities: CrudeAuthority | list[CrudeAuthority] = Authority.ABSOLUTE,
        mode: CrudeWriterMode = WriterMode.PERSIST_STREAM,
        err_on_unauthorized: bool = False,
        enable_auto_commit: bool = True,
        auto_index_persist_interval: TimeSpan = 1 * TimeSpan.SECOND,
        use_experimental_codec: bool = True,
    ) -> None:
        self.start = start
        self._adapter = adapter
        if use_experimental_codec:
            client = client.with_codec(WSWriterCodec(adapter.codec))
        self._stream = client.stream("/frame/write", WriterRequest, WriterResponse)
        config = WriterConfig(
            control_subject=Subject(name=name, key=str(uuid4())),
            keys=self._adapter.keys,
            start=TimeStamp(self.start),
            authorities=normalize(authorities),
            mode=parse_writer_mode(mode),
            err_on_unauthorized=err_on_unauthorized,
            enable_auto_commit=enable_auto_commit,
            auto_index_persist_interval=auto_index_persist_interval,
        )
        exc = self._stream.send(
            WriterRequest(command=WriterCommand.OPEN, config=config)
        )
        if exc is not None:
            raise exc
        _, exc = self._stream.receive()
        if exc is not None:
            raise exc

    @overload
    def write(
        self, channels_or_data: channel.Key | str, series: CrudeSeries
    ) -> None: ...

    @overload
    def write(
        self,
        channels_or_data: (
            list[channel.Key] | tuple[channel.Key] | list[str] | tuple[str]
        ),
        series: list[CrudeSeries],
    ) -> None: ...

    @overload
    def write(
        self,
        channels_or_data: CrudeFrame,
    ) -> None: ...

    def write(
        self,
        channels_or_data: (
            str
            | channel.Key
            | list[channel.Key]
            | tuple[channel.Key]
            | list[str]
            | tuple[str]
            | CrudeFrame
        ),
        series: CrudeSeries | list[CrudeSeries] | None = None,
    ) -> None:
        """Writes the given data to the database. The formats are listed below. Before
        we get into them, here are some important terms to know.

            1. Channel ID -> the key or name of the channel(s) you're writing to.
            2. Series or CrudeSeries -> the data for that channel, which can be
            represented as a synnax Series type, a numpy array, or a simple Python
            list. You can also provide a single numeric (or, in the case of variable
            length types, a string or JSON) value and Synnax will convert it into a
            Series for you.

        Here are the formats you can use to write data to a Synnax cluster:

            1. Channel ID and a single series: Writes the series for the given channel.
            2. A list of channel ids and their corresponding series: Assumes a
            one to one mapping of ids to series i.e. the channel id at index i
            corresponds to the series at index i.
            3. A Synnax Frame (see the Frame documentation for more).
            4. A dictionary of channel ids to series i.e. write the series for the
            given channel id.
            5. A pandas DataFrame where the columns are the channel ids and the rows
            are the series to write.
            6. A dictionary of channel ids to a single
            numeric value. Synnax will convert this into a series for you.

        There are a few important rules to keep in mind when writing data to a Synnax
        cluster:

            1. Have exactly one array for each key in the list of keys provided to the
            writer's open method.
            2. Have equal length arrays for each key.
            3. When writing to an index (i.e. TimeStamp) channel, the values must be
            monotonically increasing.

        :raises: If the writer has accumulated an error. Once write throws,
        all subsequent calls to write and/or commit will also fail. At this point,
        the writer must be closed and re-opened in order to continue writing.
        """
        if self._close_exc is not None:
            raise self._close_exc
        frame = self._adapter.adapt(channels_or_data, series)
        try:
            self._exec(
                WriterRequest(command=WriterCommand.WRITE, frame=frame.to_payload()),
                timeout=0,
            )
        except TimeoutError:
            ...

    @overload
    def set_authority(self, value: CrudeAuthority) -> None:
        """Sets the authority level for all channels the writer is writing to.

        :param value: The authority level to set (0-255). Use Authority.ABSOLUTE (255)
            for exclusive control or Authority.DEFAULT (1) for standard priority.
        """
        ...

    @overload
    def set_authority(
        self,
        value: channel.Key | str,
        authority: CrudeAuthority,
    ) -> None:
        """Sets the authority level for a single channel.

        :param value: The key or name of the channel to set authority for.
        :param authority: The authority level to set (0-255).
        """
        ...

    @overload
    def set_authority(
        self,
        value: dict[channel.Key | str | channel.Payload, CrudeAuthority],
    ) -> None:
        """Sets the authority level for multiple channels.

        :param value: A dictionary mapping channel keys/names to their authority levels.
        """
        ...

    def set_authority(
        self,
        value: (
            dict[channel.Key | str | channel.Payload, CrudeAuthority]
            | channel.Key
            | str
            | CrudeAuthority
        ),
        authority: CrudeAuthority | None = None,
    ) -> None:
        """Sets the control authority for channels the writer is writing to. Authority
        determines which writer takes precedence when multiple writers attempt to write
        to the same channel. Higher authority values take priority over lower ones.

        This method has three call signatures:

        1. Set authority for all channels:
            writer.set_authority(Authority.ABSOLUTE)

        2. Set authority for a single channel by key or name:
            writer.set_authority("my_channel", Authority.ABSOLUTE)
            writer.set_authority(channel_key, 200)

        3. Set authority for multiple channels via dictionary:
            writer.set_authority({
                "channel_a": Authority.ABSOLUTE,
                "channel_b": 100,
            })

        Authority levels range from 0 to 255:
            - Authority.ABSOLUTE (255): Highest priority. No other writer can take
              control while this writer is active.
            - Authority.DEFAULT (1): Standard priority for normal operations.
            - Custom values (0-255): Higher values take precedence over lower ones.

        :param value: Either an authority level (int) to apply to all channels, a
            channel key/name to set authority for a single channel, or a dictionary
            mapping channels to their authority levels.
        :param authority: The authority level when setting for a single channel.
            Required when value is a channel key or name, ignored otherwise.
        :raises ValueError: If a single channel is specified without an authority value.
        :raises: Any accumulated error from previous writer operations.
        """
        if self._close_exc is not None:
            raise self._close_exc
        if isinstance(value, int) and authority is None:
            cfg = WriterConfig(keys=[], authorities=[value])
        else:
            auth_map: dict[channel.Key | str | channel.Payload, CrudeAuthority]
            if isinstance(value, (channel.Key, str)):
                if authority is None:
                    raise ValueError(
                        "authority must be provided when setting a single channel"
                    )
                auth_map = {value: authority}
            elif isinstance(value, dict):
                auth_map = value
            else:
                raise ValueError(f"unexpected authority value type: {type(value)}")
            resolved = self._adapter.adapt_dict_keys(auth_map)
            cfg = WriterConfig(
                keys=list(resolved.keys()),
                authorities=list(resolved.values()),
            )
        self._exec(WriterRequest(command=WriterCommand.SET_AUTHORITY, config=cfg))

    def commit(self) -> TimeStamp:
        """Commits the written frames to the database. Commit is synchronous, meaning
        that it will not return until all frames have been committed to the database.

        :returns: The timestamp of the last sample written to the writer.
        :raises: An exception if the commit fails or any calls to previous writer
            methods have thrown an exception. Once commit fails, the writer must be
            closed and re-opened to continue use.
        """
        if self._close_exc is not None:
            raise self._close_exc
        res = self._exec(WriterRequest(command=WriterCommand.COMMIT))
        if res is None:
            raise UnexpectedError("commit did not receive a response")
        if res.end is None:
            raise UnexpectedError("commit response missing end timestamp")
        return res.end

    def close(self) -> None:
        """Closes the writer, raising any accumulated error encountered during
        operation. A writer MUST be closed after use, and this method should probably
        be placed in a 'finally' block.
        """
        return self._close(None)

    def _close(self, exc: Exception | None) -> None:
        if self._close_exc is not None:
            if isinstance(self._close_exc, WriterClosed):
                return
            raise self._close_exc
        self._close_exc = exc
        self._stream.close_send()
        while True:
            if self._close_exc is not None:
                if isinstance(self._close_exc, WriterClosed):
                    return
                raise self._close_exc
            res, recv_exc = self._stream.receive()
            if recv_exc is not None:
                self._close_exc = (
                    WriterClosed() if isinstance(recv_exc, EOF) else recv_exc
                )
            elif res is not None:
                self._close_exc = decode_exception(res.err)

    def _exec(
        self, req: WriterRequest, timeout: int | None = None
    ) -> WriterResponse | None:
        send_exc = self._stream.send(req)
        if send_exc is not None:
            self._close(send_exc)
            return None
        while True:
            res, recv_exc = self._stream.receive(timeout)
            if recv_exc is not None:
                self._close(recv_exc)
                return None
            if res is None:
                continue
            decoded_exc = decode_exception(res.err)
            if decoded_exc is not None:
                self._close(decoded_exc)
                return None
            if res.command == req.command:
                return res

    def __enter__(self) -> Writer:
        return self

    def __exit__(self, exc_type: object, exc_value: object, traceback: object) -> None:
        self.close()
