cpython/Lib/test/test_profiling/test_heatmap.py

"""Tests for the heatmap collector (profiling.sampling)."""

import os
import shutil
import tempfile
import unittest
from unittest import mock
from collections import namedtuple
from pathlib import Path

# Matches the C structseq LocationInfo from _remote_debugging
LocationInfo = namedtuple('LocationInfo', ['lineno', 'end_lineno', 'col_offset', 'end_col_offset'])

from profiling.sampling.heatmap_collector import (
    HeatmapCollector,
    get_python_path_info,
    extract_module_name,
)

from test.support import captured_stdout, captured_stderr


# =============================================================================
# Unit Tests for Public Helper Functions
# =============================================================================

class TestPathInfoFunctions(unittest.TestCase):
    """Test public helper functions for path information."""

    def test_get_python_path_info_returns_dict(self):
        """Test that get_python_path_info returns a dictionary with expected keys."""
        path_info = get_python_path_info()

        self.assertIsInstance(path_info, dict)
        self.assertIn('stdlib', path_info)
        self.assertIn('site_packages', path_info)
        self.assertIn('sys_path', path_info)

    def test_get_python_path_info_stdlib_is_path_or_none(self):
        """Test that stdlib is either a Path object or None."""
        path_info = get_python_path_info()

        if path_info['stdlib'] is not None:
            self.assertIsInstance(path_info['stdlib'], Path)

    def test_get_python_path_info_site_packages_is_list(self):
        """Test that site_packages is a list."""
        path_info = get_python_path_info()

        self.assertIsInstance(path_info['site_packages'], list)
        for item in path_info['site_packages']:
            self.assertIsInstance(item, Path)

    def test_get_python_path_info_sys_path_is_list(self):
        """Test that sys_path is a list of Path objects."""
        path_info = get_python_path_info()

        self.assertIsInstance(path_info['sys_path'], list)
        for item in path_info['sys_path']:
            self.assertIsInstance(item, Path)

    def test_extract_module_name_with_none(self):
        """Test extract_module_name with None filename."""
        path_info = get_python_path_info()
        module_name, module_type = extract_module_name(None, path_info)

        self.assertEqual(module_name, 'unknown')
        self.assertEqual(module_type, 'other')

    def test_extract_module_name_with_empty_string(self):
        """Test extract_module_name with empty filename."""
        path_info = get_python_path_info()
        module_name, module_type = extract_module_name('', path_info)

        self.assertEqual(module_name, 'unknown')
        self.assertEqual(module_type, 'other')

    def test_extract_module_name_with_stdlib_file(self):
        """Test extract_module_name with a standard library file."""
        path_info = get_python_path_info()

        # Use os module as a known stdlib file
        if path_info['stdlib']:
            stdlib_file = str(path_info['stdlib'] / 'os.py')
            module_name, module_type = extract_module_name(stdlib_file, path_info)

            self.assertEqual(module_type, 'stdlib')
            self.assertIn('os', module_name)

    def test_extract_module_name_with_project_file(self):
        """Test extract_module_name with a project file."""
        path_info = get_python_path_info()

        # Create a mock project file path
        if path_info['sys_path']:
            # Use current directory as project path
            project_file = '/some/project/path/mymodule.py'
            module_name, module_type = extract_module_name(project_file, path_info)

            # Should classify as 'other' if not in sys.path
            self.assertIn(module_type, ['project', 'other'])

    def test_extract_module_name_removes_py_extension(self):
        """Test that .py extension is removed from module names."""
        path_info = get_python_path_info()

        # Test with a simple .py file
        module_name, module_type = extract_module_name('/path/to/test.py', path_info)

        # Module name should not contain .py
        self.assertNotIn('.py', module_name)

    def test_extract_module_name_with_special_files(self):
        """Test extract_module_name with special filenames like <string>."""
        path_info = get_python_path_info()

        special_files = ['<string>', '<stdin>', '[eval]']
        for special_file in special_files:
            module_name, module_type = extract_module_name(special_file, path_info)
            self.assertEqual(module_type, 'other')


# =============================================================================
# Unit Tests for HeatmapCollector Public API
# =============================================================================

