"""
Command-line interface utilities
"""
import sys
import multiprocessing
import threading
import logging
import numbers
from typing import Any, Optional
from functools import lru_cache
from datetime import datetime
from pathlib import Path
import pickle


def print_nicqs_logo():
    """Print NICQS logo in light green color."""

    # ANSI color code for light green
    light_green = '\033[92m'
    reset = '\033[0m'

    logo = f"""{light_green}
 ███╗   ██╗██╗ ██████╗ ██████╗ ███████╗
 ████╗  ██║██║██╔════╝██╔═══██╗██╔════╝
 ██╔██╗ ██║██║██║     ██║   ██║███████╗
 ██║╚██╗██║██║██║     ██║▄▄ ██║╚════██║
 ██║ ╚████║██║╚██████╗╚██████╔╝███████║
 ╚═╝  ╚═══╝╚═╝ ╚═════╝ ╚══▀▀═╝ ╚══════╝
{reset}"""
    
    try:
        print(logo)
    except UnicodeEncodeError:
        # Fallback to ASCII version if Unicode characters can't be encoded
        ascii_logo = f"""{light_green}
 ███   █ █  ██████  ██████  ███████
 ████  █ █ █        █     █ █      
 █ ██  █ █ █        █     █ ███████
 █  ██ █ █ █        █  ▄▄ █       █
 █   ███ █  ██████   ██████  ███████
                       ▀▀          
{reset}"""
        try:
            print(ascii_logo)
        except UnicodeEncodeError:
            # If even the ASCII version fails, print a simple text version
            print(f"{light_green}NICQS - No-Insulation Coil Quench Simulator{reset}")


def var_dump(output_path, dump_data, name):
    with (output_path / f'{name}.pkl').open('wb') as f:
        pickle.dump(dump_data, f, protocol=pickle.HIGHEST_PROTOCOL)

class OmittedInfoFormatter(logging.Formatter):
    """Custom formatter to show log level for all messages except INFO."""
    def format(self, record):
        if record.levelno == logging.INFO:
            return record.getMessage()
        else:
            return super().format(record)


def start_logger(log: Path, verbose: bool = False):
    """
    Starts the logger.
    """
    print(f'Logging to {log}')
    logger = logging.getLogger('nicqs')
    logger.propagate = False

    logfile_handler = logging.FileHandler(log, mode="w", encoding='utf-8')
    logfile_handler.setFormatter(
        logging.Formatter("%(asctime)s.%(msecs)03d %(levelname)-8s %(message)s", datefmt='%Y-%m-%d %H:%M:%S')
    )
    logfile_handler.setLevel(logging.DEBUG)
    logger.addHandler(logfile_handler)

    console_handler = logging.StreamHandler()
    console_handler.setFormatter(OmittedInfoFormatter('%(levelname)-8s %(message)s'))
    console_handler.setLevel(logging.DEBUG if verbose else logging.INFO)
    logger.setLevel(logging.DEBUG)
    logger.addHandler(console_handler)


def start_solver_logger(log: Path):
    """
    Starts the solver logger. It logs all NICQS events, and is attached to the root `nicqs` logger.
    """
    logger = logging.getLogger('nicqs')
    solverlog_handler = logging.FileHandler(log, mode="a")
    solverlog_handler.setFormatter(
        logging.Formatter("%(asctime)s.%(msecs)03d %(levelname)-8s %(message)s", datefmt='%Y-%m-%d %H:%M:%S')
    )
    solverlog_handler.setLevel(logging.DEBUG)
    logger.addHandler(solverlog_handler)


