from __future__ import annotations

import json

import pyarrow.flight as flight

from tempora.utils.batch import TSSegment
from tempora.utils.dataclass import get_args_dict
from tempora.utils.server import get_connection
from tempora.utils.time import Length


def get_stream_info() -> dict[str, str]:
    action = flight.Action('get_stream_info', b'')
    client, auth_token = get_connection()
    call_options = flight.FlightCallOptions(headers=[auth_token])

    result = next(client.do_action(action, call_options))
    stream_info_str = result.body.to_pybytes().decode()
    return json.loads(stream_info_str)


def get_segments(context_len: Length, num_segments: int) -> list[TSSegment]:
    client, auth_token = get_connection()
    args = {'context_len': get_args_dict(context_len), 'num_segments': num_segments}
    json_cmd = json.dumps(args)
    desc = flight.FlightDescriptor.for_command(json_cmd.encode())
    call_options = flight.FlightCallOptions(headers=[auth_token])
    _, flight_reader = client.do_exchange(desc, call_options)

    segments: list[TSSegment] = []
    while True:
        try:
            chunk = flight_reader.read_chunk()
            segment = TSSegment.from_chunk(chunk=chunk)
            if segment:
                segments.append(segment)
        except StopIteration:
            break
    return segments
