"""
This module defines the compiler interface. Look at the Builder class.
"""

from io import IOBase
from threading import Lock
from typing import Any, Dict, Iterator

from ...logger import logger
from .event import Event
from .graph import Graph
from .types.execution.entities import Call, Process
from .utils import chrono

__all__ = ["Builder", "build", "BuildingPhase"]





class BuildingPhase(Event):
    
    """
    This class of event is made to notify the system that we are entering a new buidling phase.
    The major attribute is the name of the phase.
    The minor attribute indicates the iteration number of this phase.
    """

    __lock = Lock()
    __extra_phases = 0

    __slots__ = {
        "major" : "The name of the phase of which the execution is starting.",
        "minor" : "The iteration number of this phase."
    }

    def __init__(self, major : str, minor : int) -> None:
        self.major = major
        self.minor = minor
    
    @staticmethod
    def request_finalizing_phase() -> int:
        """
        Requests a new finalizing phase to be performed by the Builder.
        Returns the integer assigned to the minor of the scheduled finalizer phase.
        """
        with BuildingPhase.__lock:
            n = BuildingPhase.__extra_phases
            BuildingPhase.__extra_phases += 1
            return n
    
    @staticmethod
    def finalizing_steps() -> int:
        """
        Returns the number of finalizing steps tp be performed.
        """
        with BuildingPhase.__lock:
            return BuildingPhase.__extra_phases





