#!/usr/bin/env python
# \file   Model.py
# \brief  High-level Python interface to Iris models.
#
# \date   Copyright ARM Limited 2020-2024 All Rights Reserved.

import iris.iris as iris
from .Target import AsyncTarget, SyncTarget
from .Exceptions import TargetBusyError, TimeoutError

from collections import namedtuple
from warnings import warn
import threading
import time
import os

TargetInfo = namedtuple("TargetInfo", ["instance_name", "target_name", "version",
                                       "description", "component_type",
                                       "executes_software", "is_hardware_model"])

class ModelBase(object):
    """
    This class wraps an Iris model.
    """

    def __init__(self, client, verbose, synchronous, socketpairpool=None):
        self.client = client
        self.verbose = verbose
        self.synchronous = synchronous
        self.socketpairpool = socketpairpool
        self._pid = None

        # Get instance infos including properties for all targets.
        all_instance_infos = self.client.irisCall().instanceRegistry_getList(getProperties = True)
        self.instance_infos = {inst_info.instName: inst_info for inst_info in all_instance_infos}

        # Targets indexed by instId. Each target will be created at most once.
        self.targets = {}

        # List of all CPU targets.
        self.cpus = []

    def __repr__(self):
        return "<%s(%s) object at %s>" % (
                self.__class__.__name__, str(self.client), hex(id(self)))

    def _new_target(self, target_info):
        # see SyncModel and AsyncModel for implementations
        raise NotImplementedError

    def get_target(self, instance_name):
        """
        Obtain an interface to a target

        :param instance_name: The instance name corresponding to the desired target
        """
        if instance_name not in self.instance_infos:
            raise ValueError("Model has no target named \"%s\"" % instance_name)

        instId = self.instance_infos[instance_name].instId
        if instId not in self.targets:
            self.targets[instId] = self._new_target(self.instance_infos[instance_name])
        return self.targets[instId]

    def get_cpus(self):
        """
        Returns all targets which have executesSoftware set or have
        componentType = 'Core'
        """
        if not self.cpus:
            for info in self.instance_infos.values():
                if (info.properties.get('componentType') == 'Core') or \
                            (info.properties.get('executesSoftware', 0) != 0):
                    self.cpus.append(self.get_target(info.instName))
            # Sort by instance id, which is compatible with fm.debug
            self.cpus.sort(key = lambda x: x.instId)

        return self.cpus

    def get_global_instance(self):
        """Returns global instance target.
        """
        if 0 not in self.targets:
            self.targets[0] = self._new_target([x for x in self.instance_infos.values() if x.instId == 0][0])
        return self.targets[0]

    def get_target_info(self):
        """Returns an iterator over namedtuples containing information about
           all of the target instances contained in the model"""
        for info in self.instance_infos.values():
            if not info.instName.startswith("component."):
                continue
            yield TargetInfo(info.instName,
                             info.properties.get('componentName'),  # target_name
                             info.properties.get('version'),
                             info.properties.get('description'),
                             info.properties.get('componentType'),
                             info.properties.get('executesSoftware', 0) != 0,
                             info.properties.get('isHardwareModel', 0) != 0)

    def get_targets(self, prefix = 'component'):
        """Generator function to iterate all targets in the simulation"""
        for info in self.instance_infos.values():
            if (not prefix) or info.instName.startswith(prefix + "."):
                yield self.get_target(info.instName)

    @property
    def is_running(self):
        try:
            return self.client.irisCall().simulationTime_get(instId=1).running
        except iris.error.IrisError:
            raise ValueError("Failed to get simulation time information")

    def _shutdown_model(self):
        """Request shutdown without disconnecting the TCP connection.
        """
        # Delete event streams in model before shutting down the model.
        for target in self.targets.values():
            if isinstance(target, SyncTarget):
                target.destroy_event_buffer()

        self.client.irisCall().simulation_requestShutdown(instId=1)

    def reset(self, allow_partial_reset=False):
        """
        Reset the simulation to exactly the same state it had after instantiation.
        Reason why we might need the argument allowPartialReset:
        We do not necessarily control all components in a SystemC platform and we are only able to
        reset Fast Models components. We can only fulfil the request if:
        1) This is an ISIM model. ISIM's are built purely from Fast Models components so we can reset
            the whole platform.
        2) allowPartialReset=true was set. This option indicates that the user knows that no all
           components will be reset and accepts the consequences.
        """

        self.client.irisCall().simulation_reset(instId=1, allowPartialReset=allow_partial_reset)

    def isConnected(self):
        """Returns whether the client is connected to the model"""
        return self.client.isConnected()

    def isFork(self):
        """This model was created by fork().
        """
        return self._pid is not None

    def release(self, shutdown=False):
        """
        End the simulation and release the model.

        :param shutdown:
            For TCP connections: If True, the simulation will be shutdown and any other scripts or debuggers
            will have to disconnect. If False, a simulation might be kept alive after disconnection.

            For UNIX domain socket connections: This should always be set to False as disconnecting from 
            the model will automatically shut down the model. Setting this to True will work but will 
            result in a non-deterministic shutdown sequence which is undesirable.
        """
        if self.isConnected():
            # If we are a fork, ask model explicitly to shut itself down via simulation_requestShutdown()
            # because fork clients do not close their sockets (which would automatically shut the model down)
            # because sockets are recycled for future forks. 
            if shutdown or self.isFork():
                self._shutdown_model()
            sock = self.client.sock
            self.client.disconnect(closeSocket = not self.isFork())
            if self.isFork():
                self.socketpairpool.releaseSocketpair(sock)

        if self.isFork():
            # If we are a fork, wait until our associated process exited.
            # The most reliable way is to do this via the parent client
            # and parent model, which implies that all child processes must
            # terminate before their parents. Except for this "waitpid()"
            # there is no other parent/child-dependency.
            exitStatus = self._parent_model.forkWait(self._pid)
            self._pid = None
            return exitStatus

    def fork(self, client_name = None):
        """Fork model and return a connected client.

        Use release() on the returned model to shut the model down and wait for the subprocess to exit.

        :param client_name
            Instance name of the returned Python client. This should begin with "client.".
            The default is the client name of the main client plus a suffix indicating the fork tree path using the socket file descriptors.
            Note that when having multiple forks, all client_names can be set to an identical name and the Iris instance name uniquification
            will not make them unique (different) because all client names live in independent model processes.
            Setting client_name to different values for different forks is useful for debugging/logging purposes but is not necessary otherwise.
        """
        if os.name == 'nt':
            raise NotImplementedError()
        lhs, rhs = self.socketpairpool.allocSocketpair()
        connectionSpec = "socketfd={}".format(rhs.fileno())
        try:
            simulationForkResult = self.client.irisCall().simulation_fork(instId=0, connect=connectionSpec)
        except Exception:
            raise
        try:
            if client_name == None:
                client_name = self.client.instanceName + "_f{}".format(lhs.fileno())
            model = NewUnixDomainSocketModel(lhs, timeoutInMs=self.client.timeoutInMs, client_name=client_name, synchronous=self.synchronous, verbose=self.verbose, waitForInstantiation=False, socketpairpool=self.socketpairpool)
            model._pid = simulationForkResult.childPid
            model._parent_model = self
            return model
        except Exception:
            raise

    def forkWait(self, pid):
        """Wait for child process to exit.

        Return exit status of child process.
        """
        simulationForkWaitResult = self.client.irisCall().simulation_forkWait(instId=0, pid=pid)
        return simulationForkWaitResult.exitStatus


