"""
XmodemOperation - Reusable xmodem upload operations

This class provides reusable xmodem upload functionality that can be
shared across different task implementations.
"""

from __future__ import annotations
import time
import xmodem
import threading
from typing import Optional
from pathlib import Path
from datetime import datetime
from enum import Enum
from pydantic import BaseModel, Field
from logging import getLogger
import os

from lib.bytes_utils import bytes_to_hex
from deviceio.serial_io_module import SerialIoModule
from csv_report.csv_report_helper import CsvReportHelper
from csv_report_struct.xmodem_spend_time import XmodemSpendTimeReport
from csv_report_struct.program_full_spend_time import ProgramFullSpendTimeReport
from Task.TaskOperation.operation_base import OperationBase
from pydantic_argparse.argparse_helper_funcs import ArgExtra
from error.atmosic_error import XmodemOperationError, DeviceIoModuleError
from error.errorcodes import XmodemOperationErrorCode

logger = getLogger(__name__)
START_XMODEM = b'U' #0x55
ASK_VERSION = b'V' #0x56
XMODEM_VERSION = b"\x33"

class UploadFirstAction(Enum):
    """First action to perform after xmodem upload"""
    WAIT_CRC = 'wait_crc'
    ASK_VERSION = 'ask_ver'
    SKIP = 'skip'

class ReportContext(BaseModel):
    start_timestamp: datetime
    full_report: ProgramFullSpendTimeReport

