import os
import time
import functools


class time_exec:
    def __init__(self, func):
        functools.update_wrapper(self, func)
        self.func = func

    def __call__(self, *args, **kwargs):
        start_time = time.time()
        result = self.func(*args, **kwargs)
        end_time = time.time()
        execution_time = end_time - start_time
        timing_log = (
            f"Function '{self.func.__name__}' executed in {execution_time:.2f} seconds"
        )
        print(timing_log)
        return result, execution_time


class Logger:
    def __init__(self, args):
        self.log_file = args.log_file
        self.ds_domain = args.domain
        self.seed = args.seed
        self.datastore_size = args.sample_size
        self.stride = args.stride
        self.max_seq_length = args.max_seq_length
        self.merge = args.merge
        self.prefix = f"{self.ds_domain}\t{self.seed}\t{self.datastore_size}\t{self.stride}\t{self.max_seq_length}\t{self.merge}"

    def log_results(
        self,
        time_sample=None,
        time_chunk=None,
        time_index=None,
        time_eval=None,
        num_eval=None,
        perplexity=None,
    ):
        # Create the log entry
        log_entry = f"{self.prefix}\t{time_sample}\t{time_chunk}\t{time_index}\t{time_eval}\t{num_eval}\t{perplexity}\n"

        # Open the file in append mode. Creates the file if it doesn't exist.
        with open(self.log_file, "a") as file:
            file.write(log_entry)

    def log_string(self, log_string):
        with open(self.log_file, "a") as file:
            file.write(log_string)