
class MockModel:
    def generate(self, prompt, max_tokens, temperature):
        # Determine intent based on prompt for testing purposes
        # This is a stub for the real LLM service
        
        # If testing safety, we might see specific inputs
        if "test_safety.txt" in prompt:
             return '{"action": "tool", "tool_name": "write_file", "args": {"path": "test_safety.txt", "content": "This is a test of the safety system."}}'
             
        # Hive Test Candidates (Stricter matching to avoid context pollution)
        if "USER: We need a plan to launch" in prompt:
             return '{"action": "tool", "tool_name": "create_plan", "args": {"goal": "Launch the application"}}'
             
        if "USER: Write a python script called hello.py" in prompt:
             return '{"action": "tool", "tool_name": "write_file", "args": {"path": "hello.py", "content": "print(\'Hello World\')"}}'
             
        if "USER: Audit this code for bugs" in prompt:
             return '{"action": "reply", "content": "Auditing code... No vulnerabilities found."}'

        if "list_dir" in prompt or "list files" in prompt:
             return '{"action": "tool", "tool_name": "list_dir", "args": {"path": "."}}'

        return '{"action": "reply", "content": "I am a mocked brain. I cannot think yet."}'

class ModelFactory:
    @staticmethod
    def create(model_name):
        import os
        from dotenv import load_dotenv
        
        # Ensure env vars are loaded
        load_dotenv()
        
        gemini_key = os.getenv("GEMINI_API_KEY")
        
        if gemini_key:
            # 🚀 Use Real Intelligence
            from squadron.services.llm.gemini import GeminiProvider
            return GeminiProvider(gemini_key)
        else:
            # 🤡 Fallback to Mock
            print("⚠️  Warning: GEMINI_API_KEY not found. Using Mock Brain.")
            return MockModel()