class AsyncModel(ModelBase):
    """An Iris model optimized for asynchronous processing.

    The latency is reduced when running the model many instruction at a time,
    processing most events as the model is running.
    """
    def __init__(self, client, verbose, socketpairpool=None):
        ModelBase.__init__(self, client, verbose, synchronous=False, socketpairpool=socketpairpool)

        self.get_global_instance().add_event_callback("IRIS_SIMULATION_TIME_EVENT", self.ec_IRIS_SIMULATION_TIME_EVENT)
        self.get_global_instance().add_event_callback("IRIS_SIM_PHASE_END_OF_SIMULATION", self.ec_IRIS_SIM_PHASE_END_OF_SIMULATION)

        self._stop_event = threading.Event() # Triggered when simulation time stops or ends.
        self._run_event = threading.Event() # Triggered when simulation time starts.

    def _new_target(self, target_info):
        return AsyncTarget(target_info, self)

    def ec_IRIS_SIM_PHASE_END_OF_SIMULATION(self, time, fields):
          #Simulation is ending
          self._stop_event.set()

    # this method is called by the receive thread in iris library
    def ec_IRIS_SIMULATION_TIME_EVENT(self, esId, fields):
        if fields['RUNNING']:
            # Simulation time has started
            # Reset the state of all cpu breakpoints.
            for c in self.get_cpus():
                c.clear_bpts()

            # Notify anyone waiting on this event that the simulation has started running.
            self._run_event.set()
        else:
            # Simulation time has stopped
            self._stop_event.set()

    # Checks if the model is already running and clears _stop_event
    def _request_run(self):
        if self.is_running:
            raise TargetBusyError("The model is already running")
        self._stop_event.clear()
        self.client.irisCall().simulationTime_run(instId=1)

    # Checks if the model is already stopped and clears _stop_event
    def _request_stop(self):
        if not self.is_running:
            raise TargetBusyError("The model is already stopped")
        self._stop_event.clear()
        self.client.irisCall().simulationTime_stop(instId=1)

    def run(self, blocking = True, timeout = None):
        """
        Starts the model executing

        :param blocking:
            If True, this call will block until the model stops executing
            (typically due to a breakpoint).
            If False, this call will return once that target has started executing.

        :param timeout:
            If None, this call will wait indefinitely for the target to enter
            the correct state.
            If set to a float or int, this parameter gives the maximum number of
            seconds to wait.

        :raises TimeoutError: If the timeout expires.
                TargetBusyError: If the model is running.

        If blocking, returns a list of hit breakpoints that contributed to the
        stoppage (i.e. breakpoints configured as `dontStop = False` that were
        hit since the `run` invocation). Returns empty list if no stopping
        breakpoints were hit - which indicates end of simulation has been reached.

        If not blocking, returns `None`.
        """
        if not blocking:
            self._run_event.clear()

        self._request_run()

        if blocking:
            if not self._stop_event.wait(timeout):
                raise TimeoutError()
            self._stop_event.clear()
            return [bpt for target in self.targets.values()
                        for bpt in target.get_hit_breakpoints()
                        if hasattr(bpt.info, 'dontStop') and not bpt.info.dontStop]
        else:
            if not self._run_event.wait(timeout):
                raise TimeoutError()
            self._run_event.clear()
            return None

    def stop(self, timeout = None):
        """
        Stops the model executing.

        Silently return when the model is already stopped.

        :param timeout:
            If None, this call will wait indefinitely for the target to enter
            the correct state.
            If set to a float or int, this parameter gives the maximum number of
            seconds to wait.

        :raises TimeoutError: If the timeout expires.
        """

        if not self.is_running:
            return

        self._request_stop()

        # Wait for simulation time stop event after calling simulationTime_stop()
        if not self._stop_event.wait(timeout):
            raise TimeoutError()


