from __future__ import annotations

import gc
import os
import shutil
import time
from typing import List

import pytest
from inference.core.cache import air_gapped
from inference.core.cache.air_gapped import (get_cached_foundation_models,
                                             has_cached_model_variant)


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
get_cached_foundation_models = air_gapped.get_cached_foundation_models
has_cached_model_variant = air_gapped.has_cached_model_variant

class SimpleBlock:

    def __init__(self, manifest_class, identifier: str):
        self.manifest_class = manifest_class
        self.identifier = identifier
        self.block_source = 'test_source'

def make_manifest_class(supported_variants: List[str]=None, schema: dict=None, schema_raises: bool=False):
    """
    Create a minimal manifest class object with the required classmethods:
    - get_supported_model_variants: returns supported_variants (or None)
    - model_json_schema: returns schema or raises if schema_raises True

    We return an actual class (constructed at runtime) so that attribute access
    and isinstance checks behave like real classes.
    """

    class Manifest:

        @classmethod
        def get_supported_model_variants(cls):
            return supported_variants

        @classmethod
        def model_json_schema(cls):
            if schema_raises:
                raise RuntimeError('schema failed')
            return schema or {}
    return Manifest

def test_basic_traditional_cache_detected_and_schema_used(tmp_path, monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    monkeypatch.setattr(air_gapped, 'MODEL_CACHE_DIR', str(tmp_path))
    monkeypatch.setattr(air_gapped, 'USE_INFERENCE_MODELS', False)
    model_id = 'foundation/model-a'
    model_dir = tmp_path / model_id
    model_dir.mkdir(parents=True)
    (model_dir / 'weights.bin').write_bytes(b'\x00')
    schema = {'name': 'FancyModelBlock', 'properties': {'type': {'const': 'foundation.block'}}}
    Manifest = make_manifest_class(supported_variants=[model_id], schema=schema)
    block = SimpleBlock(manifest_class=Manifest, identifier='fallback.identifier')
    codeflash_return_value = codeflash_wrap(get_cached_foundation_models, 'tests.inference.unit_tests.test_get_cached_foundation_models__unit_test_0', None, 'test_basic_traditional_cache_detected_and_schema_used', 'get_cached_foundation_models', '9', codeflash_loop_index, blocks=[block])
    codeflash_output = codeflash_return_value
    results = codeflash_output
    entry = results[0]

def test_schema_errors_and_type_fallback_to_identifier(tmp_path, monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    monkeypatch.setattr(air_gapped, 'MODEL_CACHE_DIR', str(tmp_path))
    monkeypatch.setattr(air_gapped, 'USE_INFERENCE_MODELS', False)
    model_id = 'model-b'
    (tmp_path / model_id).mkdir(parents=True)
    (tmp_path / model_id / 'file').write_text('data')
    Manifest = make_manifest_class(supported_variants=[model_id], schema_raises=True)
    block = SimpleBlock(manifest_class=Manifest, identifier='explicit.identifier')
    codeflash_return_value = codeflash_wrap(get_cached_foundation_models, 'tests.inference.unit_tests.test_get_cached_foundation_models__unit_test_0', None, 'test_schema_errors_and_type_fallback_to_identifier', 'get_cached_foundation_models', '7', codeflash_loop_index, blocks=[block])
    codeflash_output = codeflash_return_value
    results = codeflash_output
    entry = results[0]

def test_block_with_no_supported_variants_is_skipped(tmp_path, monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    monkeypatch.setattr(air_gapped, 'MODEL_CACHE_DIR', str(tmp_path))
    monkeypatch.setattr(air_gapped, 'USE_INFERENCE_MODELS', False)
    Manifest = make_manifest_class(supported_variants=None, schema={'name': 'X'})
    block = SimpleBlock(manifest_class=Manifest, identifier='id.x')
    codeflash_return_value = codeflash_wrap(get_cached_foundation_models, 'tests.inference.unit_tests.test_get_cached_foundation_models__unit_test_0', None, 'test_block_with_no_supported_variants_is_skipped', 'get_cached_foundation_models', '4', codeflash_loop_index, blocks=[block])
    codeflash_output = codeflash_return_value
    results = codeflash_output

def test_empty_blocks_list_returns_empty(tmp_path, monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    monkeypatch.setattr(air_gapped, 'MODEL_CACHE_DIR', str(tmp_path))
    monkeypatch.setattr(air_gapped, 'USE_INFERENCE_MODELS', False)
    codeflash_return_value = codeflash_wrap(get_cached_foundation_models, 'tests.inference.unit_tests.test_get_cached_foundation_models__unit_test_0', None, 'test_empty_blocks_list_returns_empty', 'get_cached_foundation_models', '2', codeflash_loop_index, blocks=[])
    codeflash_output = codeflash_return_value
    results = codeflash_output

def test_load_blocks_failure_returns_empty_and_is_silent(tmp_path, monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    monkeypatch.setattr(air_gapped, 'MODEL_CACHE_DIR', str(tmp_path))
    monkeypatch.setattr(air_gapped, 'USE_INFERENCE_MODELS', False)

    def _bad_loader():
        raise RuntimeError("can't load blocks")
    monkeypatch.setattr(air_gapped, '_load_blocks', _bad_loader)
    codeflash_return_value = codeflash_wrap(get_cached_foundation_models, 'tests.inference.unit_tests.test_get_cached_foundation_models__unit_test_0', None, 'test_load_blocks_failure_returns_empty_and_is_silent', 'get_cached_foundation_models', '4', codeflash_loop_index, blocks=None)
    codeflash_output = codeflash_return_value
    results = codeflash_output

def test_inference_models_layout_detected_when_flag_enabled(tmp_path, monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    monkeypatch.setattr(air_gapped, 'MODEL_CACHE_DIR', str(tmp_path))
    monkeypatch.setattr(air_gapped, 'USE_INFERENCE_MODELS', True)
    model_id = 'simple-model'
    models_cache_dir = tmp_path / 'models-cache' / model_id
    models_cache_dir.mkdir(parents=True)
    (models_cache_dir / 'weights').write_text('x')
    Manifest = make_manifest_class(supported_variants=[model_id], schema={'name': 'S'})
    block = SimpleBlock(manifest_class=Manifest, identifier='id.s')
    codeflash_return_value = codeflash_wrap(get_cached_foundation_models, 'tests.inference.unit_tests.test_get_cached_foundation_models__unit_test_0', None, 'test_inference_models_layout_detected_when_flag_enabled', 'get_cached_foundation_models', '8', codeflash_loop_index, blocks=[block])
    codeflash_output = codeflash_return_value
    results = codeflash_output

def test_directory_with_only_hidden_files_is_not_considered_cached(tmp_path, monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    monkeypatch.setattr(air_gapped, 'MODEL_CACHE_DIR', str(tmp_path))
    monkeypatch.setattr(air_gapped, 'USE_INFERENCE_MODELS', False)
    model_id = 'hidden-only'
    model_dir = tmp_path / model_id
    model_dir.mkdir(parents=True)
    (model_dir / '.gitkeep').write_text('')
    (model_dir / '.hidden_folder').mkdir()
    Manifest = make_manifest_class(supported_variants=[model_id], schema={'name': 'H'})
    block = SimpleBlock(manifest_class=Manifest, identifier='id.h')
    codeflash_return_value = codeflash_wrap(get_cached_foundation_models, 'tests.inference.unit_tests.test_get_cached_foundation_models__unit_test_0', None, 'test_directory_with_only_hidden_files_is_not_considered_cached', 'get_cached_foundation_models', '9', codeflash_loop_index, blocks=[block])
    codeflash_output = codeflash_return_value
    results = codeflash_output

def test_has_cached_model_variant_handles_none_and_empty(tmp_path, monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    monkeypatch.setattr(air_gapped, 'MODEL_CACHE_DIR', str(tmp_path))
    monkeypatch.setattr(air_gapped, 'USE_INFERENCE_MODELS', False)

def test_large_scale_many_blocks_with_mixed_cache(tmp_path, monkeypatch):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Create 1000 block specs where every even-indexed model is cached and odd\n    ones are not. Confirm that the function returns the correct number of\n    cached entries and that their data is consistent.\n    '
    monkeypatch.setattr(air_gapped, 'MODEL_CACHE_DIR', str(tmp_path))
    monkeypatch.setattr(air_gapped, 'USE_INFERENCE_MODELS', False)
    num_blocks = 1000
    blocks = []
    cached_count = 0
    for i in range(num_blocks):
        model_id = f'bulk-model-{i}'
        schema = {'name': f'bulk-name-{i}', 'properties': {'type': {'const': f'bulk.type.{i}'}}}
        Manifest = make_manifest_class(supported_variants=[model_id], schema=schema)
        block = SimpleBlock(manifest_class=Manifest, identifier=f'bulk.identifier.{i}')
        blocks.append(block)
        if i % 2 == 0:
            model_dir = tmp_path / model_id
            model_dir.mkdir(parents=True)
            (model_dir / 'w').write_text('x')
            cached_count += 1
    codeflash_return_value = codeflash_wrap(get_cached_foundation_models, 'tests.inference.unit_tests.test_get_cached_foundation_models__unit_test_0', None, 'test_large_scale_many_blocks_with_mixed_cache', 'get_cached_foundation_models', '7', codeflash_loop_index, blocks=blocks)
    codeflash_output = codeflash_return_value
    results = codeflash_output
    seen_ids = {entry['model_id'] for entry in results}
    for i in range(num_blocks):
        model_id = f'bulk-model-{i}'
        if i % 2 == 0:
            pass
        else:
            pass
    if cached_count > 0:
        first = results[0]