#########################################################################################
#
# @file app.py
#
# @brief UART loader host Tool - Refactored with task-based architecture
#
# Copyright (C) Atmosic Technologies, Inc, 2025
#
#########################################################################################

from __future__ import annotations
import sys
from argparse import ArgumentParser, Namespace, SUPPRESS
from logging import getLogger
from os import getpid
from typing import Any, Type
from pathlib import Path
import tomllib


from pydantic_argparse.argparse_helper_funcs import ArgumentSubParserHelper
from lib.polling_thread_base import PollingThreadBase
from Task.task_program_bin_file import (
    TaskProgramBinFile, TaskProgramBinFileContext
)
from Task.task_program_atm_file import (
    TaskProgramAtmFile, TaskProgramAtmFileContext
)
from Task.task_dump_bin_file import (
    TaskDumpBinFile, TaskDumpBinFileContext
)
from Task.task_atm_append_flash_bin import (
    TaskAtmAppendFlashBin, TaskAtmAppendFlashBinContext
)
from Task.task_atm_append_flash_erase import (
    TaskAtmAppendFlashErase, TaskAtmAppendFlashEraseContext
)
from Task.task_atm_append_rram_bin import (
    TaskAtmAppendRramBin, TaskAtmAppendRramBinContext
)
from Task.task_atm_append_rram_erase import (
    TaskAtmAppendRramErase, TaskAtmAppendRramEraseContext
)
from Task.task_atm_show import (
    TaskAtmShow, TaskAtmShowContext
)
from Task.task_atm_export_files import (
    TaskExportAtmFiles, TaskExportAtmFilesContext
)
from Task.task_diagno_tput import (
    TaskDiagnoTput, TaskDiagnoTputContext
)
from Task.task_diagno_latency import (
    TaskDiagnoLatency, TaskDiagnoLatencyContext
)
from Task.task_enter_xmodem import (
    TaskEnterXmodem, TaskEnterXmodemContext
)
from error.atmosic_error import AtmosicError
from error.errorcodes import SUCCESS_EXIT_CODE, RETRY_SUCCESS
from logger import start_new_session, start_retry_session, LogCompressorType,\
    AsyncFileHandler
from app_context import AppContext

__version__ = "25.09.01.0"
logger = getLogger(__name__)

CLI_COMMAND_FIELD = "cmd_task_name"
DEFAULT_LOG_COMPRESS = LogCompressorType.ZSTD
DEFAULT_MAX_RETRY_TIMES = 1
CONFIG_FILE = "config.toml"