class TestHeatmapCollectorInit(unittest.TestCase):
    """Test HeatmapCollector initialization."""

    def test_init_creates_empty_data_structures(self):
        """Test that __init__ creates empty data structures."""
        collector = HeatmapCollector(sample_interval_usec=100)

        # Check that data structures are initialized
        self.assertIsInstance(collector.line_samples, dict)
        self.assertIsInstance(collector.file_samples, dict)
        self.assertIsInstance(collector.line_self_samples, dict)
        self.assertIsInstance(collector.file_self_samples, dict)
        self.assertIsInstance(collector.call_graph, dict)
        self.assertIsInstance(collector.callers_graph, dict)
        self.assertIsInstance(collector.function_definitions, dict)
        self.assertIsInstance(collector.edge_samples, dict)

        # Check that they're empty
        self.assertEqual(len(collector.line_samples), 0)
        self.assertEqual(len(collector.file_samples), 0)
        self.assertEqual(len(collector.line_self_samples), 0)
        self.assertEqual(len(collector.file_self_samples), 0)

    def test_init_sets_total_samples_to_zero(self):
        """Test that total samples starts at zero."""
        collector = HeatmapCollector(sample_interval_usec=100)
        self.assertEqual(collector._total_samples, 0)

    def test_init_gets_path_info(self):
        """Test that path info is retrieved during init."""
        collector = HeatmapCollector(sample_interval_usec=100)
        self.assertIsNotNone(collector._path_info)
        self.assertIn('stdlib', collector._path_info)


class TestHeatmapCollectorSetStats(unittest.TestCase):
    """Test HeatmapCollector.set_stats() method."""

    def test_set_stats_stores_all_parameters(self):
        """Test that set_stats stores all provided parameters."""
        collector = HeatmapCollector(sample_interval_usec=100)

        collector.set_stats(
            sample_interval_usec=500,
            duration_sec=10.5,
            sample_rate=99.5,
            error_rate=0.5
        )

        self.assertEqual(collector.stats['sample_interval_usec'], 500)
        self.assertEqual(collector.stats['duration_sec'], 10.5)
        self.assertEqual(collector.stats['sample_rate'], 99.5)
        self.assertEqual(collector.stats['error_rate'], 0.5)

    def test_set_stats_includes_system_info(self):
        """Test that set_stats includes Python and platform info."""
        collector = HeatmapCollector(sample_interval_usec=100)
        collector.set_stats(sample_interval_usec=100, duration_sec=1.0, sample_rate=100.0)

        self.assertIn('python_version', collector.stats)
        self.assertIn('python_implementation', collector.stats)
        self.assertIn('platform', collector.stats)

    def test_set_stats_accepts_kwargs(self):
        """Test that set_stats accepts additional kwargs."""
        collector = HeatmapCollector(sample_interval_usec=100)

        collector.set_stats(
            sample_interval_usec=100,
            duration_sec=1.0,
            sample_rate=100.0,
            custom_key='custom_value',
            another_key=42
        )

        self.assertEqual(collector.stats['custom_key'], 'custom_value')
        self.assertEqual(collector.stats['another_key'], 42)

    def test_set_stats_with_none_error_rate(self):
        """Test set_stats with error_rate=None."""
        collector = HeatmapCollector(sample_interval_usec=100)
        collector.set_stats(sample_interval_usec=100, duration_sec=1.0, sample_rate=100.0)

        self.assertIn('error_rate', collector.stats)
        self.assertIsNone(collector.stats['error_rate'])


