# minimal_connector.py
from __future__ import annotations

import asyncio
import uuid
from dynalab_protocol.models import (
    ConnectorHello,
    DynaLabHello,
    HandshakeAccepted,
    HandshakeRejected,
    Heartbeat,
    SignalDescriptor,
    machine_timestamp_ms,
)
from dynalab_protocol.wire import read_message, write_message


HOST = "127.0.0.1"
PORT = 8765


async def main() -> None:
    reader, writer = await asyncio.open_connection(HOST, PORT)

    dynalab_hello: DynaLabHello
    handshake: HandshakeAccepted

    peer = writer.get_extra_info("peername")
    print(f"Connected to DynaLab at {peer}")

    try:
        message = await asyncio.wait_for(read_message(reader), timeout=5.0)

        if not isinstance(message, DynaLabHello):
            raise ValueError(f"Expected dynalab_hello, received {message.type}")

        dynalab_hello = message
        print(f"Received DynaLab Hello\n {dynalab_hello}")

        signals: list[SignalDescriptor] = [
            SignalDescriptor(
                id=uuid.uuid4(), name="lambda", type="number", min_value=0, max_value=5
            )
        ]

        connector_hello = ConnectorHello(
            connector_uuid=uuid.uuid4(),
            connector_name="Minimal Connector",
            connector_version="a0.0.1",
            signals=signals,
        )
        await write_message(writer, connector_hello)

        message = await asyncio.wait_for(read_message(reader), timeout=5.0)

        if isinstance(message, HandshakeAccepted):
            handshake = message
            print(f"Received connection handshake: {handshake}")
        elif isinstance(message, HandshakeRejected):
            print(f"Connector handshake rejected: {message.reason}")
            return

        while True:
            message = await read_message(reader)
            print(f"Received {message}")

            if isinstance(message, Heartbeat):
                if message.sequence < 10:
                    message.return_timestamp = machine_timestamp_ms()
                    await write_message(writer, message)

    except asyncio.CancelledError:
        raise

    finally:
        print("Disconnecting")
        writer.close()
        await writer.wait_closed()


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        print("\nConnector stopped")