class App:
    """Main application class with task-based architecture"""

    def __init__(self):
        self.sub_parser_helper = ArgumentSubParserHelper()
        self.args : Namespace | None = None
        self.log_compress = LogCompressorType.ZSTD
        self.is_del_log_after_compress = True
        self.subcommand : str = ""
        self.serial : str = ""
        self.context : Any | None = None
        self.task_type : Type[Any] | None = None
        self.test_time = 0
        self.max_retry_times = DEFAULT_MAX_RETRY_TIMES
        self.is_return_retry_error_code = False
        self.argument_defaults = {}

    def load_default_args(self) -> None:
        """Load default arguments from TOML config file"""
        config_path = Path(CONFIG_FILE)
        
        if not config_path.exists():
            return
        
        with open(config_path, "rb") as f:
            config = tomllib.load(f)
        
        self.argument_defaults =  config.get("launch_arguments", {})
        print(f"Default arguments: {self.argument_defaults}")

    def setup_argument(self):
        """Setup command line arguments to App object"""
        assert self.args is not None
        if hasattr(self.args, "log_compress"):
            self.log_compress = LogCompressorType[
                self.args.log_compress.upper()]
        if hasattr(self.args, "is_del_log_after_compress"):
            self.is_del_log_after_compress = self.args.is_del_log_after_compress
        if hasattr(self.args, CLI_COMMAND_FIELD):
            self.subcommand = getattr(self.args, CLI_COMMAND_FIELD)
        else:
            raise ValueError("Missing command line field"
                             f" `{CLI_COMMAND_FIELD}`")
        if hasattr(self.args, "serial"):
            self.serial = self.args.serial
        if hasattr(self.args, "max_full_process_retry_times"):
            self.max_retry_times = self.args.max_full_process_retry_times
        if hasattr(self.args, "return_retry_error_code"):
            self.is_return_retry_error_code = self.args.return_retry_error_code


    def configure_logger(self):
        AsyncFileHandler.set_compressor_type(self.log_compress)
        AsyncFileHandler.set_del_log_after_compress(
            self.is_del_log_after_compress)

    def start_logging_session(self):
        log_section_name = self.subcommand
        if self.serial:
            log_section_name += f"_{self.serial}"
        else:
            try:
                log_section_name += f"_PID{getpid()}"
            except BaseException:
                pass
        start_new_session(log_section_name)

    def generate_task_type_context(self) -> None:
        assert self.args is not None
        assert self.sub_parser_helper is not None
        self.sub_parser_helper.set_parsed_args(self.args)
        self.context = self.sub_parser_helper.get_subcommand_context()
        self.task_type = self.sub_parser_helper.get_subcommand_task_type()

    def run_task(self) -> None:
        assert self.args is not None
        assert self.context is not None
        assert self.task_type is not None
        self.test_times = 0
        while self.test_times < self.max_retry_times + 1:
            self.test_times += 1
            if self.test_times > 1:
                start_retry_session(f"retry_{self.test_times}")
            try:
                app_ctx = AppContext()
                app_ctx.test_times = self.test_times
                app_ctx.max_test_times = self.max_retry_times + 1
                task = self.task_type(self.context, app_ctx)
                task.run()
                return
            except AtmosicError as e:
                if self.test_times != self.max_retry_times + 1:
                    logger.warning(f"Task failed: {e}, then retry")
                else:
                    raise e

    def run(self) -> int:
        """Main execution function"""
        exit_code = SUCCESS_EXIT_CODE
        try:
            self.parse_args()
            assert self.args is not None
            self.setup_argument()
            self.configure_logger()
            self.start_logging_session()
            logger.debug(f"Args: {self.args}")
            self.generate_task_type_context()
            self.run_task()
        except AtmosicError as e:
            logger.error(f"Task failed: {e}")
            exit_code = e.error_code
        finally:
            PollingThreadBase.is_need_all_polling_thread_stopped = True
        if self.is_return_retry_error_code and self.test_times > 1:
            exit_code = RETRY_SUCCESS + exit_code
            logger.info(f"Repath error code: {exit_code}")
        return exit_code

    def parse_args(self) -> None:
        """Parse command line arguments"""
        self.load_default_args()
        is_engineer_mode = False
        if "--engineer-mode" in sys.argv:
            is_engineer_mode = True
            sys.argv.remove("--engineer-mode")

        basic_parser = ArgumentParser(add_help=False)
        
        basic_parser.add_argument(
            "--version", action="version",
            version=f"%(prog)s {__version__}"
        )
        basic_parser.add_argument(
            "--log-debug", "-v",
            action="store_true",
            help="Enable debug logging"
        )
        basic_parser.add_argument(
            "--log-compress",
            help="Compress log files",
            default="zstd",
            choices=["zstd", "7z", "none"],
        )
        basic_parser.add_argument(
            "--no-del-log",
            help="Do not delete log file after compress"\
                if is_engineer_mode else SUPPRESS,
            dest="is_del_log_after_compress",
            action="store_false",
        )
        basic_parser.add_argument(
            "--raw-debug-serial",
            help="Serial port for HCI raw data output"\
                if is_engineer_mode else SUPPRESS,
            default="",
        )
        basic_parser.add_argument(
            "--raw-debug-baudrate",
            help="Baudrate for HCI raw data output"\
                if is_engineer_mode else SUPPRESS,
            default=115200,
            type=int,
        )
        basic_parser.add_argument(
            "--raw-debug-tcp-port",
            help="TCP port for HCI raw data output"\
                if is_engineer_mode else SUPPRESS,
            default=0,
            type=int,
        )
        basic_parser.add_argument(
            "--max-full-process-retry-times",
            help="Max retry times for full process"\
                if is_engineer_mode else SUPPRESS,
            default=self.argument_defaults.get(
                "max-full-process-retry-times", DEFAULT_MAX_RETRY_TIMES),
            type=int,
        )
        basic_parser.add_argument(
            "--return-retry-error-code",
            help="Error code returned by the retry group"\
                if is_engineer_mode else SUPPRESS,
            action="store_true",
        )
        self.parser = ArgumentParser(
            description="UART loader host application",
            parents=[basic_parser],
        )
        self.sub_parsers = self.parser.add_subparsers(
            title='task',
            dest=CLI_COMMAND_FIELD,
            help='Task to execute',
        )
        self.sub_parser_helper.set_sub_parser(self.sub_parsers)
        self.sub_parser_helper.add_subcommand_class(
            "prog_bin",
            TaskProgramBinFileContext,
            TaskProgramBinFile,
            help="",
            parents=[basic_parser],
            is_engineer_mode=is_engineer_mode,
        )
        self.sub_parser_helper.add_subcommand_class(
            "prog_atm",
            TaskProgramAtmFileContext,
            TaskProgramAtmFile,
            help="",
            parents=[basic_parser],
            is_engineer_mode=is_engineer_mode,
        )
        self.sub_parser_helper.add_subcommand_class(
            "dump_bin",
            TaskDumpBinFileContext,
            TaskDumpBinFile,
            help="",
            parents=[basic_parser],
            is_engineer_mode=is_engineer_mode,
        )
        self.sub_parser_helper.add_subcommand_class(
            "atm_append_flash_bin",
            TaskAtmAppendFlashBinContext,
            TaskAtmAppendFlashBin,
            help="",
            parents=[basic_parser],
            is_engineer_mode=is_engineer_mode,
        )
        self.sub_parser_helper.add_subcommand_class(
            "atm_append_flash_erase",
            TaskAtmAppendFlashEraseContext,
            TaskAtmAppendFlashErase,
            help="",
            parents=[basic_parser],
            is_engineer_mode=is_engineer_mode,
        )
        self.sub_parser_helper.add_subcommand_class(
            "atm_append_rram_bin",
            TaskAtmAppendRramBinContext,
            TaskAtmAppendRramBin,
            help="",
            parents=[basic_parser],
            is_engineer_mode=is_engineer_mode,
        )
        self.sub_parser_helper.add_subcommand_class(
            "atm_append_rram_erase",
            TaskAtmAppendRramEraseContext,
            TaskAtmAppendRramErase,
            help="",
            parents=[basic_parser],
            is_engineer_mode=is_engineer_mode,
        )
        self.sub_parser_helper.add_subcommand_class(
            "atm_show",
            TaskAtmShowContext,
            TaskAtmShow,
            help="",
            parents=[basic_parser],
            is_engineer_mode=is_engineer_mode,
        )
        self.sub_parser_helper.add_subcommand_class(
            "atm_export",
            TaskExportAtmFilesContext,
            TaskExportAtmFiles,
            help="",
            parents=[basic_parser],
            is_engineer_mode=is_engineer_mode,
        )
        self.sub_parser_helper.add_subcommand_class(
            "diagno_tput",
            TaskDiagnoTputContext,
            TaskDiagnoTput,
            help="",
            parents=[basic_parser],
            is_engineer_mode=is_engineer_mode,
        )
        self.sub_parser_helper.add_subcommand_class(
            "diagno_latency",
            TaskDiagnoLatencyContext,
            TaskDiagnoLatency,
            help="",
            parents=[basic_parser],
            is_engineer_mode=is_engineer_mode,
        )
        self.sub_parser_helper.add_subcommand_class(
            "xmodem_sync",
            TaskEnterXmodemContext,
            TaskEnterXmodem,
            help="",
            parents=[basic_parser],
            is_engineer_mode=is_engineer_mode,
        )
        self.args = self.parser.parse_args()