from __future__ import annotations

import gc
import os
import time
from datetime import datetime, timedelta

import pytest
from inference.core.managers.inference_models_cache_watchdog import (
    RECENT_THRESHOLD_DAYS, SECONDS_IN_DAY, STALE_THRESHOLD_DAYS,
    WARM_THRESHOLD_DAYS, FileInfo, StalenessGroup, rank_for_deletion)

def codeflash_wrap(
    codeflash_wrapped: Callable[..., Any],
    codeflash_test_module_name: str,
    codeflash_test_class_name: Optional[str],
    codeflash_test_name: str,
    codeflash_function_name: str,
    codeflash_line_id: str,
    codeflash_loop_index: int,
    *args: Any,
    **kwargs: Any,
) -> str:
    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}"
    exception = None
    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}######$!")
    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 _instantiate_fileinfo(cls, *, modified_at, size_mb, extra_fields=None):
    """
    Attempt to create a real instance of the FileInfo class exported by the module.
    We try to call the constructor with common field names; if that fails we
    fall back to creating an instance without calling __init__ and then set
    attributes directly on the instance. This still produces a real instance
    of the imported FileInfo class (satisfying the requirement to use real classes).
    """
    if extra_fields is None:
        extra_fields = {}
    candidate_kwargs = {'modified_at': modified_at, 'size_mb': size_mb, 'mtime': modified_at, 'modifiedAt': modified_at, 'timestamp': modified_at, 'size': size_mb, 'size_mb': size_mb, 'name': extra_fields.get('name', 'f'), 'path': extra_fields.get('path', 'f'), 'relative_path': extra_fields.get('relative_path', 'f'), 'filename': extra_fields.get('filename', 'f')}
    try:
        sig = inspect.signature(cls)
        allowed = set(sig.parameters)
        filtered = {k: v for k, v in candidate_kwargs.items() if k in allowed}
        try:
            return cls(**filtered)
        except TypeError:
            params = list(sig.parameters)
            pos_args = []
            for p in params:
                if p in ('self',):
                    continue
                if p in filtered:
                    pos_args.append(filtered[p])
                elif p in candidate_kwargs:
                    pos_args.append(candidate_kwargs[p])
                else:
                    pos_args = None
                    break
            if pos_args is not None:
                return cls(*pos_args)
    except (ValueError, TypeError):
        pass
    inst = object.__new__(cls)
    setattr(inst, 'modified_at', modified_at)
    setattr(inst, 'size_mb', size_mb)
    for k, v in (('mtime', modified_at), ('size', size_mb)):
        try:
            setattr(inst, k, v)
        except Exception:
            pass
    return inst

def _expected_group_value(file_info, now, recent_days, warm_days, stale_days):
    """
    Replicates the staleness group logic from the function under test using the
    imported StalenessGroup enum values. This allows tests to compute the
    expected ordering deterministically without relying on internal function state.
    """
    age_days = (now - file_info.modified_at).total_seconds() / SECONDS_IN_DAY
    if age_days > stale_days:
        return StalenessGroup.ABANDONED.value
    if age_days > warm_days:
        return StalenessGroup.STALE.value
    if age_days > recent_days:
        return StalenessGroup.WARM.value
    return StalenessGroup.RECENT.value

def test_empty_input_returns_empty_list():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    codeflash_return_value = codeflash_wrap(rank_for_deletion, 'tests.inference.unit_tests.test_rank_for_deletion__unit_test_0', None, 'test_empty_input_returns_empty_list', 'rank_for_deletion', '0', codeflash_loop_index, [])
    codeflash_output = codeflash_return_value

def test_basic_sorting_by_group_then_size():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    now = datetime.utcnow()
    f_recent_big = _instantiate_fileinfo(FileInfo, modified_at=now - timedelta(hours=1), size_mb=50.0)
    f_recent_small = _instantiate_fileinfo(FileInfo, modified_at=now - timedelta(hours=2), size_mb=10.0)
    f_warm = _instantiate_fileinfo(FileInfo, modified_at=now - timedelta(days=2), size_mb=100.0)
    f_stale = _instantiate_fileinfo(FileInfo, modified_at=now - timedelta(days=10), size_mb=200.0)
    f_abandoned = _instantiate_fileinfo(FileInfo, modified_at=now - timedelta(days=40), size_mb=5.0)
    files = [f_warm, f_recent_small, f_abandoned, f_stale, f_recent_big]
    codeflash_return_value = codeflash_wrap(rank_for_deletion, 'tests.inference.unit_tests.test_rank_for_deletion__unit_test_0', None, 'test_basic_sorting_by_group_then_size', 'rank_for_deletion', '7', codeflash_loop_index, files, now=now)
    codeflash_output = codeflash_return_value
    ranked = codeflash_output
    expected = sorted(files, key=lambda f: (_expected_group_value(f, now, RECENT_THRESHOLD_DAYS, WARM_THRESHOLD_DAYS, STALE_THRESHOLD_DAYS), -f.size_mb))

