# Copyright (c) 2025, UChicago Argonne, LLC
# BSD OPEN SOURCE LICENSE. Full license can be found in LICENSE
# Copyright (c) 2025, UChicago Argonne, LLC
# BSD OPEN SOURCE LICENSE. Full license can be found in LICENSE
from os.path import dirname, join
from datetime import datetime
from time import sleep
from pathlib import Path
import shlex
import logging

from qgis.PyQt import uic
from qgis.PyQt.QtWidgets import QDialog
from qgis.PyQt.QtCore import pyqtSignal, QProcess, QThread, QObject
from qgis.core import (
    Qgis,
    QgsBlockingProcess,
    QgsRunProcess,
    QgsApplication,
    QgsVectorFileWriter,
    QgsProcessingFeedback,
    QgsProcessingUtils,
    QgsMessageLog,
    QgsSettings,
    QgsCredentials,
    QgsDataSourceUri,
    QgsProjUtils,
    QgsCoordinateReferenceSystem,
    QgsProcessingException,
)

from polaris.utils.env_utils import is_windows
from polaris.network.utils.worker_thread import WorkerThread
from polaris.runs.convergence.convergence_runner import run_polaris_convergence
from polaris.runs.convergence.convergence_callback_functions import default_end_of_loop_fn, default_start_of_loop_fn

from .is_test import running_on_ci


FORM_CLASS, _ = uic.loadUiType(join(dirname(__file__), "forms/model_run.ui"))


class ModelRunner(QDialog, FORM_CLASS):
    def __init__(self, _PQgis):
        QDialog.__init__(self)
        self.iface = _PQgis.iface
        self.setupUi(self)
        self._PQgis = _PQgis

        self.progress_frame.setVisible(False)
        # self.execute()
        self.but_run.clicked.connect(self.execute)

    def execute(self):
        self.progressBar1.setMaximum(len(self._PQgis.polaris_project.run_config.iterations()))
        self.progressBar2.setMaximum(86400)
        self.progress_frame.setVisible(True)
        self.setup_frame.setVisible(False)
        self.lbl_iter_name.setText("Loading data. Please wait.")

        self.thread = QThread()
        self.worker = Updater(self._PQgis)

        self.worker.moveToThread(self.thread)

        self.thread.started.connect(self.worker.run)
        self.worker.update.connect(self.signal_handler)
        self.thread.start()

    def signal_handler(self, val):
        if val[0] == "iter_finished":
            self.progressBar1.setValue(val[1])
        elif val[0] == "iter_start":
            print("ITER START", val)
            self.lbl_iter_name.setText(val[1])
        elif val[0] == "time_update":
            self.progressBar2.setValue(val[1])
            self.lbl_tod.setText(val[2])
        elif val[0] == "log_update":
            self.lbl_log.setText(val[1])
        elif val[0] == "done":
            print("done")
            self.close()


class Updater(QObject):
    update = pyqtSignal(list)

    def __init__(self, _PQgis):
        QObject.__init__(self)
        self._PQgis = _PQgis
        self.iteration_folders = []
        self.iter_count = 0
        self.run_config = (
            self._PQgis.polaris_project.run_config
        )  # type: polaris.runs.convergence.convergence_config.ConvergenceConfig
        self.run_config.num_retries = 0

    def run(self):
        print("running")
        self.run_config.user_data = self
        run_polaris_convergence(
            self.run_config,
            start_of_loop_fn=self.start_iteration,
            end_of_loop_fn=self.finish_iteration,
            cmd_runner=self.cmd_runner,
        )

        self.update.emit(["done"])
        self.finished.emit()

    def start_iteration(self, config, current_iteration, mods, scenario_file):
        self.iter_count += 1
        self.update.emit(["iter_start", str(current_iteration)])
        self.update.emit(["time_update", 0, "Starting the day", "."])
        default_start_of_loop_fn(config, current_iteration, mods, scenario_file)

    def finish_iteration(self, config, current_iteration, output_dir, polaris_inputs):
        self.update.emit(["iter_finished", self.iter_count])
        self.update.emit(["time_update", 86400, "Finished day", "."])
        default_end_of_loop_fn(config, current_iteration, output_dir, polaris_inputs)

    def cmd_runner(self, cmd, working_dir, ignore_errors=False, stderr_buf=None):
        print("Running command: ", cmd)

        def on_stdout(ba):
            all_val = ba.data().decode("UTF-8")
            for val in all_val.split("\n"):
                print("OUT", val)
                line = val.split("[NOTICE]")
                if len(line) > 1 and "departed" in line[1] and "arrived" in line[1]:
                    time_str = line[1].strip().split(",")[0].strip()
                    self.update.emit(["time_update", self.time_to_seconds(time_str), time_str])
                line = val.strip()
                if len(line) > 50:
                    self.update.emit(["log_update", line])

        def on_stderr(ba):
            val = ba.data().decode("UTF-8")
            line = val.split("[NOTICE]")
            if len(line) > 1 and "departed" in line[1] and "arrived" in line[1]:
                time_str = line[1].strip().split(",")[0].strip()
                self.update.emit(["time_update", self.time_to_seconds(time_str), time_str])
            else:
                print("NOT VALID_LINE", line)
            line = val.split("]")
            if len(line) > 1:
                self.update.emit(["log_update", line[1]])

        if isinstance(cmd, str):
            cmd = shlex.split(cmd, posix=not is_windows())

        # On windows we get problems with file not found when running an exe with backslashes in
        cmd = [str(x).replace("\\", "/") for x in cmd]

        command, *arguments = cmd
        logging.info(command, arguments)
        proc = QgsBlockingProcess(command, arguments)
        proc.setStdOutHandler(on_stdout)
        proc.setStdErrHandler(on_stdout)

        feedback = QgsProcessingFeedback()
        feedback.pushInfo("Polaris Running command:")
        feedback.pushCommandInfo(" ".join(cmd))
        feedback.pushInfo("Polaris command output")
        res = proc.run(feedback)

        if res == 0:
            logging.info("Model ran successfully")
        elif feedback.isCanceled() and res != 0:
            logging.info("Model run was canceled and did not complete")
        else:
            logging.critical(f"Non-zero exit code ({ proc.exitStatus()}) returned for cmd: {''.join(cmd)}")
            if not feedback.isCanceled() and proc.exitStatus() == QProcess.ExitStatus.CrashExit and not ignore_errors:
                raise QgsProcessingException("Process was unexpectedly terminated")
            elif proc.processError() == QProcess.ProcessError.FailedToStart and not ignore_errors:
                raise QgsProcessingException(
                    f"Process {command} failed to start. Either {command} is missing, or you may have insufficient permissions to run the program."
                )
            else:
                feedback.reportError("Process returned error code {}".format(res))
        return proc.exitStatus()

    @staticmethod
    def time_to_seconds(time_str):
        time_format = datetime.strptime(time_str, "%H:%M:%S")
        return time_format.hour * 3600 + time_format.minute * 60 + time_format.second
