from __future__ import annotations
from logging import getLogger

from enum import Enum
from pydantic import BaseModel, Field
from reset_support.jlinkob_reset_support import JlinkobResetSupport
from reset_support.rts_reset_support import RtsResetSupport
from reset_support.rts_relay_support import RtsRelaySupport
from reset_support.plugin_reset_support import PluginResetSupport
from pydantic_argparse.argparse_helper_funcs import ArgExtra

logger = getLogger(__name__)

class UartLoaderBootSource(Enum):
    """Uart loader boot source enumeration"""
    RAM = "ram"      # Reset to bboot mode -> xmodem upload -> uart loader
    RRAM = "rram"    # Reset to normal mode -> direct uart loader


class ResetOperation:
    """Reusable reset operations"""
    class InitContext(BaseModel):
        reset_support:str = Field(
            description="Reset support implementation (jlinkob, rts)"
        )
        uart_loader_mode: UartLoaderBootSource = Field(
            description="Uart loader mode (ram, rram)"
        )
        jlink_sn: str = Field(
            description="J-Link serial number (for jlinkob reset support)"
        )
        reset_serial: str = Field(
            description="Serial port for RTS reset support"
        )
        plugin_start_py: str = Field(
            description="Python script to reset device before start"
        )
        plugin_end_py: str = Field(
            description="Python script to reset device after end"
        )
        test_times: int = Field(
            description="Test times"
        )
        max_test_times: int = Field(
            description="Max test times"
        )
    
    def __init__(self, ctx: InitContext):
        self.ctx = ctx
        if self.ctx.reset_support == "none":
            self.reset_support = None
        elif self.ctx.reset_support == "jlinkob":
            self.reset_support = JlinkobResetSupport(self.ctx.jlink_sn)
        elif self.ctx.reset_support == "rts":
            self.reset_support = RtsResetSupport(self.ctx.reset_serial)
        elif self.ctx.reset_support == "rts_relay":
            self.reset_support = RtsRelaySupport(self.ctx.reset_serial)
        elif self.ctx.reset_support == "plugin":
            self.reset_support = PluginResetSupport(self.ctx.plugin_start_py,
                                                    self.ctx.plugin_end_py,
                                                    self.ctx.test_times,
                                                    self.ctx.max_test_times)
        else:
            raise ValueError(f"Unknown reset support: {self.ctx.reset_support}")
        
    def execute_reset_phase(self):
        """Execute reset phase based on uart loader mode"""
        if self.ctx.uart_loader_mode == UartLoaderBootSource.RAM:
            logger.info("Executing RAM mode reset phase")
            self._reset_to_bboot_mode()
        elif self.ctx.uart_loader_mode == UartLoaderBootSource.RRAM:
            logger.info("Executing RRAM mode reset phase")
            self._reset_to_normal_mode()

    def _reset_to_bboot_mode(self):
        """Reset device to bboot mode (for RAM mode)"""
        if self.reset_support is None:
            logger.info("Please ensure device is in bboot mode")
            return

        logger.info("Resetting device to bboot mode...")
        self.reset_support.reset_bboot_mode()
        logger.info("Device reset to bboot mode successfully")

    def _reset_to_normal_mode(self):
        """Reset device to normal mode (for RRAM mode)"""
        if self.reset_support is None:
            logger.info("Please ensure device is in normal mode")
            return

        logger.info("Resetting device to normal mode...")
        self.reset_support.reset_normal_boot()
        logger.info("Device reset to normal mode successfully")

    def module_stop(self, is_test_success: bool):
        if self.reset_support is not None:
            self.reset_support.module_stop(is_test_success)