class TestHeatmapCollectorProcessFrames(unittest.TestCase):
    """Test HeatmapCollector.process_frames() method."""

    def test_process_frames_increments_total_samples(self):
        """Test that process_frames increments total samples count."""
        collector = HeatmapCollector(sample_interval_usec=100)

        initial_count = collector._total_samples
        frames = [('file.py', (10, 10, -1, -1), 'func', None)]
        collector.process_frames(frames, thread_id=1)

        self.assertEqual(collector._total_samples, initial_count + 1)

    def test_process_frames_records_line_samples(self):
        """Test that process_frames records line samples."""
        collector = HeatmapCollector(sample_interval_usec=100)

        frames = [('test.py', (5, 5, -1, -1), 'test_func', None)]
        collector.process_frames(frames, thread_id=1)

        # Check that line was recorded
        self.assertIn(('test.py', 5), collector.line_samples)
        self.assertEqual(collector.line_samples[('test.py', 5)], 1)

    def test_process_frames_records_multiple_lines_in_stack(self):
        """Test that process_frames records all lines in a stack."""
        collector = HeatmapCollector(sample_interval_usec=100)

        frames = [
            ('file1.py', (10, 10, -1, -1), 'func1', None),
            ('file2.py', (20, 20, -1, -1), 'func2', None),
            ('file3.py', (30, 30, -1, -1), 'func3', None)
        ]
        collector.process_frames(frames, thread_id=1)

        # All frames should be recorded
        self.assertIn(('file1.py', 10), collector.line_samples)
        self.assertIn(('file2.py', 20), collector.line_samples)
        self.assertIn(('file3.py', 30), collector.line_samples)

    def test_process_frames_distinguishes_self_samples(self):
        """Test that process_frames distinguishes self (leaf) samples."""
        collector = HeatmapCollector(sample_interval_usec=100)

        frames = [
            ('leaf.py', (5, 5, -1, -1), 'leaf_func', None),  # This is the leaf (top of stack)
            ('caller.py', (10, 10, -1, -1), 'caller_func', None)
        ]
        collector.process_frames(frames, thread_id=1)

        # Leaf should have self sample
        self.assertIn(('leaf.py', 5), collector.line_self_samples)
        self.assertEqual(collector.line_self_samples[('leaf.py', 5)], 1)

        # Caller should NOT have self sample
        self.assertNotIn(('caller.py', 10), collector.line_self_samples)

    def test_process_frames_accumulates_samples(self):
        """Test that multiple calls accumulate samples."""
        collector = HeatmapCollector(sample_interval_usec=100)

        frames = [('file.py', (10, 10, -1, -1), 'func', None)]

        collector.process_frames(frames, thread_id=1)
        collector.process_frames(frames, thread_id=1)
        collector.process_frames(frames, thread_id=1)

        self.assertEqual(collector.line_samples[('file.py', 10)], 3)
        self.assertEqual(collector._total_samples, 3)

    def test_process_frames_ignores_invalid_frames(self):
        """Test that process_frames ignores invalid frames."""
        collector = HeatmapCollector(sample_interval_usec=100)

        # These should be ignored
        invalid_frames = [
            ('<string>', (1, 1, -1, -1), 'test', None),
            ('[eval]', (1, 1, -1, -1), 'test', None),
            ('', (1, 1, -1, -1), 'test', None),
            (None, (1, 1, -1, -1), 'test', None),
            ('__init__', (0, 0, -1, -1), 'test', None),  # Special invalid frame
        ]

        for frame in invalid_frames:
            collector.process_frames([frame], thread_id=1)

        # Should not record these invalid frames
        for frame in invalid_frames:
            if frame[0]:
                self.assertNotIn((frame[0], frame[1][0]), collector.line_samples)

    def test_process_frames_builds_call_graph(self):
        """Test that process_frames builds call graph relationships."""
        collector = HeatmapCollector(sample_interval_usec=100)

        frames = [
            ('callee.py', (5, 5, -1, -1), 'callee_func', None),
            ('caller.py', (10, 10, -1, -1), 'caller_func', None)
        ]
        collector.process_frames(frames, thread_id=1)

        # Check that call relationship was recorded
        caller_key = ('caller.py', 10)
        self.assertIn(caller_key, collector.call_graph)

        # Check callers graph
        callee_key = ('callee.py', 5)
        self.assertIn(callee_key, collector.callers_graph)

    def test_process_frames_records_function_definitions(self):
        """Test that process_frames records function definition locations."""
        collector = HeatmapCollector(sample_interval_usec=100)

        frames = [('module.py', (42, 42, -1, -1), 'my_function', None)]
        collector.process_frames(frames, thread_id=1)

        self.assertIn(('module.py', 'my_function'), collector.function_definitions)
        self.assertEqual(collector.function_definitions[('module.py', 'my_function')], 42)

    def test_process_frames_tracks_edge_samples(self):
        """Test that process_frames tracks edge sample counts."""
        collector = HeatmapCollector(sample_interval_usec=100)

        frames = [
            ('callee.py', (5, 5, -1, -1), 'callee', None),
            ('caller.py', (10, 10, -1, -1), 'caller', None)
        ]

        # Process same call stack multiple times
        collector.process_frames(frames, thread_id=1)
        collector.process_frames(frames, thread_id=1)

        # Check that edge count is tracked
        self.assertGreater(len(collector.edge_samples), 0)

    def test_process_frames_weight_applies_to_identical_samples(self):
        collector = HeatmapCollector(sample_interval_usec=100)

        frames = [
            ('callee.py', (5, 5, -1, -1), 'callee', None),
            ('caller.py', (10, 10, -1, -1), 'caller', None),
        ]

        collector.process_frames(frames, thread_id=1, weight=5)

        edge_key = (('caller.py', 10), ('callee.py', 5))
        self.assertEqual(collector.edge_samples[edge_key], 5)
        self.assertEqual(collector.line_samples[('callee.py', 5)], 5)
        self.assertEqual(collector.line_samples[('caller.py', 10)], 5)

    def test_process_frames_handles_empty_frames(self):
        """Test that process_frames handles empty frame list."""
        collector = HeatmapCollector(sample_interval_usec=100)

        initial_count = collector._total_samples
        collector.process_frames([], thread_id=1)

        # Should still increment total samples
        self.assertEqual(collector._total_samples, initial_count + 1)

    def test_process_frames_with_file_samples_dict(self):
        """Test that file_samples dict is properly populated."""
        collector = HeatmapCollector(sample_interval_usec=100)

        frames = [('test.py', (10, 10, -1, -1), 'func', None)]
        collector.process_frames(frames, thread_id=1)

        self.assertIn('test.py', collector.file_samples)
        self.assertIn(10, collector.file_samples['test.py'])
        self.assertEqual(collector.file_samples['test.py'][10], 1)


