import time
import os
import threading

class ProfileResult:
    def __init__(
        self,
        name: str,
        start,
        end,
        process_id: int = 0,
        thread_id: int = 0,
    ) -> None:
        self.name: str = name
        self.start = start
        self.end = end
        self.process_id = process_id
        self.thread_id = thread_id

class Instrumentor:

    def __init__(self) -> None:
        self.file = None
        self.profile_count: int = 0

    def begin_session(
        self,
        name: str,
        file_path: str = 'results.json'
    ) -> None:
        self.file = open(file_path, "w")
        self.write_header()

    def end_session(self) -> None:
        self.write_footer()
        self.file.close()
        self.session = None
        self.profile_count = 0

    def write_profile(
        self,
        result: ProfileResult,
    ) -> None:

        if self.profile_count > 0:
            self.file.write(',')
        self.profile_count += 1

        name: str = result.name
        name = name.replace('"', "'")

        self.file.write('{')
        self.file.write('"cat":"function",')
        self.file.write(f'"dur":{result.end - result.start},')
        self.file.write(f'"name":"{name}",')
        self.file.write('"ph":"X",')
        self.file.write(f'"pid":{result.process_id},')
        self.file.write(f'"tid":{result.thread_id},')
        self.file.write(f'"ts":{result.start}')
        self.file.write('}')

    def write_header(self) -> None:
        self.file.write('{"otherData": {},"traceEvents":[')

    def write_footer(self) -> None:
        self.file.write(']}')

class InstrumentationTimer:

    def __init__(
        self,
        name: str,
        process_id = None,
        thread_id = None,
    ) -> None:
        self.name = name

        self.process_id = os.getpid()
        if process_id is not None:
            self.process_id = process_id
        
        self.thread_id = threading.get_native_id()
        if thread_id is not None:
            self.thread_id = thread_id

        self.stopped: bool = False
        self.start = time.time_ns() / 1000

    def __del__(self) -> None:
        if not self.stopped:
            self.stop()

    def stop(self) -> None:
        end = time.time_ns() / 1000

        instance.write_profile(ProfileResult(self.name, self.start, end, self.process_id, self.thread_id))
        self.stopped = True

instance = Instrumentor()

def configure_profiling(
    save_location: str,
    file_name: str = 'results',
) -> None:
    instance.begin_session(os.path.join(save_location, file_name + '.json'))

def profile_decorator(func, pid = None, tid = None):
    def wrapper(*args, **kwargs):
        timer = InstrumentationTimer(func.__name__, pid, tid)
        ret = func(*args, **kwargs)
        timer.stop()
        return ret
    return wrapper

def end_profiling() -> None:
    instance.end_session()