def test_thresholds_must_be_ascending_raises_value_error():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    now = datetime.utcnow()
    file = _instantiate_fileinfo(FileInfo, modified_at=now - timedelta(days=1), size_mb=1.0)
    with pytest.raises(ValueError):
        codeflash_return_value = codeflash_wrap(rank_for_deletion, 'tests.inference.unit_tests.test_rank_for_deletion__unit_test_0', None, 'test_thresholds_must_be_ascending_raises_value_error', 'rank_for_deletion', '2_0', codeflash_loop_index, [file], now=now, recent_threshold_days=10, warm_threshold_days=5, stale_threshold_days=20)
    with pytest.raises(ValueError):
        codeflash_return_value = codeflash_wrap(rank_for_deletion, 'tests.inference.unit_tests.test_rank_for_deletion__unit_test_0', None, 'test_thresholds_must_be_ascending_raises_value_error', 'rank_for_deletion', '3_0', codeflash_loop_index, [file], now=now, recent_threshold_days=1, warm_threshold_days=1, stale_threshold_days=30)

def test_boundary_ages_are_handled_exclusively():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    now = datetime.utcnow()
    at_recent = _instantiate_fileinfo(FileInfo, modified_at=now - timedelta(days=RECENT_THRESHOLD_DAYS), size_mb=1.0)
    just_over_recent = _instantiate_fileinfo(FileInfo, modified_at=now - timedelta(days=RECENT_THRESHOLD_DAYS) - timedelta(seconds=1), size_mb=1.0)
    codeflash_return_value = codeflash_wrap(rank_for_deletion, 'tests.inference.unit_tests.test_rank_for_deletion__unit_test_0', None, 'test_boundary_ages_are_handled_exclusively', 'rank_for_deletion', '3', codeflash_loop_index, [just_over_recent, at_recent], now=now)
    codeflash_output = codeflash_return_value
    ranked = codeflash_output
    expected = sorted([just_over_recent, at_recent], key=lambda f: (_expected_group_value(f, now, RECENT_THRESHOLD_DAYS, WARM_THRESHOLD_DAYS, STALE_THRESHOLD_DAYS), -f.size_mb))

def test_invalid_modified_at_types_raise_exception():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    now = datetime.utcnow()
    broken_file = _instantiate_fileinfo(FileInfo, modified_at='not-a-datetime', size_mb=1.0)
    with pytest.raises(Exception):
        codeflash_return_value = codeflash_wrap(rank_for_deletion, 'tests.inference.unit_tests.test_rank_for_deletion__unit_test_0', None, 'test_invalid_modified_at_types_raise_exception', 'rank_for_deletion', '2_0', codeflash_loop_index, [broken_file], now=now)

def test_large_scale_sorting_deterministic_and_efficient():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    now = datetime.utcnow()
    files = []
    n = 1000
    for i in range(n):
        cycle = i % 4
        if cycle == 0:
            delta = timedelta(hours=1 + i % 24)
        elif cycle == 1:
            delta = timedelta(days=RECENT_THRESHOLD_DAYS + 0.5 + i % 3)
        elif cycle == 2:
            delta = timedelta(days=WARM_THRESHOLD_DAYS + 1 + i % 5)
        else:
            delta = timedelta(days=STALE_THRESHOLD_DAYS + 5 + i % 10)
        size = float(i * 37 % 1000) + 0.5
        fi = _instantiate_fileinfo(FileInfo, modified_at=now - delta, size_mb=size, extra_fields={'name': f'f{i}'})
        files.append(fi)
    codeflash_return_value = codeflash_wrap(rank_for_deletion, 'tests.inference.unit_tests.test_rank_for_deletion__unit_test_0', None, 'test_large_scale_sorting_deterministic_and_efficient', 'rank_for_deletion', '4', codeflash_loop_index, list(files), now=now)
    codeflash_output = codeflash_return_value
    ranked = codeflash_output
    expected = sorted(files, key=lambda f: (_expected_group_value(f, now, RECENT_THRESHOLD_DAYS, WARM_THRESHOLD_DAYS, STALE_THRESHOLD_DAYS), -f.size_mb))
    first_group = _expected_group_value(ranked[0], now, RECENT_THRESHOLD_DAYS, WARM_THRESHOLD_DAYS, STALE_THRESHOLD_DAYS)
    last_group = _expected_group_value(ranked[-1], now, RECENT_THRESHOLD_DAYS, WARM_THRESHOLD_DAYS, STALE_THRESHOLD_DAYS)