def frame(filename, line, func):
    """Create a frame tuple: (filename, location, funcname, opcode)."""
    return (filename, (line, line, -1, -1), func, None)


class TestHeatmapCollectorNavigationButtons(unittest.TestCase):
    """Test navigation button behavior for caller/callee relationships.

    For every call stack:
    - Root frames (entry points): only DOWN arrow (callees)
    - Middle frames: both UP and DOWN arrows
    - Leaf frames: only UP arrow (callers)
    """

    def collect(self, *stacks):
        """Create collector and process frame stacks."""
        collector = HeatmapCollector(sample_interval_usec=100)
        for stack in stacks:
            collector.process_frames(stack, thread_id=1)
        return collector

    def test_deep_call_stack_relationships(self):
        """Test root/middle/leaf navigation in a 5-level call stack."""
        # Stack: root -> A -> B -> C -> leaf
        stack = [
            frame('leaf.py', 5, 'leaf'),
            frame('c.py', 10, 'func_c'),
            frame('b.py', 15, 'func_b'),
            frame('a.py', 20, 'func_a'),
            frame('root.py', 25, 'root'),
        ]
        c = self.collect(stack)

        # Root: only callees (no one calls it)
        self.assertIn(('root.py', 25), c.call_graph)
        self.assertNotIn(('root.py', 25), c.callers_graph)

        # Middle frames: both callers and callees
        for key in [('a.py', 20), ('b.py', 15), ('c.py', 10)]:
            self.assertIn(key, c.call_graph)
            self.assertIn(key, c.callers_graph)

        # Leaf: only callers (doesn't call anyone)
        self.assertNotIn(('leaf.py', 5), c.call_graph)
        self.assertIn(('leaf.py', 5), c.callers_graph)

    def test_all_lines_in_function_see_callers(self):
        """Test that interior lines map to their function for caller lookup."""
        # Same function sampled at different lines (12, 15, 10)
        c = self.collect(
            [frame('mod.py', 12, 'my_func'), frame('caller.py', 100, 'caller')],
            [frame('mod.py', 15, 'my_func'), frame('caller.py', 100, 'caller')],
            [frame('mod.py', 10, 'my_func'), frame('caller.py', 100, 'caller')],
        )

        # All lines should map to same function
        for line in [10, 12, 15]:
            self.assertEqual(c.line_to_function[('mod.py', line)], 'my_func')

        # Function definition line should have callers
        func_def = c.function_definitions[('mod.py', 'my_func')]
        self.assertIn(('mod.py', func_def), c.callers_graph)

    def test_multiple_callers_and_callees(self):
        """Test multiple callers/callees are recorded correctly."""
        # Two callers -> target, and caller -> two callees
        c = self.collect(
            [frame('target.py', 10, 'target'), frame('caller1.py', 20, 'c1')],
            [frame('target.py', 10, 'target'), frame('caller2.py', 30, 'c2')],
            [frame('callee1.py', 5, 'f1'), frame('dispatcher.py', 40, 'dispatch')],
            [frame('callee2.py', 6, 'f2'), frame('dispatcher.py', 40, 'dispatch')],
        )

        # Target has 2 callers
        callers = c.callers_graph[('target.py', 10)]
        self.assertEqual({x[0] for x in callers}, {'caller1.py', 'caller2.py'})

        # Dispatcher has 2 callees
        callees = c.call_graph[('dispatcher.py', 40)]
        self.assertEqual({x[0] for x in callees}, {'callee1.py', 'callee2.py'})

    def test_edge_samples_counted(self):
        """Test that repeated calls accumulate edge counts."""
        stack = [frame('callee.py', 10, 'callee'), frame('caller.py', 20, 'caller')]
        c = self.collect(stack, stack, stack)

        edge_key = (('caller.py', 20), ('callee.py', 10))
        self.assertEqual(c.edge_samples[edge_key], 3)