class ProgressBar:
    """
    A progress bar context that is stationary and updates in-place.
    """

    DEFAULT_LENGTH = 16

    def __init__(self, steps, prefix, interval=None, length=DEFAULT_LENGTH, eta_delay=4, show_final_time=True):
        self.steps = steps
        self.prefix = prefix
        if interval:
            self.interval = interval
        elif steps < 1000:
            self.interval = 1
        elif steps < 10000:
            self.interval = 10
        elif steps < 100000:
            self.interval = 100
        elif steps < 1000000:
            self.interval = 1000
        else:
            self.interval = 5000
        self.length = length
        # Tracks the starting datetime to show an ETA 
        self.start_time = None
        # The number of seconds before we start to display the ETA
        self.eta_delay = eta_delay
        self.show_final_time = show_final_time

        self._previous_step = multiprocessing.Value('i', -1)

        self._step_lock = threading.RLock()
        self._display_lock = threading.Lock()


    def update(self, i):
        """
        Update the progress bar.
        Don't draw every update if an interval has been set. Always draw first and last update.

        `i` should not surpass `steps - 1`

        Thread safe.
        """
        # Comparing i+1 instead of i gives neat multiples in the output
        if ((i + 1) % self.interval) == 0 or i+1 == self.steps or i == 0:
            self.force_update(i)

        with self._step_lock:
            self._previous_step.value = i


    @lru_cache(maxsize=DEFAULT_LENGTH)
    def bar_str(self, stage):
        """ String representation of the loading bar. Cached because string operations are expensive in Python. """
        return '[' + '#' * stage + ' ' * (self.length - stage) + ']'


    def force_update(self, i):
        """
        Ignores specified interval and will force a progress bar redraw.
        This is useful if the user wants to update based on time instead of step count, for example.

        Thread safe
        """
        with self._step_lock:
            if i >= self.steps:
                print(f'Warning: progress bar went out of bounds: {i}>={self.steps}')
                i = self.steps - 1

            if self.start_time is None:
                self.start_time = datetime.now()

            stage = int(self.length * (i + 1) / self.steps)
            display_str = \
                f"\r{self.prefix} {self.bar_str(stage)} {i + 1}/{self.steps} {(i+1)*100 / self.steps:5.1f}%"
            # Check if we should display ETA
            runtime = (datetime.now() - self.start_time).total_seconds()
            if (i+1) == self.steps and self.show_final_time:
                # Final step in progress bar. Show final elapsed time
                display_str += f'  ({runtime:.2f}s)           '
            elif runtime > self.eta_delay:
                # Intermediate step, calculate ETA
                if i != 0:
                    eta = int((runtime / i) * (self.steps - i))
                    display_str += f' ({int(eta / 60):02}:{eta % 60:02} remaining)'
                else: 
                    display_str += '  (???s remaining)'

            # Spaces at the end to overwrite any potential shifts in output
            display_str += "      "
            if (i+1) == self.steps:
                # Final progress bar line
                display_str += "\n"
            sys.stdout.write(display_str)
            sys.stdout.flush()


    def increment(self):
        """
        Equivalent to calling update() with the previous value incremented by one. The first
        increment calls update(0). It is important to call imcrement as many times as the specified
        number of steps. i.e. for steps=2, increment() should be called twice.

        Safe to use with update calls.

        Thread safe.
        """
        with self._step_lock:
            nexti = self._previous_step.value + 1
            if ((nexti + 1) % self.interval) == 0 or nexti+1 == self.steps or nexti == 0:
                self.force_update(nexti)
            self._previous_step.value = nexti


    def __enter__(self):
        pass


    def __exit__(self, exc_type, exc_val, exc_tb):
        # Finish drawing the progress bar if used with a context manager
        self.force_update(self.steps - 1)


