from __future__ import annotations

import gc
import hashlib
import json
import os
import time
from pathlib import Path

import inference.core.env as env
import pytest
from inference.core.env import MODEL_CACHE_DIR, SINGLE_TENANT_WORKFLOW_CACHE
from inference.core.roboflow_api import (
    delete_cached_workflow_response_if_exists, get_workflow_cache_file,
    load_cached_workflow_response)


def codeflash_wrap(codeflash_wrapped, codeflash_test_module_name, codeflash_test_class_name, codeflash_test_name, codeflash_function_name, codeflash_line_id, codeflash_loop_index, *args, **kwargs):
    test_id = f'{codeflash_test_module_name}:{codeflash_test_class_name}:{codeflash_test_name}:{codeflash_line_id}:{codeflash_loop_index}'
    if not hasattr(codeflash_wrap, 'index'):
        codeflash_wrap.index = {}
    if test_id in codeflash_wrap.index:
        codeflash_wrap.index[test_id] += 1
    else:
        codeflash_wrap.index[test_id] = 0
    codeflash_test_index = codeflash_wrap.index[test_id]
    invocation_id = f'{codeflash_line_id}_{codeflash_test_index}'
    test_stdout_tag = f"{codeflash_test_module_name}:{(codeflash_test_class_name + '.' if codeflash_test_class_name else '')}{codeflash_test_name}:{codeflash_function_name}:{codeflash_loop_index}:{invocation_id}"
    print(f'!$######{test_stdout_tag}######$!')
    exception = None
    gc.disable()
    try:
        counter = time.perf_counter_ns()
        return_value = codeflash_wrapped(*args, **kwargs)
        codeflash_duration = time.perf_counter_ns() - counter
    except Exception as e:
        codeflash_duration = time.perf_counter_ns() - counter
        exception = e
    gc.enable()
    print(f'!######{test_stdout_tag}:{codeflash_duration}######!')
    if exception:
        raise exception
    return return_value

def test_load_cached_workflow_response_returns_dict_for_valid_json_single_tenant(tmp_path):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    original_cache_dir = getattr(env, 'MODEL_CACHE_DIR', None)
    original_single_tenant = getattr(env, 'SINGLE_TENANT_WORKFLOW_CACHE', None)
    env.MODEL_CACHE_DIR = str(tmp_path)
    env.SINGLE_TENANT_WORKFLOW_CACHE = True
    try:
        workspace_id = 'my-workspace'
        workflow_id = 'my_workflow-1'
        api_key = 'irrelevant_for_single_tenant'
        workflow_version_id = None
        cache_file = get_workflow_cache_file(workspace_id=workspace_id, workflow_id=workflow_id, api_key=api_key, workflow_version_id=workflow_version_id)
        Path(cache_file).parent.mkdir(parents=True, exist_ok=True)
        expected_payload = {'hello': 'world', 'value': 123}
        with open(cache_file, 'w', encoding='utf-8') as f:
            json.dump(expected_payload, f)
        codeflash_return_value = codeflash_wrap(load_cached_workflow_response, 'tests.inference.unit_tests.test_load_cached_workflow_response__unit_test_0', None, 'test_load_cached_workflow_response_returns_dict_for_valid_json_single_tenant', 'load_cached_workflow_response', '4_8', codeflash_loop_index, workspace_id=workspace_id, workflow_id=workflow_id, api_key=api_key, workflow_version_id=workflow_version_id)
        codeflash_output = codeflash_return_value
        result = codeflash_output
    finally:
        if original_cache_dir is None:
            delattr(env, 'MODEL_CACHE_DIR')
        else:
            env.MODEL_CACHE_DIR = original_cache_dir
        if original_single_tenant is None:
            delattr(env, 'SINGLE_TENANT_WORKFLOW_CACHE')
        else:
            env.SINGLE_TENANT_WORKFLOW_CACHE = original_single_tenant

def test_load_cached_workflow_response_returns_none_when_file_missing(tmp_path):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    original_cache_dir = getattr(env, 'MODEL_CACHE_DIR', None)
    original_single_tenant = getattr(env, 'SINGLE_TENANT_WORKFLOW_CACHE', None)
    env.MODEL_CACHE_DIR = str(tmp_path)
    env.SINGLE_TENANT_WORKFLOW_CACHE = True
    try:
        workspace_id = 'no_such_workspace'
        workflow_id = 'no_such_workflow'
        api_key = None
        workflow_version_id = None
        codeflash_return_value = codeflash_wrap(load_cached_workflow_response, 'tests.inference.unit_tests.test_load_cached_workflow_response__unit_test_0', None, 'test_load_cached_workflow_response_returns_none_when_file_missing', 'load_cached_workflow_response', '4_4', codeflash_loop_index, workspace_id=workspace_id, workflow_id=workflow_id, api_key=api_key, workflow_version_id=workflow_version_id)
        codeflash_output = codeflash_return_value
        result = codeflash_output
    finally:
        if original_cache_dir is None:
            delattr(env, 'MODEL_CACHE_DIR')
        else:
            env.MODEL_CACHE_DIR = original_cache_dir
        if original_single_tenant is None:
            delattr(env, 'SINGLE_TENANT_WORKFLOW_CACHE')
        else:
            env.SINGLE_TENANT_WORKFLOW_CACHE = original_single_tenant