class SyncModel(ModelBase):
    """An Iris model optimized for synchronous processing.

    The latency is reduced when stepping the model one or few instructions at a
    time.
    """
    def __init__(self, client, verbose, socketpairpool=None):
        ModelBase.__init__(self, client, verbose, synchronous=True, socketpairpool=socketpairpool)

    def _new_target(self, target_info):
        return SyncTarget(target_info, self)


class LibraryModel(AsyncModel):
    """A model that is not supported in Iris"""

    def __init__(self, filename, parameters=None, verbose=False):
        """Raises NotImplementedError in Iris"""
        raise NotImplementedError("LibraryModel is not supported by Iris")


def _NewModel(synchronous, client, verbose, socketpairpool=None):
    if synchronous:
        return SyncModel(client, verbose, socketpairpool)
    else:
        return AsyncModel(client, verbose, socketpairpool)


def _NewTCPClient(host, port, client_name, timeoutInMs, verbose, expectedServerPid = 0):
    client = iris.IrisTcpClient(client_name, iris.IrisInstance(), verbose, None, timeoutInMs)
    client.connect(host, port, timeoutInMs = timeoutInMs, expectedServerPid = expectedServerPid)
    # Ensure that the simulation has been instantiated and target components are available.
    try:
        client.irisCall().simulation_waitForInstantiation(instId=1)
    except iris.error.IrisError as e:
        if e.code == iris.error.E_function_not_supported_by_instance:
            # The simulation does not support simulation_waitForInstantiation(), carry on.
            warn("Simulation does not support simulation_waitForInstantiation: unable to guarantee instantiation")
        else:
            # A different error occurred. Propagate it so it doesn't get lost.
            raise e
    return client