class Table:
    """ For printing tabular data to the CLI. """

    def __init__(self,
                 columns: list[str],
                 column_widths: list[int],
                 reprint_header_after: int = 40, *,
                 logger: Optional[logging.Logger] = None,
                 vertical_rules: Optional[list[int]] = None,
                 only_positive: Optional[list[bool]] = None):
        """
        Creates a new Table instance for printing tabular data to the CLI.
        If `logger` is specified, output is printed to the logger instead of stdout.
        Vertical rules can be added via the `vertical_rules` parameter, which adds a rule after the specified column.
        `only_positive` can contain a mask specifing which columns will only ever be positive, to save one character
        for the sign.
        """
        self.columns = columns
        self.column_widths = column_widths
        self.reprint_header_after = reprint_header_after
        self.logger = logger
        self.vrules = vertical_rules
        if only_positive is not None:
            self.only_positive = only_positive
        else:
            self.only_positive = [False for _ in self.columns]

        self.last_header = reprint_header_after
        """ Counts how many rows have been printed since the last header. """


        if len(self.columns) != len(self.column_widths) or len(self.columns) != len(self.only_positive):
            raise ValueError
        if len(self.columns) == 0:
            raise ValueError
        for rule in self.vrules:
            if rule < 0 or rule > (len(self.columns) - 1):
                raise ValueError


    @property
    def width(self) -> int:
        """ Width of table in characters. """
        #  | and a space for the leftmost column
        return 3 + sum(self.column_widths) + len(self.vrules) * 2


    def _print(self, s: str):
        """ Prints string to the correct output device. """
        if self.logger:
            self.logger.info(s)
        else:
            print(s)


    def print_rule(self):
        """ Prints a horizontal rule. """
        row = '+-'
        for i, w in enumerate(self.column_widths):
            row += '-' * w
            if i in self.vrules:
                row += '+-'
        row += '+'
        self._print(row)


    def print_header(self):
        """ Prints the table header. """
        self.print_rule()
        self.print_row(self.columns, no_header=True)
        self.print_rule()
        self.last_header = 0


    def print_row(self, data: list[Any] | dict[str, Any], no_header: bool = False):
        """ Prints a row to the table. Automatically reprints the header if necessary. """
        # Convert dict to list
        if isinstance(data, dict):
            data = [data[col] for col in self.columns]
        if len(data) != len(self.columns):
            raise ValueError(f'Expected list of length {len(self.columns)}, got {len(data)}')

        if self.last_header >= self.reprint_header_after and not no_header:
            self.print_header()

        row = '| '
        for i, datum in enumerate(data):
            col_width = self.column_widths[i]
            if isinstance(datum, numbers.Integral) and self.only_positive[i]:
                col = f'{datum:<{col_width - 1}d} '
            elif isinstance(datum, numbers.Integral):
                col = f'{datum:< {col_width - 1}d} '
            elif isinstance(datum, numbers.Real) and self.only_positive[i]:
                digits = col_width - 7
                if digits < 0:
                    raise ValueError('The column width must be at least 7 to print a float column.')
                col = f'{datum:<{col_width - 1}.{digits}e} '
            elif isinstance(datum, numbers.Real):
                digits = col_width - 8
                if digits < 0:
                    raise ValueError('The column width must be at least 8 to print a float column.')
                col = f'{datum:< {col_width - 1}.{digits}e} '
            else:
                col = f'{datum:<{col_width - 1}} '
            if len(col) > col_width:
                col = col[:col_width - 3] + '.. '
            row += col
            if i in self.vrules:
                row += '| '
        row += '|'
        self._print(row)
        self.last_header += 1


    def __enter__(self):
        pass


    def __exit__(self, exc_type, exc_val, exc_tb):
        # Finish drawing the table
        self.print_rule()


def geometry_main():
    """
    Entry point for nicqs-geometry command-line tool.
    Creates geometry files from YAML input for NI coil simulations.

    Usage:
        nicqs-geometry <input_yaml> [output_path]
    """
    # Import here to avoid circular dependencies and reduce startup time
    import sys
    from nicqs.geometry import NI_coil_geometry

    print_nicqs_logo()

    if len(sys.argv) > 2:
        print(f'Loading the provided yaml file: {sys.argv[1]}')
        print(f'Output path: {sys.argv[2]}')
        NICO = NI_coil_geometry(sys.argv[1], sys.argv[2])
    elif len(sys.argv) > 1:
        print(f'Loading the provided yaml file: {sys.argv[1]}')
        NICO = NI_coil_geometry(sys.argv[1], 'simulation_result')
    else:
        print("Usage: nicqs-geometry <input_yaml> [output_path]")
        print("\nCreates geometry files from YAML input for NI coil simulations.")
        print("\nArguments:")
        print("  input_yaml   Path to YAML configuration file")
        print("  output_path  Output directory (default: 'simulation_result')")
        sys.exit(1)


def solver_main():
    """
    Entry point for nicqs-solver command-line tool.
    Runs the NI coil thermo-electromagnetic solver.

    Usage:
        nicqs-solver <solver_input_yaml> [geometry_yaml]
    """
    # Import here to avoid circular dependencies and reduce startup time
    import sys
    from nicqs.solver_torch import NI_Coil_Solver

    print_nicqs_logo()

    if len(sys.argv) == 2:
        print(f'Loading solver configuration: {sys.argv[1]}')
        NI_sim = NI_Coil_Solver(sys.argv[1])
    elif len(sys.argv) == 3:
        print(f'Loading solver configuration: {sys.argv[1]}')
        print(f'Geometry file: {sys.argv[2]}')
        NI_sim = NI_Coil_Solver(sys.argv[1], sys.argv[2])
    else:
        print("Usage: nicqs-solver <solver_input_yaml> [geometry_yaml]")
        print("\nRuns the NI coil thermo-electromagnetic solver.")
        print("\nArguments:")
        print("  solver_input_yaml  Path to solver input YAML file")
        print("  geometry_yaml      Path to geometry YAML file (optional)")
        sys.exit(1)