def test_corrupted_json_file_is_deleted_and_returns_none(tmp_path):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    original_cache_dir = getattr(env, 'MODEL_CACHE_DIR', None)
    original_single_tenant = getattr(env, 'SINGLE_TENANT_WORKFLOW_CACHE', None)
    env.MODEL_CACHE_DIR = str(tmp_path)
    env.SINGLE_TENANT_WORKFLOW_CACHE = True
    try:
        workspace_id = 'workspaceX'
        workflow_id = 'corrupt_workflow'
        api_key = None
        cache_file = get_workflow_cache_file(workspace_id=workspace_id, workflow_id=workflow_id, api_key=api_key)
        Path(cache_file).parent.mkdir(parents=True, exist_ok=True)
        with open(cache_file, 'w', encoding='utf-8') as f:
            f.write('{ this is not : valid JSON,,, }')
        codeflash_return_value = codeflash_wrap(load_cached_workflow_response, 'tests.inference.unit_tests.test_load_cached_workflow_response__unit_test_0', None, 'test_corrupted_json_file_is_deleted_and_returns_none', 'load_cached_workflow_response', '4_6', codeflash_loop_index, workspace_id=workspace_id, workflow_id=workflow_id, api_key=api_key)
        codeflash_output = codeflash_return_value
        result = codeflash_output
    finally:
        if original_cache_dir is None:
            delattr(env, 'MODEL_CACHE_DIR')
        else:
            env.MODEL_CACHE_DIR = original_cache_dir
        if original_single_tenant is None:
            delattr(env, 'SINGLE_TENANT_WORKFLOW_CACHE')
        else:
            env.SINGLE_TENANT_WORKFLOW_CACHE = original_single_tenant

def test_non_single_tenant_cache_uses_api_key_in_filename_hash(tmp_path):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    original_cache_dir = getattr(env, 'MODEL_CACHE_DIR', None)
    original_single_tenant = getattr(env, 'SINGLE_TENANT_WORKFLOW_CACHE', None)
    env.MODEL_CACHE_DIR = str(tmp_path)
    env.SINGLE_TENANT_WORKFLOW_CACHE = False
    try:
        workspace_id = 'workspace-hash'
        workflow_id = 'wf-hash'
        api_key_a = 'KEY_A'
        api_key_b = 'KEY_B'
        cache_file_a = get_workflow_cache_file(workspace_id=workspace_id, workflow_id=workflow_id, api_key=api_key_a)
        cache_file_b = get_workflow_cache_file(workspace_id=workspace_id, workflow_id=workflow_id, api_key=api_key_b)
        Path(cache_file_a).parent.mkdir(parents=True, exist_ok=True)
        with open(cache_file_a, 'w', encoding='utf-8') as f:
            json.dump({'which': 'A'}, f)
        with open(cache_file_b, 'w', encoding='utf-8') as f:
            json.dump({'which': 'B'}, f)
        codeflash_return_value = codeflash_wrap(load_cached_workflow_response, 'tests.inference.unit_tests.test_load_cached_workflow_response__unit_test_0', None, 'test_non_single_tenant_cache_uses_api_key_in_filename_hash', 'load_cached_workflow_response', '4_9', codeflash_loop_index, workspace_id=workspace_id, workflow_id=workflow_id, api_key=api_key_a)
        codeflash_output = codeflash_return_value
        result_a = codeflash_output
        codeflash_return_value = codeflash_wrap(load_cached_workflow_response, 'tests.inference.unit_tests.test_load_cached_workflow_response__unit_test_0', None, 'test_non_single_tenant_cache_uses_api_key_in_filename_hash', 'load_cached_workflow_response', '4_12', codeflash_loop_index, workspace_id=workspace_id, workflow_id=workflow_id, api_key=api_key_b)
        codeflash_output = codeflash_return_value
        result_b = codeflash_output
        codeflash_return_value = codeflash_wrap(load_cached_workflow_response, 'tests.inference.unit_tests.test_load_cached_workflow_response__unit_test_0', None, 'test_non_single_tenant_cache_uses_api_key_in_filename_hash', 'load_cached_workflow_response', '4_15', codeflash_loop_index, workspace_id=workspace_id, workflow_id=workflow_id, api_key='NON_EXISTENT_KEY')
        codeflash_output = codeflash_return_value
        result_none = codeflash_output
    finally:
        if original_cache_dir is None:
            delattr(env, 'MODEL_CACHE_DIR')
        else:
            env.MODEL_CACHE_DIR = original_cache_dir
        if original_single_tenant is None:
            delattr(env, 'SINGLE_TENANT_WORKFLOW_CACHE')
        else:
            env.SINGLE_TENANT_WORKFLOW_CACHE = original_single_tenant

