import time
from threading import Thread

from llm_server import opts
from llm_server.database.database import weighted_average_column_for_model
from llm_server.llm.info import get_running_model
from llm_server.routes.cache import redis


def main_background_thread():
    redis.set('average_generation_elapsed_sec', 0)
    redis.set('estimated_avg_tps', 0)
    redis.set('average_output_tokens', 0)
    redis.set('backend_online', 0)
    redis.set_dict('backend_info', {})

    while True:
        # TODO: unify this
        if opts.mode == 'oobabooga':
            running_model, err = get_running_model()
            if err:
                print(err)
                redis.set('backend_online', 0)
            else:
                redis.set('running_model', running_model)
                redis.set('backend_online', 1)
        elif opts.mode == 'vllm':
            running_model, err = get_running_model()
            if err:
                print(err)
                redis.set('backend_online', 0)
            else:
                redis.set('running_model', running_model)
                redis.set('backend_online', 1)
        else:
            raise Exception

        # exclude_zeros=True filters out rows where an error message was returned. Previously, if there was an error, 0
        # was entered into the column. The new code enters null instead but we need to be backwards compatible for now.
        average_generation_elapsed_sec = weighted_average_column_for_model('prompts', 'generation_time', running_model, opts.mode, opts.backend_url, exclude_zeros=True, include_system_tokens=opts.include_system_tokens_in_stats) or 0
        if average_generation_elapsed_sec:  # returns None on exception
            redis.set('average_generation_elapsed_sec', average_generation_elapsed_sec)

        # overall = average_column_for_model('prompts', 'generation_time', opts.running_model)
        # print(f'Weighted: {average_generation_elapsed_sec}, overall: {overall}')

        average_output_tokens = weighted_average_column_for_model('prompts', 'response_tokens', running_model, opts.mode, opts.backend_url, exclude_zeros=True, include_system_tokens=opts.include_system_tokens_in_stats) or 0
        if average_generation_elapsed_sec:
            redis.set('average_output_tokens', average_output_tokens)

        # overall = average_column_for_model('prompts', 'response_tokens', opts.running_model)
        # print(f'Weighted: {average_output_tokens}, overall: {overall}')

        estimated_avg_tps = round(average_output_tokens / average_generation_elapsed_sec, 2) if average_generation_elapsed_sec > 0 else 0  # Avoid division by zero
        redis.set('estimated_avg_tps', estimated_avg_tps)
        time.sleep(60)