import unittest
import os
import tempfile
import tools.aggregate_regression as ar

class TestAggregateRegression(unittest.TestCase):

    def setUp(self):
        # Reset the global dictionaries in the aggregator module
        ar.SCENARIOS_MAP.clear()
        ar.SCENARIO_IDS_SEEN.clear()
        ar.QUARANTINED_INFO.clear()
        ar.SCENARIO_ID_TO_TIER.clear()
        
        # Save original directories
        self.orig_test_src = ar.TEST_SRC_DIR
        self.orig_benchmark_src = ar.BENCHMARK_SRC_DIR
        self.orig_history_json = ar.HISTORY_JSON
        self.orig_dashboard_md = ar.DASHBOARD_MD
        self.orig_additional_output = ar.ADDITIONAL_OUTPUT_DIR
        
        # Create temp directories and files
        self.temp_dir = tempfile.TemporaryDirectory()
        ar.TEST_SRC_DIR = os.path.join(self.temp_dir.name, "test_src")
        ar.BENCHMARK_SRC_DIR = os.path.join(self.temp_dir.name, "benchmark_src")
        os.makedirs(ar.TEST_SRC_DIR)
        os.makedirs(ar.BENCHMARK_SRC_DIR)
        
        ar.HISTORY_JSON = os.path.join(self.temp_dir.name, "history.json")
        ar.DASHBOARD_MD = os.path.join(self.temp_dir.name, "dashboard.md")
        ar.ADDITIONAL_OUTPUT_DIR = os.path.join(self.temp_dir.name, "additional_output")
        os.makedirs(ar.ADDITIONAL_OUTPUT_DIR)

    def tearDown(self):
        # Restore original paths
        ar.TEST_SRC_DIR = self.orig_test_src
        ar.BENCHMARK_SRC_DIR = self.orig_benchmark_src
        ar.HISTORY_JSON = self.orig_history_json
        ar.DASHBOARD_MD = self.orig_dashboard_md
        ar.ADDITIONAL_OUTPUT_DIR = self.orig_additional_output
        self.temp_dir.cleanup()

    def write_kotlin_file(self, content, filename="TestFile.kt", is_benchmark=False):
        dir_path = ar.BENCHMARK_SRC_DIR if is_benchmark else ar.TEST_SRC_DIR
        file_path = os.path.join(dir_path, filename)
        with open(file_path, "w", encoding="utf-8") as f:
            f.write(content)
        return file_path

    def test_parse_kotlin_sources_multiple_classes(self):
        # Write a Kotlin file containing multiple classes, with different annotations and tiers
        kotlin_content = """
        package com.yorvana.regression
        
        import com.yorvana.testsupport.tiers.Regression
        import com.yorvana.testsupport.tiers.RegressionFull
        import com.yorvana.testsupport.tiers.ScenarioId
        import com.yorvana.testsupport.tiers.Quarantined
        import org.junit.Test

        @Regression
        class FirstClassTest {
            @Test
            fun rV01_firstTest() = Unit
            
            @Quarantined(reason = "mixed-order", since = "2026-05-22", issue = 120)
            @Test
            fun rV02_secondTest() = Unit
        }
        
        @RegressionFull
        class SecondClassTest {
            @ScenarioId("CUSTOM-S01")
            @Test
            fun customScenarioTest() = Unit
            
            fun helperMethod() = Unit
        }
        
        class UnannotatedClass {
            @Test
            fun rV03_testInUnannotated() = Unit
        }
        """
        self.write_kotlin_file(kotlin_content)
        ar.parse_kotlin_sources()
        
        # Verify first class test methods
        self.assertIn("com.yorvana.regression.FirstClassTest#rV01_firstTest", ar.SCENARIOS_MAP)
        self.assertEqual(ar.SCENARIOS_MAP["com.yorvana.regression.FirstClassTest#rV01_firstTest"]["scenario_id"], "R-V01")
        self.assertEqual(ar.SCENARIOS_MAP["com.yorvana.regression.FirstClassTest#rV01_firstTest"]["tier"], "Regression")
        self.assertFalse(ar.SCENARIOS_MAP["com.yorvana.regression.FirstClassTest#rV01_firstTest"]["quarantined"])

        # Verify quarantined secondTest (with out-of-order quarantine args)
        self.assertIn("com.yorvana.regression.FirstClassTest#rV02_secondTest", ar.SCENARIOS_MAP)
        self.assertTrue(ar.SCENARIOS_MAP["com.yorvana.regression.FirstClassTest#rV02_secondTest"]["quarantined"])
        self.assertEqual(ar.SCENARIOS_MAP["com.yorvana.regression.FirstClassTest#rV02_secondTest"]["q_meta"]["reason"], "mixed-order")
        self.assertEqual(ar.SCENARIOS_MAP["com.yorvana.regression.FirstClassTest#rV02_secondTest"]["q_meta"]["since"], "2026-05-22")
        self.assertEqual(ar.SCENARIOS_MAP["com.yorvana.regression.FirstClassTest#rV02_secondTest"]["q_meta"]["issue"], 120)

        # Verify second class test methods
        self.assertIn("com.yorvana.regression.SecondClassTest#customScenarioTest", ar.SCENARIOS_MAP)
        self.assertEqual(ar.SCENARIOS_MAP["com.yorvana.regression.SecondClassTest#customScenarioTest"]["scenario_id"], "CUSTOM-S01")
        self.assertEqual(ar.SCENARIOS_MAP["com.yorvana.regression.SecondClassTest#customScenarioTest"]["tier"], "RegressionFull")
        
        # Verify that helper method (not annotated @Test/@Benchmark) is skipped
        self.assertNotIn("com.yorvana.regression.SecondClassTest#helperMethod", ar.SCENARIOS_MAP)
        
        # Verify unannotated class is parsed and maps to default Regression tier
        self.assertIn("com.yorvana.regression.UnannotatedClass#rV03_testInUnannotated", ar.SCENARIOS_MAP)
        self.assertEqual(ar.SCENARIOS_MAP["com.yorvana.regression.UnannotatedClass#rV03_testInUnannotated"]["tier"], "Regression")

    def test_streak_calculations(self):
        # We test streak logic manually by simulating outcomes_history
        # Green streak: pass/flake increments, fail resets, skip/missing neutral.
        # Pass streak: pass increments, fail/flake resets, skip/missing neutral.
        
        stats = {
            "outcomes_history": ["pass", "pass", "skip", "pass", "flake", "missing", "pass", "fail", "pass"]
        }
        
        # Simulating streak logic in main
        # pass, pass -> green=2, pass=2
        # skip -> green=2, pass=2
        # pass -> green=3, pass=3
        # flake -> green=4, pass=0 (resets)
        # missing -> green=4, pass=0
        # pass -> green=5, pass=1
        # fail -> green=0, pass=0 (resets)
        # pass -> green=1, pass=1
        
        consec_green = 0
        consec_passes = 0
        for outcome in stats["outcomes_history"]:
            if outcome == "fail":
                consec_green = 0
            elif outcome in ("pass", "flake"):
                consec_green += 1

            if outcome in ("fail", "flake"):
                consec_passes = 0
            elif outcome == "pass":
                consec_passes += 1
                
        self.assertEqual(consec_green, 1)
        self.assertEqual(consec_passes, 1)
        
        # Test another sequence
        stats2 = {
            "outcomes_history": ["pass", "flake", "pass", "pass", "skip", "pass"]
        }
        # pass -> green=1, pass=1
        # flake -> green=2, pass=0
        # pass, pass -> green=4, pass=2
        # skip -> green=4, pass=2
        # pass -> green=5, pass=3
        consec_green = 0
        consec_passes = 0
        for outcome in stats2["outcomes_history"]:
            if outcome == "fail":
                consec_green = 0
            elif outcome in ("pass", "flake"):
                consec_green += 1

            if outcome in ("fail", "flake"):
                consec_passes = 0
            elif outcome == "pass":
                consec_passes += 1
                
        self.assertEqual(consec_green, 5)
        self.assertEqual(consec_passes, 3)

    def test_check_flake_file(self):
        classname = "com.yorvana.regression.ExampleTest"
        methodname = "testMethod"
        
        # Before creating file, check_flake_file is False
        self.assertFalse(ar.check_flake_file(classname, methodname))
        
        # Create a flake file in ADDITIONAL_OUTPUT_DIR
        test_key = f"{classname}#{methodname}"
        sanitized_key = test_key.replace("[", "_").replace("]", "_").replace("=", "_").replace(" ", "_").replace("#", "_")
        flake_file_path = os.path.join(ar.ADDITIONAL_OUTPUT_DIR, f"flake-{sanitized_key}.txt")
        with open(flake_file_path, "w") as f:
            f.write("flaked")
            
        self.assertTrue(ar.check_flake_file(classname, methodname))

    def test_parse_kotlin_sources_class_regex_robustness(self):
        kotlin_content = """
        package com.yorvana.regression
        
        // This is a comment discussing some class which shouldn't match.
        /* Multi-line comment referring to class AnotherClass also ignored. */
        
        @Regression
        abstract class AbstractClassTest {
            @Test
            fun rV01_abstractTest() = Unit
        }
        
        sealed class SealedClass {
            @Test
            fun rV02_sealedTest() = Unit
        }
        
        open class OpenClass {
            @Test
            fun rV03_openTest() = Unit
        }
        
        @RegressionFull
        inner class InnerClassTest {
            @Test
            fun rV04_innerTest() = Unit
        }
        
        value class ValueClass {
            @Test
            fun rV05_valueTest() = Unit
        }
        
        class NormalClass {
            val someString = "This is a string literal containing the class word."
            @Test
            fun rV06_normalTest() = Unit
        }
        """
        self.write_kotlin_file(kotlin_content, filename="ClassRobustnessTest.kt")
        ar.parse_kotlin_sources()
        
        # Verify valid classes are parsed
        self.assertIn("com.yorvana.regression.AbstractClassTest#rV01_abstractTest", ar.SCENARIOS_MAP)
        self.assertIn("com.yorvana.regression.SealedClass#rV02_sealedTest", ar.SCENARIOS_MAP)
        self.assertIn("com.yorvana.regression.OpenClass#rV03_openTest", ar.SCENARIOS_MAP)
        self.assertIn("com.yorvana.regression.InnerClassTest#rV04_innerTest", ar.SCENARIOS_MAP)
        self.assertIn("com.yorvana.regression.ValueClass#rV05_valueTest", ar.SCENARIOS_MAP)
        self.assertIn("com.yorvana.regression.NormalClass#rV06_normalTest", ar.SCENARIOS_MAP)
        
        # Verify that comments and strings containing "class" did NOT create phantom classes
        # (By checking the parsed class names - only the valid classes should be in the map keys)
        parsed_classes = {k.split('#')[0] for k in ar.SCENARIOS_MAP.keys()}
        self.assertEqual(parsed_classes, {
            "com.yorvana.regression.AbstractClassTest",
            "com.yorvana.regression.SealedClass",
            "com.yorvana.regression.OpenClass",
            "com.yorvana.regression.InnerClassTest",
            "com.yorvana.regression.ValueClass",
            "com.yorvana.regression.NormalClass"
        })

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