class TestHeatmapCollectorExport(unittest.TestCase):
    """Test HeatmapCollector.export() method."""

    def setUp(self):
        """Set up test directory."""
        self.test_dir = tempfile.mkdtemp()
        self.addCleanup(shutil.rmtree, self.test_dir)

    def test_export_creates_output_directory(self):
        """Test that export creates the output directory."""
        collector = HeatmapCollector(sample_interval_usec=100)

        # Add some data
        frames = [('test.py', (10, 10, -1, -1), 'func', None)]
        collector.process_frames(frames, thread_id=1)

        output_path = os.path.join(self.test_dir, 'heatmap_output')

        with captured_stdout(), captured_stderr():
            collector.export(output_path)

        self.assertTrue(os.path.exists(output_path))
        self.assertTrue(os.path.isdir(output_path))

    def test_export_creates_index_html(self):
        """Test that export creates index.html."""
        collector = HeatmapCollector(sample_interval_usec=100)

        frames = [('test.py', (10, 10, -1, -1), 'func', None)]
        collector.process_frames(frames, thread_id=1)

        output_path = os.path.join(self.test_dir, 'heatmap_output')

        with captured_stdout(), captured_stderr():
            collector.export(output_path)

        index_path = os.path.join(output_path, 'index.html')
        self.assertTrue(os.path.exists(index_path))

    def test_export_creates_file_htmls(self):
        """Test that export creates individual file HTMLs."""
        collector = HeatmapCollector(sample_interval_usec=100)

        frames = [('test.py', (10, 10, -1, -1), 'func', None)]
        collector.process_frames(frames, thread_id=1)

        output_path = os.path.join(self.test_dir, 'heatmap_output')

        with captured_stdout(), captured_stderr():
            collector.export(output_path)

        # Check for file_XXXX.html files
        html_files = [f for f in os.listdir(output_path)
                      if f.startswith('file_') and f.endswith('.html')]
        self.assertGreater(len(html_files), 0)

    def test_export_with_empty_data(self):
        """Test export with no data collected."""
        collector = HeatmapCollector(sample_interval_usec=100)

        output_path = os.path.join(self.test_dir, 'empty_output')

        # Should handle empty data gracefully
        with captured_stdout(), captured_stderr():
            collector.export(output_path)

    def test_export_handles_html_suffix(self):
        """Test that export handles .html suffix in output path."""
        collector = HeatmapCollector(sample_interval_usec=100)

        frames = [('test.py', (10, 10, -1, -1), 'func', None)]
        collector.process_frames(frames, thread_id=1)

        # Path with .html suffix should be stripped
        output_path = os.path.join(self.test_dir, 'output.html')

        with captured_stdout(), captured_stderr():
            collector.export(output_path)

        # Should create directory without .html
        expected_dir = os.path.join(self.test_dir, 'output')
        self.assertTrue(os.path.exists(expected_dir))

    def test_export_with_multiple_files(self):
        """Test export with multiple files."""
        collector = HeatmapCollector(sample_interval_usec=100)

        # Add samples for multiple files
        collector.process_frames([('file1.py', (10, 10, -1, -1), 'func1', None)], thread_id=1)
        collector.process_frames([('file2.py', (20, 20, -1, -1), 'func2', None)], thread_id=1)
        collector.process_frames([('file3.py', (30, 30, -1, -1), 'func3', None)], thread_id=1)

        output_path = os.path.join(self.test_dir, 'multi_file')

        with captured_stdout(), captured_stderr():
            collector.export(output_path)

        # Should create HTML for each file
        html_files = [f for f in os.listdir(output_path)
                      if f.startswith('file_') and f.endswith('.html')]
        self.assertGreaterEqual(len(html_files), 3)

    def test_export_index_contains_file_references(self):
        """Test that index.html contains references to profiled files."""
        collector = HeatmapCollector(sample_interval_usec=100)
        collector.set_stats(sample_interval_usec=100, duration_sec=1.0, sample_rate=100.0)

        frames = [('mytest.py', (10, 10, -1, -1), 'my_func', None)]
        collector.process_frames(frames, thread_id=1)

        output_path = os.path.join(self.test_dir, 'test_output')

        with captured_stdout(), captured_stderr():
            collector.export(output_path)

        index_path = os.path.join(output_path, 'index.html')
        with open(index_path, 'r', encoding='utf-8') as f:
            content = f.read()

        # Should contain reference to the file
        self.assertIn('mytest', content)

    def test_export_file_html_has_line_numbers(self):
        """Test that exported file HTML contains line numbers."""
        collector = HeatmapCollector(sample_interval_usec=100)

        # Create a temporary Python file
        temp_file = os.path.join(self.test_dir, 'temp_source.py')
        with open(temp_file, 'w') as f:
            f.write('def test():\n    pass\n')

        frames = [(temp_file, (1, 1, -1, -1), 'test', None)]
        collector.process_frames(frames, thread_id=1)

        output_path = os.path.join(self.test_dir, 'line_test')

        with captured_stdout(), captured_stderr():
            collector.export(output_path)

        # Find the generated file HTML
        html_files = [f for f in os.listdir(output_path)
                      if f.startswith('file_') and f.endswith('.html')]

        self.assertEqual(len(html_files), 1)
        with open(os.path.join(output_path, html_files[0]), 'r', encoding='utf-8') as f:
            content = f.read()

        # Should have line-related content
        self.assertIn('line-', content)

    def test_export_skips_nonexistent_source(self):
        self.check_unavailable_source(os.path.join(self.test_dir, 'missing.py'))

    def test_export_skips_directory_source(self):
        self.check_unavailable_source(self.test_dir)

    @unittest.skipUnless(hasattr(os, 'mkfifo'), 'requires os.mkfifo')
    def test_export_skips_fifo_source(self):
        filename = os.path.join(self.test_dir, 'source.fifo')
        os.mkfifo(filename)
        self.check_unavailable_source(filename)

    def check_unavailable_source(self, filename):
        collector = HeatmapCollector(sample_interval_usec=100)
        frames = [(filename, (1, 1, -1, -1), 'f', None)]
        collector.process_frames(frames, thread_id=1)
        output_path = Path(self.test_dir) / 'unavailable_source'
        read_text = Path.read_text

        def check_read(path, *args, **kwargs):
            self.assertNotEqual(path, Path(filename),
                                'Non-regular source must not be opened')
            return read_text(path, *args, **kwargs)

        with (captured_stdout(), captured_stderr(),
              mock.patch.object(Path, 'read_text', check_read)):
            collector.export(output_path)
        content = (output_path / collector.file_index[filename]).read_text(
            encoding='utf-8')
        self.assertIn('Source file not available', content)