class XmodemOperation(OperationBase):
    """Reusable xmodem upload operations"""
    class InitContext(BaseModel):
        xmodem_serial:str = Field(
            description="Serial port for xmodem communication"
        )
        ram_image: str = Field(
            default="binfile/uart_loader.bin",
            description="Path to RAM image file to upload via xmodem"
        )
        xmodem_128: bool = Field(
            default=False,
            description="Use XMODEM128 instead of XMODEM1k"
        )
        upload_first_action: UploadFirstAction = Field(
            default=UploadFirstAction.ASK_VERSION,
            description="First action to perform after xmodem upload"
        )
        wait_time_before_upload_ms: int = Field(
            default=200,
            description="Wait time before upload in milliseconds"
        )
        upload_timeout_ms: int = Field(
            default=10000,
            description="Timeout for xmodem upload in milliseconds"
        )
        check_join_timeout_ms: int = Field(
            default=11,
            description="Timeout for join check in milliseconds"
        )
        xmodem_sync_timeout_s: float = Field(
            default=5,
            description="Timeout for xmodem sync in seconds"
        )
        ask_version_timeout_ms: int = Field(
            default=100,
            description="Timeout for ask version in milliseconds"
        )

        @property
        def wait_time_before_upload_s(self) -> float:
            return self.wait_time_before_upload_ms / 1000
        @property
        def upload_timeout_s(self) -> float:
            return self.upload_timeout_ms / 1000
        @property
        def check_join_interval_s(self) -> float:
            return self.check_join_timeout_ms / 1000

    def __init__(self, ctx: InitContext,
                 start_timestamp: datetime | None = None,
                 full_report: ProgramFullSpendTimeReport | None = None):
        """Initialize XmodemOperation"""

        self.ctx = ctx
        self.start_timestamp = start_timestamp
        self.full_report = full_report

        self.xmodem_serial_module: Optional[SerialIoModule] = None
        self.send_fake_response_c = False

        logger.info(f"XmodemOperation initialized:")
        logger.info(f"  Xmodem serial: {ctx.xmodem_serial}")
        logger.info(f"  RAM image: {ctx.ram_image}")
        logger.info(f"  XMODEM128 mode: {ctx.xmodem_128}")
        logger.info(f"  Upload first action: {ctx.upload_first_action}")
        logger.info(f"  Wait time before upload:"
                      f" {ctx.wait_time_before_upload_ms}ms")
    
    def execute_xmodem_sync(self) -> None:
        """Execute xmodem sync phase"""
        try:
            self._init_xmodem_serial()
            assert self.xmodem_serial_module is not None, \
                "Xmodem serial module must be initialized"
            stime = datetime.now()
            while True:
                if (datetime.now() - stime).total_seconds() > \
                    self.ctx.xmodem_sync_timeout_s:
                    raise XmodemOperationError(
                        "Timeout waiting for version response",
                        XmodemOperationErrorCode.XMODEM_ASK_VERSION_TIMEOUT
                    )
                logger.info("Send START_XMODEM")
                self.xmodem_serial_module.write(START_XMODEM)
                time.sleep(0.1)
                logger.info("Send CAN")
                self.xmodem_serial_module.write(xmodem.CAN)
                time.sleep(0.1)
                self.xmodem_serial_module.clear_read_buffer()
                logger.info("Send START_XMODEM")
                self.xmodem_serial_module.write(START_XMODEM)
                time.sleep(self.ctx.wait_time_before_upload_s)
                logger.info("Send ASK_VERSION")
                self.xmodem_serial_module.write(ASK_VERSION)
                try:
                    version_rsp = self.xmodem_serial_module.read(
                        size=1, timeout= self.ctx.ask_version_timeout_ms / 1000)
                except DeviceIoModuleError:
                    logger.debug("Timeout waiting for version response")
                    time.sleep(0.05)
                    logger.debug("Retrying for ask version...")
                    continue
                if version_rsp is not None:
                    break
            logger.info(f"Version response: {bytes_to_hex(version_rsp)}")
            
        finally:
            self._close_xmodem_serial()

    
    def execute_ram_mode_phase(self) -> None:
        """Execute RAM mode phase (xmodem upload only - reset should be done by caller)

        Args:
            start_timestamp: Command start timestamp for reporting
        """
        logger.info("=== Executing RAM mode phase ===")
        xmodem_start_time = datetime.now()

        try:
            # Initialize phase (caller should have already reset device to bboot mode)
            init_stime = datetime.now()
            self._init_xmodem_serial()
            init_spend_time = int((datetime.now() - init_stime).total_seconds() * 1000)
            
            # Upload phase
            upload_stime = datetime.now()
            self._upload_ram_code()
            upload_spend_time = int((datetime.now() - upload_stime).total_seconds() * 1000)
            
            # Create and populate report
            report = XmodemSpendTimeReport()
            report.cmd_start_timestamp = self.start_timestamp
            report.serial = self.ctx.xmodem_serial
            report.ram_file = self.ctx.ram_image
            report.ram_size = Path(self.ctx.ram_image).stat().st_size
            report.xmodem_flow_start_timestamp = xmodem_start_time
            report.init_ms = init_spend_time
            report.upload_ms = upload_spend_time
            CsvReportHelper.append(report)

            if self.full_report is not None:
                self.full_report.xmodem_init_ms = init_spend_time
                self.full_report.xmodem_upload_ms = upload_spend_time
            
            logger.info("=== RAM mode phase completed successfully ===")
            
        finally:
            self._close_xmodem_serial()

    def _init_xmodem_serial(self):
        """Initialize xmodem serial port"""
        logger.info("=== Initialize xmodem serial port ===")
        self.xmodem_serial_module = SerialIoModule(
            port=self.ctx.xmodem_serial,
            baudrate=115200,
            rtscts=False,
            stream_type="ascii",
            timeout=3
        )
        logger.info(f"Opening xmodem serial port: {self.ctx.xmodem_serial}")
        self.xmodem_serial_module.open()

        # Clear any remaining data in buffer
        self.xmodem_serial_module.clear_read_buffer()

        logger.info(f"Xmodem serial port {self.ctx.xmodem_serial} initialized successfully")

    def setup_fake_response_c(self):
        """Setup fake response for 'C' character"""
        self.send_fake_response_c =\
            self.ctx.upload_first_action != UploadFirstAction.WAIT_CRC

    def getc(self, size, timeout=5):
        assert self.xmodem_serial_module is not None,\
            "Xmodem serial module must be initialized"
        if self.send_fake_response_c:
            self.send_fake_response_c = False
            logger.debug(
                "XMODEM getc: for upload enhance, respond with 'C'")
            return b'C'
        start_time = time.time()
        result = self.xmodem_serial_module.read(size, timeout)
        if result is None:
            raise XmodemOperationError(
                f"XMODEM getc timeout after {timeout}s, size={size}",
                XmodemOperationErrorCode.XMODEM_GETC_TIMEOUT
            )
            
        logger.debug(f"XMODEM getc: read {len(result)} bytes"
                     f" in {time.time() - start_time:.3f}s")
        return result

    def putc(self, data, timeout=5):
        assert self.xmodem_serial_module is not None,\
            "Xmodem serial module must be initialized"
        # Note: timeout parameter is required by xmodem library interface
        # but not used by our SerialIoModule implementation
        start_time = time.time()
        bytes_written = self.xmodem_serial_module.write(data)
        logger.debug(f"XMODEM putc: wrote {bytes_written} bytes"
                     f" in {time.time() - start_time:.3f}s")
        return bytes_written

    def xmodem_upload_with_timeout(self, modem: xmodem.XMODEM, file_handle,
                                   timeout_seconds):
        """Upload file using xmodem with timeout protection."""
        result = {'success': False, 'error': None}

        def upload_thread():
            try:
                self.setup_fake_response_c()
                if modem.send(file_handle):
                    result['success'] = True
                    logger.info("Xmodem upload completed successfully")
                else:
                    logger.error("Xmodem upload failed")
                    result['error'] = XmodemOperationError(
                        "Xmodem upload failed",
                        XmodemOperationErrorCode.XMODEM_UPLOAD_FAILED
                    )
            except Exception as e:
                result['error'] = e

        thread = threading.Thread(target=upload_thread)
        thread.daemon = True
        thread.start()
        join_start_time = datetime.now()
        while thread.is_alive():
            elapsed_time =\
                (datetime.now() - join_start_time).total_seconds()
            if elapsed_time > timeout_seconds:
                raise XmodemOperationError(
                    f"Xmodem upload timed out after {timeout_seconds} seconds",
                    XmodemOperationErrorCode.XMODEM_UPLOAD_TIMEOUT)
            thread.join(timeout=self.ctx.check_join_interval_s)

        if result['error']:
            raise result['error']

        if not result['success']:
            raise XmodemOperationError(
                "Xmodem upload failed for unknown reason",
                XmodemOperationErrorCode.XMODEM_UPLOAD_FAILED
            )

    def _ask_version(self) -> None:
        """Ask version from device"""
        assert self.xmodem_serial_module is not None, \
            "Xmodem serial module must be initialized"
        logger.debug("Send START_XMODEM")
        self.xmodem_serial_module.write(START_XMODEM)
        time.sleep(self.ctx.wait_time_before_upload_s)
        logger.debug("Send ASK_VERSION")
        self.xmodem_serial_module.clear_read_buffer()
        self.xmodem_serial_module.write(ASK_VERSION)
        try:
            version_rsp = self.xmodem_serial_module.read(
                size=1, timeout= self.ctx.ask_version_timeout_ms / 1000)
            if version_rsp == XMODEM_VERSION:
                logger.debug(f"Version response: {bytes_to_hex(version_rsp)}")
                return
            else:
                logger.warning(
                    f"Unexpected version response: {bytes_to_hex(version_rsp)}")
        except DeviceIoModuleError:
            logger.warning("Timeout waiting for version response")
        
        time.sleep(self.ctx.wait_time_before_upload_s)
        logger.debug("Send ASK_VERSION again")
        self.xmodem_serial_module.clear_read_buffer()
        self.xmodem_serial_module.write(ASK_VERSION)
        try:
            version_rsp = self.xmodem_serial_module.read(
                size=1, timeout= self.ctx.ask_version_timeout_ms / 1000)
            if version_rsp == XMODEM_VERSION:
                logger.debug(f"Version response: {bytes_to_hex(version_rsp)}")
                return
            else:
                logger.error(
                    f"Unexpected version response: {bytes_to_hex(version_rsp)}")
                logger.debug("Send CAN")
                self.xmodem_serial_module.write(xmodem.CAN)
                raise XmodemOperationError(
                    "Unexpected version response",
                    XmodemOperationErrorCode.XMODEM_UNEXPECTED_VERSION
                )
        except DeviceIoModuleError:
            logger.error("Timeout waiting for version response")
            logger.debug("Send CAN")
            self.xmodem_serial_module.write(xmodem.CAN)
            FILE_NAME = "debug_time.txt"
            # If file does not exist, create it with default value 180
            if os.path.exists(FILE_NAME):
                with open(FILE_NAME, "r", encoding="utf-8") as f:
                    content = f.read().strip()  # Remove whitespace/newline
                debug_time = int(content)
                if debug_time > 0:
                    logger.warning(f"For debugging, sleep {debug_time}s"
                                   " then continue")
                    time.sleep(debug_time)
                elif debug_time < 0:
                    logger.error("Debug time < 0, keep looping until Ctrl+C")
                    while True:
                        time.sleep(1)
            raise XmodemOperationError(
                "Timeout waiting for version response",
                XmodemOperationErrorCode.XMODEM_ASK_VERSION_TIMEOUT
            )


    def _upload_ram_code(self):
        """Upload RAM code via xmodem with timeout protection"""
        logger.info("=== Upload RAM code via xmodem ===")

        assert self.xmodem_serial_module is not None, \
            "Xmodem serial module must be initialized"

        logger.info(f"Starting RAM code upload: {self.ctx.ram_image}")
        if self.ctx.upload_first_action == UploadFirstAction.ASK_VERSION:
            self._ask_version()
        if not self.ctx.xmodem_128:
            logger.info("Using XMODEM1k")
            modem = xmodem.XMODEM1k(self.getc, self.putc)
        else:
            logger.info("Using XMODEM128")
            modem = xmodem.XMODEM(self.getc, self.putc)

        try:
            start_time = time.time()
            with open(self.ctx.ram_image, 'rb') as f:
                file_size = Path(self.ctx.ram_image).stat().st_size
                logger.info(f"Uploading file: {self.ctx.ram_image}"
                              f" ({file_size} bytes)")
                self.xmodem_upload_with_timeout(modem, f,
                                                self.ctx.upload_timeout_s)
            spent_time = time.time() - start_time
            logger.info(f"RAM code upload completed successfully, took"
                          f" {spent_time:.3f}s")
        except TimeoutError as e:
            raise XmodemOperationError(
                f"RAM code upload timed out: {e}\n"
                "Device may not be ready for xmodem transfer or "
                "there is a hardware issue",
                XmodemOperationErrorCode.XMODEM_UPLOAD_TIMEOUT
            )

    def _close_xmodem_serial(self):
        """Close xmodem serial port"""
        if self.xmodem_serial_module and self.xmodem_serial_module.is_open:
            self.xmodem_serial_module.close()
            logger.info("Xmodem serial port closed")
            self.xmodem_serial_module = None