class Builder:

    """
    This class handles the building of BAGUETTEs. Just give the source file to the constructor and call build().
    """

    def __init__(self, data : str | IOBase | dict) -> None:
        from io import IOBase
        from ...logger import logger
        if isinstance(data, str):
            try:
                data = open(data, "r")
            except:
                raise 
        if isinstance(data, IOBase):
            from json import JSONDecodeError, load
            try:
                data = load(data)
                logger.info("File loaded and structured.")
            except JSONDecodeError:
                raise
        if not isinstance(data, dict):
            raise TypeError("Expected dict or json file, got " + repr(data.__class__.__name__))
        self.data = data
        self._progress = 0
        self._target = 1
    
    @property
    def progress(self) -> float:
        """
        Returns the progress (between 0.0 and 1.0) of the current task.
        """
        return self._progress / self._target

    def calls(self) -> Iterator[Dict[str, Any]]:
        """
        Yields all the calls' data from the source file.
        """
        for process in self.data["behavior"]["processes"]:
            for call in process["calls"]:
                yield call

    def machines(self) -> Iterator[Dict[str, str]]:
        """
        Yields all the machines' data from the source file.
        """
        if "domains" not in self.data["network"]:
            return
        for machine in self.data["network"]["domains"]:
            yield machine

    def hostIP(self) -> str:
        """
        Returns the host machine's IP address from the source file.
        """
        import re
        expr = re.compile(r"\[cuckoo.core.resultserver\] DEBUG: Now tracking machine (\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}) for task #\d+\n")
        for line in self.data["debug"]["cuckoo"]:
            if expr.search(line):
                match = expr.search(line)
                if match:
                    return match.group(1)
        logger.warning("Could not find the host's IP address.")
        return ""
    
    def sample_file_path(self) -> str:
        """
        Finds the name of the sample file in the report and returns its absolute path in the execution environment.
        """
        from ..source.types.execution import Process
        from ..source.utils import parse_command_line
        fname = self.data["target"]["file"]["name"]
        pid = None
        for p in self.data["behavior"]["processtree"]:      # Look only at the processes that Cuckoo started itself.
            for arg in parse_command_line(p["command_line"]):
                if arg.endswith(fname):
                    pid = p["pid"]
        if pid != None:
            for p in Process:
                if p.PID == pid:
                    return p.executable
        raise RuntimeError("Could not find malware process")
        

    @chrono
    def build(self) -> None:
        """
        Builds the BAGUETTE from the source file.
        """
        BuildingPhase("Initialization", 0).throw()
        from .graph import Graph, find_or_create
        from .types.execution import (Call, FollowedBy, NextSignificantCall,
                                      Process, Thread)
        from .types.execution.utils import CallHandler
        from .types.filesystem.integration import declare_existing_file
        from .types.network import Host, SpawnedProcess
        from ...logger import logger
        self.graph = Graph()

        with self.graph:

            BuildingPhase("Network Discovery", 0).throw()
            logger.info("Identifying machines.")

            for machine in self.machines():
                h = Host()
                h.address = machine["ip"]
                h.domain = machine["domain"]
            
            logger.info("Identifying host machine.")
            
            self.host = find_or_create(Host, address = self.hostIP())[0]
            Host.current = self.host
            self.host.domain = "host"
            self.host.platform = self.data["info"]["platform"]
            self.graph.data["platform"] = self.data["info"]["platform"]

            BuildingPhase("Input Parsing", 0).throw()
            logger.debug("Discovering work on calls.")

            self._target = 0
            for process in self.data["behavior"]["processes"]:
                self._target += len(process["calls"])
            
            N_calls = self._target
            logger.info("{} Calls to find.".format(self._target))

            # Building the basic execution graph
            BuildingPhase("Graph Building", 0).throw()
            logger.info("Building execution star.")
            for process in self.data["behavior"]["processes"]:
                self.graph.append(self.process_to_vertex(process), False)

            # Creating a file node for the sample file.
            logger.debug("Creating sample file.")
            try:
                path = self.sample_file_path()
                declare_existing_file(path)
            except RuntimeError:
                logger.error("Could not find the process executing the sample.")

            # Discovering calls
            logger.info("Discovering {} calls.".format(N_calls))
            calls : list[Call] = []
            self._target += len(Call) * 2
            seen = set()
            for t in Thread:
                a = t.first
                while a != None:
                    if a in seen:
                        logger.warning("A call has been seen twice while listing.")
                        break
                    seen.add(a)
                    calls.append(a)
                    _target = None
                    for l in a.edges:
                        if isinstance(l, FollowedBy) and l.destination is not a:
                            _target = l.destination
                            break
                    a = _target
            if len(calls) < N_calls:
                logger.warning("{} calls have been forgotten!".format(N_calls - len(calls)))
            
            # Sorting calls based on time
            logger.info("Sorting calls based on time.")
            calls.sort(key = lambda a : a.time)

            # Interpreting each Call
            BuildingPhase("Call Interpretation", 0).throw()
            logger.info("Integrating {} calls.".format(N_calls))
            for a in calls:
                CallHandler.integrate_chain(a)
                self._progress += 1
            
            # Linking processes to host
            BuildingPhase("Process Attribution", 0).throw()
            logger.info("Linking root processes to host machine.")
            p : Process
            for p in Process:
                if not p.parent_process:
                    SpawnedProcess(self.host, p)
            
            # Making skip links
            BuildingPhase("Call Skip-Linking", 0).throw()
            logger.info("Adding skip links in system call sequences")
            for t in Thread:
                last = t.first
                new = last
                while new and last:
                    if new is not last:
                        for l in new.edges:
                            if not isinstance(l, FollowedBy):
                                NextSignificantCall(last, new)
                                last = new
                                break
                    next_call = None
                    for l in new.edges:
                        if isinstance(l, FollowedBy) and l.destination is not new:
                            next_call = l.destination
                            break
                    new = next_call
            
            # Extra building phases
            for i in range(BuildingPhase.finalizing_steps()):
                logger.info("Running finalization phase #{}.".format(i + 1))
                BuildingPhase("Finalizer", i).throw()


    @chrono
    def process_to_vertex(self, process : dict) -> Process:
        """
        Given the source data of a process, creates the corresponding Process node.
        """
        from ..source.graph import find_or_create
        from ..source.types.execution import (Call, FollowedBy,
                                              HasChildProcess, HasFirstCall,
                                              HasThread, Process, Thread)
        from ..source.types.imports import HasImport, Import
        p : Process
        t : Thread
        a : Call
        p, _ = find_or_create(Process, PID = process["pid"])
        p.executable = process["process_path"]
        p.command = process["command_line"]
        p.start = process["first_seen"]
        p2, _ = find_or_create(Process, PID = process["ppid"])
        e = HasChildProcess(p2, p)

        for imp in process["modules"]:
            basename : str = imp["basename"]
            length : int = imp["imgsize"]
            path : str = imp["filepath"]
            if basename.lower().endswith(".dll"):
                imp, _ = find_or_create(Import, name = basename.lower()[:-4], path = path, length = length)
                HasImport(p, imp)

        for i, ci in enumerate(process["calls"]):
            t, existed = find_or_create(Thread, TID = ci["tid"])
            if not existed:
                HasThread(p, t)
            a = self.api_to_vertex(ci)
            a.thread = t
            if t.first == None:
                t.first = a
                t.last = a
                t.Ncalls = 1
                HasFirstCall(t, a)
            elif t.last:
                FollowedBy(t.last, a)
                t.last = a
                t.Ncalls += 1
            self._progress += 1

        for t in p.threads:
            if t.first and t.last:
                t.start = t.first.time
                t.stop = t.last.time
        
        p.start = min((t.start for t in p.threads), default=p.start)
        p.stop = max((t.stop for t in p.threads), default=p.start)

        return p

    @chrono
    def api_to_vertex(self, call : dict) -> Call:
        """
        Given the data from the source file, builds the corresponding Call node.
        """
        from ..source.record import record
        from ..source.types.execution import Call
        a = Call()
        a.name = call["api"]
        a.category = call["category"]
        a.stacktrace = tuple(call["stacktrace"])
        a.status = call["status"]
        a.arguments = record(**call["arguments"])
        a.return_value = call["return_value"]
        a.time = call["time"]
        a.flags = record(**call["flags"])
        return a





def build(file : str | IOBase) -> "Graph":
    """
    Shortcut function to directly build a BAGUETTE from a source file. Returns the freshly-baked BAGUETTE.
    """
    if isinstance(file, str):
        file = open(file, "rb")
    try:
        b = Builder(file)
    except:
        raise TypeError("Expected file, got " + repr(file.__class__.__name__))
    b.build()
    return b.graph





del IOBase, Lock, Any, Dict, Iterator, logger, Event, Graph, Call, Process, chrono