class MockFrameInfo:
    """Mock FrameInfo for testing.

    Frame format: (filename, location, funcname, opcode) where:
    - location is a tuple (lineno, end_lineno, col_offset, end_col_offset)
    - opcode is an int or None
    """

    def __init__(self, filename, lineno, funcname, opcode=None):
        self.filename = filename
        self.funcname = funcname
        self.opcode = opcode
        self.location = (lineno, lineno, -1, -1)

    def __iter__(self):
        return iter((self.filename, self.location, self.funcname, self.opcode))

    def __getitem__(self, index):
        return (self.filename, self.location, self.funcname, self.opcode)[index]

    def __len__(self):
        return 4

    def __repr__(self):
        return f"MockFrameInfo('{self.filename}', {self.location}, '{self.funcname}', {self.opcode})"


class MockThreadInfo:
    """Mock ThreadInfo for testing since the real one isn't accessible."""

    def __init__(self, thread_id, frame_info, status=0):
        self.thread_id = thread_id
        self.frame_info = frame_info
        self.status = status  # Thread status flags

    def __repr__(self):
        return f"MockThreadInfo(thread_id={self.thread_id}, frame_info={self.frame_info})"


class MockInterpreterInfo:
    """Mock InterpreterInfo for testing since the real one isn't accessible."""

    def __init__(self, interpreter_id, threads):
        self.interpreter_id = interpreter_id
        self.threads = threads

    def __repr__(self):
        return f"MockInterpreterInfo(interpreter_id={self.interpreter_id}, threads={self.threads})"