def NewNetworkModel(host = "localhost", port = 0, timeoutInMs = 20000, client_name = "client.iris_debug", synchronous = False, verbose = False, expectedServerPid = 0):
    """Creates an Iris model connected to an Iris server

    When port is 0 this means to scan the port range 7100..7109 for Iris servers and connect to the first server.

    Connects to an already initialised Iris server

    :param host:
        Hostname or IP address of the computer running the model.

    :param port:
        Port number that the model is listening on.

    :param timeoutInMs:
        Timeout in milliseconds for operations that depend on other threads or processes.

    :param client_name:
        Name that is used to identify this client as an Iris instance. An integer suffix is added for disambiguation if needed.

    :param synchronous:
        If True, an Iris model optimized for synchronous processing is constructed.

    :param verbose:
        If True, extra debugging information is printed.

    :param expectedServerPid:
        If non-zero, the Iris server process rejects the connection unless its PID matches this value.
    """
    client = _NewTCPClient(host, port, client_name, timeoutInMs, verbose, expectedServerPid)
    return _NewModel(synchronous, client, verbose)


# NOTE: kept for backward compatibility reason. Prefer using NewNetworkModel.
class NetworkModel(AsyncModel):
    """An Iris model connected to an Iris server

    When port is 0 this means to scan the port range 7100..7109 for Iris servers and connect to the first server.
    """

    def __init__(self, host = "localhost", port = 0, timeoutInMs = 20000, client_name = "client.iris_debug", verbose = False, expectedServerPid = 0):
        """
        Connects to an already initialised Iris server

        :param host:
            Hostname or IP address of the computer running the model.

        :param port:
            Port number that the model is listening on.

        :param timeoutInMs:
            Timeout in milliseconds for operations that depend on other threads or processes.

        :param client_name:
            Name that is used to identify this client as an Iris instance. An integer suffix is added for disambiguation if needed.

        :param verbose:
            If True, extra debugging information is printed.

        :param expectedServerPid:
            If non-zero, the Iris server process rejects the connection unless its PID matches this value.
        """
        client = _NewTCPClient(host, port, client_name, timeoutInMs, verbose, expectedServerPid)
        AsyncModel.__init__(self, client, verbose)


def NewUnixDomainSocketModel(sock, timeoutInMs=20000, client_name="client.iris_debug", synchronous=False, verbose=False, socketpairpool=None, waitForInstantiation=True):
    """
    An Iris model connected to an Iris instance via unix domain socket
    """
    client = iris.IrisUnixDomainClient.ConnectToSocket(
            sock, client_name=client_name, timeout=timeoutInMs,
            verbose=1 if verbose else 0)
    try:
        if waitForInstantiation:
            client.irisCall().simulation_waitForInstantiation(instId=1)
    except iris.error.IrisError as e:
        if e.code == iris.error.E_function_not_supported_by_instance:
            warn("Simulation does not support simulation_waitForInstantiation: unable to guarantee instantiation")
        else:
            raise

    return _NewModel(synchronous, client, verbose, socketpairpool)
