from __future__ import annotations

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

import pytest
from inference.core import roboflow_api as rf
from inference.core.cache import cache
from inference.core.env import WORKFLOWS_DEFINITION_CACHE_EXPIRY
from inference.core.exceptions import MalformedWorkflowResponseError
from inference.core.roboflow_api import (_prepare_workflow_response_cache_key,
                                         get_workflow_specification)


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 _write_local_workflow_file(tmp_path: Path, workflow_id: str, content: dict):
    """
    Helper to write the expected local workflow file for workspace_id == "local".
    It mirrors the file naming logic from get_workflow_specification.
    """
    workflow_hash = hashlib.sha256(workflow_id.encode()).hexdigest()
    local_file_dir = Path(tmp_path) / 'workflow' / 'local'
    local_file_dir.mkdir(parents=True, exist_ok=True)
    local_file_path = local_file_dir / f'{workflow_hash}.json'
    local_file_path.write_text(json.dumps(content), encoding='utf-8')
    return local_file_path

def test_local_invalid_workflow_id_raises_value_error(monkeypatch, tmp_path):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    monkeypatch.setattr(rf, 'MODEL_CACHE_DIR', str(tmp_path))
    with pytest.raises(ValueError):
        codeflash_return_value = codeflash_wrap(rf.get_workflow_specification, 'tests.inference.unit_tests.test_get_workflow_specification__unit_test_0', None, 'test_local_invalid_workflow_id_raises_value_error', 'get_workflow_specification', '1_0', codeflash_loop_index, api_key=None, workspace_id='local', workflow_id='invalid id')

def test_local_missing_file_raises_file_not_found(monkeypatch, tmp_path):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    monkeypatch.setattr(rf, 'MODEL_CACHE_DIR', str(tmp_path))
    workflow_id = 'valid-id-123'
    workflow_hash = hashlib.sha256(workflow_id.encode()).hexdigest()
    expected_path = Path(tmp_path) / 'workflow' / 'local' / f'{workflow_hash}.json'
    if expected_path.exists():
        expected_path.unlink()
    with pytest.raises(FileNotFoundError):
        codeflash_return_value = codeflash_wrap(rf.get_workflow_specification, 'tests.inference.unit_tests.test_get_workflow_specification__unit_test_0', None, 'test_local_missing_file_raises_file_not_found', 'get_workflow_specification', '5_0', codeflash_loop_index, api_key=None, workspace_id='local', workflow_id=workflow_id)

def test_local_malformed_response_missing_config_raises_malformed(monkeypatch, tmp_path):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    monkeypatch.setattr(rf, 'MODEL_CACHE_DIR', str(tmp_path))
    workflow_id = 'wf_missing_config'
    _write_local_workflow_file(tmp_path, workflow_id, {'id': 'local123', 'not_config': True})
    with pytest.raises(MalformedWorkflowResponseError):
        codeflash_return_value = codeflash_wrap(rf.get_workflow_specification, 'tests.inference.unit_tests.test_get_workflow_specification__unit_test_0', None, 'test_local_malformed_response_missing_config_raises_malformed', 'get_workflow_specification', '3_0', codeflash_loop_index, api_key=None, workspace_id='local', workflow_id=workflow_id)

def test_local_invalid_config_json_raises_malformed(monkeypatch, tmp_path):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    monkeypatch.setattr(rf, 'MODEL_CACHE_DIR', str(tmp_path))
    workflow_id = 'wf_invalid_json'
    _write_local_workflow_file(tmp_path, workflow_id, {'config': 'this is not valid json'})
    with pytest.raises(MalformedWorkflowResponseError):
        codeflash_return_value = codeflash_wrap(rf.get_workflow_specification, 'tests.inference.unit_tests.test_get_workflow_specification__unit_test_0', None, 'test_local_invalid_config_json_raises_malformed', 'get_workflow_specification', '3_0', codeflash_loop_index, api_key=None, workspace_id='local', workflow_id=workflow_id)

def test_local_successful_specification_includes_id_and_caches_ephemeral(monkeypatch, tmp_path):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    monkeypatch.setattr(rf, 'MODEL_CACHE_DIR', str(tmp_path))
    workflow_id = 'wf_success'
    inner_spec = {'specification': {'nodes': ['a', 'b'], 'value': 42}}
    local_content = {'id': 'local-identifier', 'config': json.dumps(inner_spec)}
    _write_local_workflow_file(tmp_path, workflow_id, local_content)
    codeflash_return_value = codeflash_wrap(rf.get_workflow_specification, 'tests.inference.unit_tests.test_get_workflow_specification__unit_test_0', None, 'test_local_successful_specification_includes_id_and_caches_ephemeral', 'get_workflow_specification', '5', codeflash_loop_index, api_key=None, workspace_id='local', workflow_id=workflow_id, use_cache=True, ephemeral_cache=cache)
    codeflash_output = codeflash_return_value
    result = codeflash_output
    cache_key = rf._prepare_workflow_response_cache_key(api_key=None, workspace_id='local', workflow_id=workflow_id, workflow_version_id=None)
    cached = cache.get(key=cache_key)

def test_ephemeral_cache_short_circuits_network_path(monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    workspace_id = 'workspace_cache_test'
    workflow_id = 'workflow_cache_1'
    api_key = 'APIKEY123'
    prepop_spec = {'cached': True, 'value': 7}
    cache_key = rf._prepare_workflow_response_cache_key(api_key=api_key, workspace_id=workspace_id, workflow_id=workflow_id)
    cache.set(key=cache_key, value=prepop_spec, expire=rf.WORKFLOWS_DEFINITION_CACHE_EXPIRY)
    codeflash_return_value = codeflash_wrap(rf.get_workflow_specification, 'tests.inference.unit_tests.test_get_workflow_specification__unit_test_0', None, 'test_ephemeral_cache_short_circuits_network_path', 'get_workflow_specification', '6', codeflash_loop_index, api_key=api_key, workspace_id=workspace_id, workflow_id=workflow_id, use_cache=True, ephemeral_cache=cache)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_large_scale_repeated_cached_calls_are_deterministic_and_fast(monkeypatch, tmp_path):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    monkeypatch.setattr(rf, 'MODEL_CACHE_DIR', str(tmp_path))
    workflow_id = 'big_workflow'
    inner_spec = {'specification': {'counter': 123}}
    local_content = {'id': 'big-id', 'config': json.dumps(inner_spec)}
    _write_local_workflow_file(tmp_path, workflow_id, local_content)
    codeflash_return_value = codeflash_wrap(rf.get_workflow_specification, 'tests.inference.unit_tests.test_get_workflow_specification__unit_test_0', None, 'test_large_scale_repeated_cached_calls_are_deterministic_and_fast', 'get_workflow_specification', '5', codeflash_loop_index, api_key=None, workspace_id='local', workflow_id=workflow_id, use_cache=True, ephemeral_cache=cache)
    codeflash_output = codeflash_return_value
    first = codeflash_output
    for _ in range(1000):
        codeflash_return_value = codeflash_wrap(rf.get_workflow_specification, 'tests.inference.unit_tests.test_get_workflow_specification__unit_test_0', None, 'test_large_scale_repeated_cached_calls_are_deterministic_and_fast', 'get_workflow_specification', '7_0', codeflash_loop_index, api_key=None, workspace_id='local', workflow_id=workflow_id, use_cache=True, ephemeral_cache=cache)
        codeflash_output = codeflash_return_value
        res = codeflash_output