class TestHeatmapCollector(unittest.TestCase):
    """Tests for HeatmapCollector functionality."""

    def test_heatmap_collector_basic(self):
        """Test basic HeatmapCollector functionality."""
        collector = HeatmapCollector(sample_interval_usec=100)

        # Test empty state
        self.assertEqual(len(collector.file_samples), 0)
        self.assertEqual(len(collector.line_samples), 0)

        # Test collecting sample data - frames are 4-tuples: (filename, location, funcname, opcode)
        test_frames = [
            MockInterpreterInfo(
                0,
                [MockThreadInfo(
                    1,
                    [MockFrameInfo("file.py", 10, "func1"), MockFrameInfo("file.py", 20, "func2")],
                )]
            )
        ]
        collector.collect(test_frames)

        # Should have recorded samples for the file
        self.assertGreater(len(collector.line_samples), 0)
        self.assertIn("file.py", collector.file_samples)

        # Check that line samples were recorded
        file_data = collector.file_samples["file.py"]
        self.assertGreater(len(file_data), 0)

    def test_heatmap_collector_export(self):
        """Test heatmap HTML export functionality."""
        heatmap_dir = tempfile.mkdtemp()
        self.addCleanup(shutil.rmtree, heatmap_dir)

        collector = HeatmapCollector(sample_interval_usec=100)

        # Create test data with multiple files using MockFrameInfo
        test_frames1 = [
            MockInterpreterInfo(
                0,
                [MockThreadInfo(1, [MockFrameInfo("file.py", 10, "func1"), MockFrameInfo("file.py", 20, "func2")])],
            )
        ]
        test_frames2 = [
            MockInterpreterInfo(
                0,
                [MockThreadInfo(1, [MockFrameInfo("file.py", 10, "func1"), MockFrameInfo("file.py", 20, "func2")])],
            )
        ]  # Same stack
        test_frames3 = [
            MockInterpreterInfo(0, [MockThreadInfo(1, [MockFrameInfo("other.py", 5, "other_func")])])
        ]

        collector.collect(test_frames1)
        collector.collect(test_frames2)
        collector.collect(test_frames3)

        # Export heatmap
        with (captured_stdout(), captured_stderr()):
            collector.export(heatmap_dir)

        # Verify index.html was created
        index_path = os.path.join(heatmap_dir, "index.html")
        self.assertTrue(os.path.exists(index_path))
        self.assertGreater(os.path.getsize(index_path), 0)

        # Check index contains HTML content
        with open(index_path, "r", encoding="utf-8") as f:
            content = f.read()

        # Should be valid HTML
        self.assertIn("<!doctype html>", content.lower())
        self.assertIn("<html", content)
        self.assertIn("Tachyon Profiler", content)

        # Should contain file references
        self.assertIn("file.py", content)
        self.assertIn("other.py", content)

        # Verify individual file HTMLs were created
        file_htmls = [f for f in os.listdir(heatmap_dir) if f.startswith("file_") and f.endswith(".html")]
        self.assertGreater(len(file_htmls), 0)

        # Check one of the file HTMLs
        file_html_path = os.path.join(heatmap_dir, file_htmls[0])
        with open(file_html_path, "r", encoding="utf-8") as f:
            file_content = f.read()

        # Should contain heatmap styling and JavaScript
        self.assertIn("line-sample", file_content)
        self.assertIn("nav-btn", file_content)