def test_sanitization_with_special_characters_in_ids(tmp_path):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    original_cache_dir = getattr(env, 'MODEL_CACHE_DIR', None)
    original_single_tenant = getattr(env, 'SINGLE_TENANT_WORKFLOW_CACHE', None)
    env.MODEL_CACHE_DIR = str(tmp_path)
    env.SINGLE_TENANT_WORKFLOW_CACHE = True
    try:
        workspace_id = '../../weird/workspace?name'
        workflow_id = 'workflow:with*special|chars<>'
        api_key = None
        cache_file = get_workflow_cache_file(workspace_id=workspace_id, workflow_id=workflow_id, api_key=api_key)
        Path(cache_file).parent.mkdir(parents=True, exist_ok=True)
        payload = {'ok': True}
        with open(cache_file, 'w', encoding='utf-8') as f:
            json.dump(payload, f)
        codeflash_return_value = codeflash_wrap(load_cached_workflow_response, 'tests.inference.unit_tests.test_load_cached_workflow_response__unit_test_0', None, 'test_sanitization_with_special_characters_in_ids', 'load_cached_workflow_response', '4_7', codeflash_loop_index, workspace_id=workspace_id, workflow_id=workflow_id, api_key=api_key)
        codeflash_output = codeflash_return_value
        result = codeflash_output
        prefix = os.path.abspath(os.path.join(env.MODEL_CACHE_DIR, 'workflow'))
    finally:
        if original_cache_dir is None:
            delattr(env, 'MODEL_CACHE_DIR')
        else:
            env.MODEL_CACHE_DIR = original_cache_dir
        if original_single_tenant is None:
            delattr(env, 'SINGLE_TENANT_WORKFLOW_CACHE')
        else:
            env.SINGLE_TENANT_WORKFLOW_CACHE = original_single_tenant

def test_many_workflows_load_in_loop_performance_and_correctness(tmp_path):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    original_cache_dir = getattr(env, 'MODEL_CACHE_DIR', None)
    original_single_tenant = getattr(env, 'SINGLE_TENANT_WORKFLOW_CACHE', None)
    env.MODEL_CACHE_DIR = str(tmp_path)
    env.SINGLE_TENANT_WORKFLOW_CACHE = True
    try:
        workspace_id = 'big_workspace'
        api_key = None
        n = 1000
        created_payloads = {}
        for i in range(n):
            workflow_id = f'wf_{i}'
            payload = {'index': i, 'data': f'value_{i}'}
            file_path = get_workflow_cache_file(workspace_id=workspace_id, workflow_id=workflow_id, api_key=api_key)
            Path(file_path).parent.mkdir(parents=True, exist_ok=True)
            with open(file_path, 'w', encoding='utf-8') as f:
                json.dump(payload, f)
            created_payloads[workflow_id] = payload
        for workflow_id, expected in created_payloads.items():
            codeflash_return_value = codeflash_wrap(load_cached_workflow_response, 'tests.inference.unit_tests.test_load_cached_workflow_response__unit_test_0', None, 'test_many_workflows_load_in_loop_performance_and_correctness', 'load_cached_workflow_response', '4_5_0', codeflash_loop_index, workspace_id=workspace_id, workflow_id=workflow_id, api_key=api_key)
            codeflash_output = codeflash_return_value
            loaded = codeflash_output
        for _ in range(10):
            for workflow_id, expected in created_payloads.items():
                codeflash_return_value = codeflash_wrap(load_cached_workflow_response, 'tests.inference.unit_tests.test_load_cached_workflow_response__unit_test_0', None, 'test_many_workflows_load_in_loop_performance_and_correctness', 'load_cached_workflow_response', '4_6_0_0', codeflash_loop_index, workspace_id=workspace_id, workflow_id=workflow_id, api_key=api_key)
                codeflash_output = codeflash_return_value
                loaded = codeflash_output
    finally:
        if original_cache_dir is None:
            delattr(env, 'MODEL_CACHE_DIR')
        else:
            env.MODEL_CACHE_DIR = original_cache_dir
        if original_single_tenant is None:
            delattr(env, 'SINGLE_TENANT_WORKFLOW_CACHE')
        else:
            env.SINGLE_TENANT_WORKFLOW_CACHE = original_single_tenant