from __future__ import annotations
from pydantic import Field, BaseModel

from pydantic_argparse.argparse_helper_funcs import ArgExtra
from Task.TaskContextBase.task_xmodem_context_base import TaskXmodemContextBase
from Task.TaskOperation.uart_loader_operation import UartLoaderOperation

class TaskUartLoaderContextBase(TaskXmodemContextBase):
    serial: str = Field(
        description="Serial port for uart loader communication"
    )

    uart_loader_init_baudrate: int = ArgExtra(Field(
        default=115200,
        description="Initial baudrate for uart loader communication"
    ), engineer_mode=True)

    xmodem_serial: str = ArgExtra(Field(
        default="",
        description="Serial port for XModem communication"
    ), engineer_mode=True)

    def to_uart_loader_init_context(self) -> UartLoaderOperation.InitContext:
        return UartLoaderOperation.InitContext(
            uart_loader_serial=self.serial,
            uart_loader_init_baudrate=self.uart_loader_init_baudrate
        )

    def to_xmodem_init_context(self):
        if self.xmodem_serial == "":
            self.xmodem_serial = self.serial
        return super().to_xmodem_init_context()

class TaskFuncChangeBaudrateContextBase(BaseModel):
    baudrate: int = Field(
        default=1000000,
        description="Baudrate for uart loader communication"
    )
    change_baudrate_delay_us: int = ArgExtra(Field(
        default=1000,
        description="Delay in microseconds after baudrate change"
    ), engineer_mode=True)

    def to_change_baudrate_context(self) \
        -> UartLoaderOperation.ChangeBaudrateContext:

        return UartLoaderOperation.ChangeBaudrateContext(
            baudrate=self.baudrate,
            delay_us=self.change_baudrate_delay_us
        )