class TestHeatmapCollectorLocation(unittest.TestCase):
    """Tests for HeatmapCollector location handling."""

    def test_heatmap_with_full_location_info(self):
        """Test HeatmapCollector uses full location tuple."""
        collector = HeatmapCollector(sample_interval_usec=1000)

        # Frame with full location: (lineno, end_lineno, col_offset, end_col_offset)
        frame = MockFrameInfo("test.py", 10, "func")
        # Override with full location info
        frame.location = LocationInfo(10, 15, 4, 20)
        frames = [
            MockInterpreterInfo(
                0,
                [MockThreadInfo(1, [frame])]
            )
        ]
        collector.collect(frames)

        # Verify data was collected with location info
        # HeatmapCollector uses file_samples dict with filename -> Counter of linenos
        self.assertIn("test.py", collector.file_samples)
        # Line 10 should have samples
        self.assertIn(10, collector.file_samples["test.py"])

    def test_heatmap_with_none_location(self):
        """Test HeatmapCollector handles None location gracefully."""
        collector = HeatmapCollector(sample_interval_usec=1000)

        # Synthetic frame with None location
        frame = MockFrameInfo("~", 0, "<native>")
        frame.location = None
        frames = [
            MockInterpreterInfo(
                0,
                [MockThreadInfo(1, [frame])]
            )
        ]
        # Should not raise
        collector.collect(frames)

    def test_heatmap_export_with_location_data(self):
        """Test HeatmapCollector export includes location info."""
        tmp_dir = tempfile.mkdtemp()
        self.addCleanup(shutil.rmtree, tmp_dir)

        collector = HeatmapCollector(sample_interval_usec=1000)

        frame = MockFrameInfo("test.py", 10, "process")
        frame.location = LocationInfo(10, 12, 0, 30)
        frames = [
            MockInterpreterInfo(
                0,
                [MockThreadInfo(1, [frame])]
            )
        ]
        collector.collect(frames)

        # Export should work
        with (captured_stdout(), captured_stderr()):
            collector.export(tmp_dir)
        self.assertTrue(os.path.exists(os.path.join(tmp_dir, "index.html")))

    def test_heatmap_collector_frame_format(self):
        """Test HeatmapCollector with 4-element frame format."""
        collector = HeatmapCollector(sample_interval_usec=1000)

        frames = [
            MockInterpreterInfo(
                0,
                [
                    MockThreadInfo(
                        1,
                        [
                            MockFrameInfo("app.py", 100, "main", opcode=90),
                            MockFrameInfo("utils.py", 50, "helper", opcode=100),
                            MockFrameInfo("lib.py", 25, "process", opcode=None),
                        ],
                    )
                ],
            )
        ]
        collector.collect(frames)

        # Should have recorded data for the files
        self.assertIn("app.py", collector.file_samples)
        self.assertIn("utils.py", collector.file_samples)
        self.assertIn("lib.py", collector.file_samples)


if __name__ == "__main__":
    unittest.main()