from __future__ import annotations

import base64
from io import BytesIO
from typing import TYPE_CHECKING
from urllib.parse import urlencode

import polars as pl
import zstandard as zstd

if TYPE_CHECKING:
    from viewer import AxisProperties, ColorTransform

COMPRESSOR = zstd.ZstdCompressor(level=22)
DECOMPRESSOR = zstd.ZstdDecompressor()


def encode_url(
    title: str,
    molecule_id: str,
    colors: ColorTransform,
    axis_properties: AxisProperties,
    data: pl.DataFrame,
    description: str = "",
):
    encode_dict = dict(title=title, molecule_id=molecule_id)
    encode_dict.update({**colors.model_dump(), **axis_properties.model_dump()})

    csv_str = data.write_csv(float_precision=4, float_scientific=True)
    compressed = COMPRESSOR.compress(csv_str.encode())
    base64_text = base64.b64encode(compressed).decode("utf8")
    encode_dict["data"] = base64_text
    if description:
        encode_dict["description"] = description

    return urlencode(encode_dict)


def decode_data(base64_text) -> pl.DataFrame:
    decoded_bytes = base64.b64decode(base64_text)
    decompressed = DECOMPRESSOR.decompress(decoded_bytes)

    bio = BytesIO(decompressed)
    data = pl.read_csv(bio)
    bio.close()

    return data