
import unittest
import time
import os
from dotenv import load_dotenv

load_dotenv()

from semantica.semantic_extract.ner_extractor import NERExtractor
from semantica.semantic_extract.relation_extractor import RelationExtractor
from semantica.semantic_extract.triplet_extractor import TripletExtractor
from semantica.semantic_extract.event_detector import EventDetector
from semantica.semantic_extract.semantic_network_extractor import SemanticNetworkExtractor
from semantica.semantic_extract.methods import _result_cache

class TestGroqRealWorldPerformance(unittest.TestCase):
    """
    Real-world performance test suite using Groq LLM.
    Tests parallel processing, caching, and correctness.
    """
    
    @classmethod
    def setUpClass(cls):
        cls.api_key = os.getenv("GROQ_API_KEY")
        if not cls.api_key:
            raise unittest.SkipTest("GROQ_API_KEY is not set")
        cls.metrics_file = os.path.join(os.getcwd(), "groq_metrics.txt")
        
        # Real-world sample texts (mix of tech, business, and general)
        cls.sample_texts = [
            """Apple Inc. is planning to launch a new AI-powered iPhone in late 2024. 
            CEO Tim Cook announced that the device will feature a neural engine capable of 
            processing 50 trillion operations per second. The company's stock rose 5% following the news.""",
            
            """Microsoft Corporation has acquired Activision Blizzard for $68.7 billion. 
            Satya Nadella, Microsoft's Chairman and CEO, stated that this acquisition will 
            accelerate growth in Microsoft's gaming business across mobile, PC, console, and cloud.""",
            
            """Elon Musk's SpaceX successfully launched the Starship rocket from Boca Chica, Texas. 
            The mission aims to test new heat shield technology essential for future Mars missions. 
            NASA Administrator Bill Nelson congratulated the team on the achievement.""",
            
            """Google DeepMind introduced Gemini, a new multimodal AI model. 
            Sundar Pichai emphasized that Gemini represents a significant leap forward in 
            AI capabilities, outperforming GPT-4 on several benchmarks including MMLU.""",
            
            """Amazon Web Services (AWS) announced a partnership with Anthropic to develop 
            reliable and high-performance foundation models. Amazon is investing up to $4 billion 
            in the AI safety startup founded by Dario Amodei."""
        ]
        
        # Warm up: Ensure modules are loaded
        print("\n[Setup] Initializing extractors...")
        cls.extractor = NERExtractor(
            method="llm",
            provider="groq",
            llm_model="llama-3.3-70b-versatile",
            api_key=cls.api_key,
        )
        cls.relation_extractor = RelationExtractor(
            method="llm",
            provider="groq",
            llm_model="llama-3.3-70b-versatile",
            api_key=cls.api_key,
        )
        cls.triplet_extractor = TripletExtractor(
            method="llm",
            provider="groq",
            llm_model="llama-3.3-70b-versatile",
            api_key=cls.api_key,
        )
        cls.event_detector = EventDetector(
            method="llm",
            provider="groq",
            llm_model="llama-3.3-70b-versatile",
            api_key=cls.api_key,
        )
        cls.network_extractor = SemanticNetworkExtractor(
            method="llm",
            provider="groq",
            llm_model="llama-3.3-70b-versatile",
            api_key=cls.api_key,
        )
        
    def setUp(self):
        # Clear cache before specific performance tests to ensure fair comparison
        # (Unless testing cache specifically)
        if _result_cache:
            _result_cache._caches["entities"].clear()
            _result_cache._caches["relations"].clear()
            
    def log_metrics(self, message):
        print(message)
        with open(self.metrics_file, "a") as f:
            f.write(message + "\n")
            f.flush()

    def test_01_parallel_vs_sequential_performance(self):
        """Compare sequential vs parallel extraction speed."""
        try:
            self.log_metrics("\n" + "="*60)
            self.log_metrics("TEST 1: Sequential vs Parallel Processing Performance")
            self.log_metrics("="*60)
            
            extractor = NERExtractor(method="llm", provider="groq", api_key=self.api_key, model="llama-3.3-70b-versatile")
            
            # 1. Sequential Run (Max workers = 1)
            self.log_metrics("\nStarting Sequential Extraction (5 documents)...")
            start_time = time.time()
            seq_results = extractor.extract(self.sample_texts, max_workers=1)
            seq_time = time.time() - start_time
            self.log_metrics(f"Sequential Time: {seq_time:.4f}s")
            self.log_metrics(f"Average Latency: {seq_time/len(self.sample_texts):.4f}s per doc")
            
            # Clear cache to force re-extraction for parallel test
            _result_cache._caches["entities"].clear()
            
            # 2. Parallel Run (Max workers = 5)
            self.log_metrics("\nStarting Parallel Extraction (5 documents, 5 workers)...")
            start_time = time.time()
            par_results = extractor.extract(self.sample_texts, max_workers=5)
            par_time = time.time() - start_time
            self.log_metrics(f"Parallel Time:   {par_time:.4f}s")
            self.log_metrics(f"Average Latency: {par_time/len(self.sample_texts):.4f}s per doc")
            
            # Analysis
            speedup = seq_time / par_time if par_time > 0 else 0
            self.log_metrics(f"\n>>> Performance Gain: {speedup:.2f}x speedup")
            self.log_metrics(f">>> Latency Reduction: {(seq_time - par_time):.4f}s total time saved")
            
            self.assertLess(par_time, seq_time * 1.35, "Parallel processing should not be significantly slower")
            self.assertEqual(len(seq_results), len(self.sample_texts))
            self.assertEqual(len(par_results), len(self.sample_texts))
        except Exception as e:
            self.log_metrics(f"ERROR in Test 1: {e}")
            raise

    def test_02_caching_latency_reduction(self):
        """Measure latency reduction from caching."""
        try:
            self.log_metrics("\n" + "="*60)
            self.log_metrics("TEST 2: Caching Performance & Latency Reduction")
            self.log_metrics("="*60)
            
            extractor = NERExtractor(method="llm", provider="groq", api_key=self.api_key, model="llama-3.3-70b-versatile")
            text = [self.sample_texts[0]]
            
            # 1. Cold Cache
            _result_cache._caches["entities"].clear()
            self.log_metrics("\nCold Cache Request...")
            start_time = time.time()
            extractor.extract(text)
            cold_time = time.time() - start_time
            self.log_metrics(f"Cold Cache Time: {cold_time:.4f}s")
            cache_size_after_cold = _result_cache.get_stats()["entities"]["size"]
            
            # 2. Warm Cache
            self.log_metrics("\nWarm Cache Request (Identical Query)...")
            start_time = time.time()
            extractor.extract(text)
            warm_time = time.time() - start_time
            self.log_metrics(f"Warm Cache Time: {warm_time:.6f}s")
            cache_size_after_warm = _result_cache.get_stats()["entities"]["size"]
            
            # Analysis
            reduction = (cold_time - warm_time) / cold_time * 100
            self.log_metrics(f"\n>>> Latency Reduction: {reduction:.2f}%")
            
            self.assertLess(warm_time, 1.0, "Warm cache response should be fast (<1.0s)")
            # self.assertGreater(reduction, 50, "Caching should reduce latency by >50%")
            if reduction < 30:
                self.log_metrics(f"WARNING: Caching reduction is low ({reduction:.2f}%)")
            self.assertGreater(reduction, 20, "Caching should reduce latency by >20%")
            self.assertGreater(cache_size_after_cold, 0, "Cache should store entity results")
            self.assertEqual(cache_size_after_warm, cache_size_after_cold, "Warm request should hit the cache")
        except Exception as e:
            self.log_metrics(f"ERROR in Test 2: {e}")
            raise

    def test_03_correctness_and_entity_matching(self):
        """Verify extraction correctness and data quality."""
        try:
            self.log_metrics("\n" + "="*60)
            self.log_metrics("TEST 3: Extraction Correctness & Data Quality")
            self.log_metrics("="*60)
            
            # Use a specific text with clear entities
            text = "Satya Nadella is the CEO of Microsoft."
            
            extractor = NERExtractor(method="llm", provider="groq", api_key=self.api_key, model="llama-3.3-70b-versatile")
            entities = extractor.extract([text])[0] # List of lists
            
            self.log_metrics(f"\nInput: {text}")
            self.log_metrics(f"Extracted Entities: {[e.text + '(' + e.label + ')' for e in entities]}")
            
            # Validation
            found_person = any(e.label == "PERSON" and "Satya" in e.text for e in entities)
            found_org = any(e.label == "ORG" and "Microsoft" in e.text for e in entities)
            
            if not found_org:
                self.log_metrics("FAILURE: Did not find Microsoft as ORG. Found entities:")
                for e in entities:
                    self.log_metrics(f" - {e.text}: {e.label}")

            self.assertTrue(found_person, "Failed to extract Satya Nadella as PERSON")
            self.assertTrue(found_org, "Failed to extract Microsoft as ORG")
            
            self.log_metrics("\n>>> Correctness Verification: PASS")
            self.log_metrics("    - Identified PERSON entity")
            self.log_metrics("    - Identified ORG entity")
            self.log_metrics("    - Pydantic models validated successfully")
        except Exception as e:
            self.log_metrics(f"ERROR in Test 3: {e}")
            raise

    def test_4_relation_extraction(self):
        """Test Relation Extraction capabilities"""
        print("\n" + "="*60)
        print("TEST 4: Relation Extraction")
        print("="*60)
        
        text = self.sample_texts[1]  # Microsoft acquisition text
        print(f"\nInput: {text[:100]}...")
        
        # First extract entities
        entities = self.__class__.extractor.extract_entities(text)
        self.assertTrue(len(entities) > 0, "Should extract entities first")
        
        # Extract relations
        start_time = time.time()
        relations = self.__class__.relation_extractor.extract_relations(text, entities)
        duration = time.time() - start_time
        
        print(f"Extracted {len(relations)} relations in {duration:.4f}s")
        for r in relations:
            print(f"  - {r.subject.text} -> {r.predicate} -> {r.object.text}")
            
        self.assertTrue(len(relations) > 0, "Should extract relations")
        
        # Verify specific relation (Microsoft -> acquired -> Activision Blizzard)
        found_acquisition = False
        for r in relations:
            if "Microsoft" in r.subject.text and "Activision" in r.object.text:
                found_acquisition = True
                break
        
        if not found_acquisition:
            # Fallback check - sometimes subject/object might be swapped or different wording
            for r in relations:
                if "Activision" in r.subject.text and "Microsoft" in r.object.text:
                    found_acquisition = True
                    break
                    
        self.assertTrue(found_acquisition, "Should find acquisition relation between Microsoft and Activision")

    def test_5_triplet_extraction(self):
        """Test RDF Triplet Extraction capabilities"""
        print("\n" + "="*60)
        print("TEST 5: Triplet Extraction")
        print("="*60)
        
        text = self.sample_texts[0]  # Apple text
        print(f"\nInput: {text[:100]}...")
        
        # Pipeline: Entities -> Relations -> Triplets
        entities = self.__class__.extractor.extract_entities(text)
        relations = self.__class__.relation_extractor.extract_relations(text, entities)
        
        start_time = time.time()
        triplets = self.__class__.triplet_extractor.extract_triplets(text, entities, relations)
        duration = time.time() - start_time
        
        print(f"Extracted {len(triplets)} triplets in {duration:.4f}s")
        for t in triplets:
            print(f"  - <{t.subject}> <{t.predicate}> <{t.object}>")
            
        self.assertTrue(len(triplets) > 0, "Should extract triplets")
        
        # Check for Apple related triplet
        found_apple = False
        for t in triplets:
            if "Apple" in t.subject or "Apple" in t.object:
                found_apple = True
                break
        self.assertTrue(found_apple, "Should find Apple-related triplet")

    def test_6_event_detection(self):
        """Test Event Detection capabilities"""
        print("\n" + "="*60)
        print("TEST 6: Event Detection")
        print("="*60)
        
        text = self.sample_texts[2]  # SpaceX launch text
        print(f"\nInput: {text[:100]}...")
        
        start_time = time.time()
        events = self.__class__.event_detector.detect_events(text)
        duration = time.time() - start_time
        
        print(f"Detected {len(events)} events in {duration:.4f}s")
        for e in events:
            print(f"  - [{e.event_type}] {e.text} (Participants: {e.participants})")
            
        self.assertTrue(len(events) > 0, "Should detect events")
        
        # Verify launch event
        found_launch = False
        for e in events:
            if "launch" in e.event_type.lower() or "launch" in e.text.lower():
                found_launch = True
                break
        self.assertTrue(found_launch, "Should detect launch event")

    def test_7_semantic_network(self):
        """Test Semantic Network Extraction capabilities"""
        print("\n" + "="*60)
        print("TEST 7: Semantic Network Extraction")
        print("="*60)
        
        text = self.sample_texts[3]  # Google DeepMind text
        print(f"\nInput: {text[:100]}...")
        
        # Extract base components first
        entities = self.__class__.extractor.extract_entities(text)
        relations = self.__class__.relation_extractor.extract_relations(text, entities)
        
        start_time = time.time()
        network = self.__class__.network_extractor.extract_network(text, entities=entities, relations=relations)
        duration = time.time() - start_time
        
        print(f"Extracted Network in {duration:.4f}s")
        print(f"  - Nodes: {len(network.nodes)}")
        print(f"  - Edges: {len(network.edges)}")
        
        self.assertTrue(len(network.nodes) > 0, "Should have nodes")
        self.assertTrue(len(network.edges) > 0, "Should have edges")
        
        # Verify Google/DeepMind/Gemini nodes exist
        node_labels = [n.label for n in network.nodes]
        print(f"  - Node Labels: {node_labels}")
        self.assertTrue(any("Gemini" in l for l in node_labels), "Should contain Gemini node")

    def test_8_parallel_event_detection(self):
        """Test Parallel Event Detection capabilities"""
        print("\n" + "="*60)
        print("TEST 8: Parallel Event Detection")
        print("="*60)
        
        # Create a larger batch by duplicating sample texts
        batch_texts = self.sample_texts * 2  # 10 documents
        
        # 1. Sequential Run
        print("\nStarting Sequential Event Detection (10 documents)...")
        start_time = time.time()
        seq_results = self.__class__.event_detector.extract(batch_texts, max_workers=1)
        seq_time = time.time() - start_time
        print(f"Sequential Time: {seq_time:.4f}s")
        
        # 2. Parallel Run
        print("\nStarting Parallel Event Detection (10 documents, 5 workers)...")
        start_time = time.time()
        par_results = self.__class__.event_detector.extract(batch_texts, max_workers=5)
        par_time = time.time() - start_time
        print(f"Parallel Time:   {par_time:.4f}s")
        
        # Analysis
        speedup = seq_time / par_time if par_time > 0 else 0
        print(f"\n>>> Performance Gain: {speedup:.2f}x speedup")
        
        self.assertEqual(len(seq_results), len(batch_texts))
        self.assertEqual(len(par_results), len(batch_texts))
        
        # Verify results match (order should be preserved)
        for i in range(len(batch_texts)):
            self.assertEqual(len(seq_results[i]), len(par_results[i]), f"Result count mismatch at index {i}")

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