import re
import random

class VariantGenerator:
    def __init__(self):
        print("[System] Initializing FIB and True/False Generator...")

    def generate_fib(self, sentence, answer):
        """
        Creates a Fill-in-the-Blanks question by replacing the answer phrase.
        """
        if not answer or answer.lower() not in sentence.lower():
            return None, None

        # Use regex for a case-insensitive replacement, replacing only the first occurrence
        # This prevents breaking the sentence if the word appears twice
        pattern = re.compile(re.escape(answer), re.IGNORECASE)
        fib_question = pattern.sub("________", sentence, count=1)
        
        return fib_question, answer

    def generate_tf(self, sentence, answer, top_distractor):
        """
        Creates a True/False statement.
        It randomly decides whether to output the True (original) statement
        or a False (manipulated) statement using the best distractor.
        """
        if not answer or answer.lower() not in sentence.lower():
            return None, None

        # Randomly decide if this will be a True statement or a False statement
        is_true = random.choice([True, False])

        if is_true or not top_distractor:
            # The True Statement is simply the original factual sentence
            return sentence, "True"
        else:
            # The False Statement replaces the correct answer with the best distractor
            pattern = re.compile(re.escape(answer), re.IGNORECASE)
            false_statement = pattern.sub(top_distractor, sentence, count=1)
            
            return false_statement, "False"

# ==========================================
# TEST EXECUTION
# ==========================================
if __name__ == "__main__":
    generator = VariantGenerator()
    
    # Mock data routed from your 3_Question_Gen.py and 6_Distractors_Eval.py
    original_sentence = "Capitalism is an economic system based on the private ownership of the means of production and their operation for profit."
    extracted_answer = "capitalism"
    
    # Assume this is the top distractor that survived your Stage 3 pipeline in File 6
    best_distractor = "socialism"

    print("\n" + "="*50)
    print("🛠️  TESTING FIB GENERATION")
    print("="*50)
    fib_q, fib_a = generator.generate_fib(original_sentence, extracted_answer)
    print(f"Question: {fib_q}")
    print(f"Answer: {fib_a}")

    print("\n" + "="*50)
    print("⚖️  TESTING TRUE/FALSE GENERATION")
    print("="*50)
    
    # Run it a few times to see both True and False outcomes
    for i in range(3):
        tf_q, tf_a = generator.generate_tf(original_sentence, extracted_answer, best_distractor)
        print(f"Statement {i+1}: {tf_q}")
        print(f"Answer: {tf_a}\n")