# ai_provider.py
from abc import ABC, abstractmethod

class AIProvider(ABC):
    @abstractmethod
    def generate_response(self, prompt: str, model: str = "gpt-4") -> str:
        pass

    @abstractmethod
    def generate_structured(self, prompt: str, schema: dict) -> dict:
        pass

# mock_provider.py
import re
import json
from ai_provider import AIProvider

class MockAIProvider(AIProvider):
    def __init__(self):
        self.response_patterns = {
            r'priority.*score': '{"priority_score": 750}',
            r'decompose.*task': '''{
                "tasks": [
                    {"name": "Research", "priority": "high"},
                    {"name": "Analysis", "priority": "medium"}
                ]
            }''',
            r'team.*composition': '''{
                "team": [
                    {"name": "John", "role": "Developer"},
                    {"name": "Sarah", "role": "Designer"}
                ]
            }''',
            r'.*': 'Mock response for testing purposes'
        }

    def generate_response(self, prompt: str, model: str = "gpt-4") -> str:
        prompt_lower = prompt.lower()
        for pattern, response in self.response_patterns.items():
            if re.search(pattern, prompt_lower):
                return response
        return self.response_patterns[r'.*']

    def generate_structured(self, prompt: str, schema: dict) -> dict:
        response = self.generate_response(prompt)
        try:
            return json.loads(response)
        except json.JSONDecodeError:
            return {"mock": True, "response": response}

# openai_provider.py
import openai
from ai_provider import AIProvider
import json

class OpenAIProvider(AIProvider):
    def __init__(self, api_key: str):
        self.client = openai.OpenAI(api_key=api_key)

    def generate_response(self, prompt: str, model: str = "gpt-4") -> str:
        response = self.client.chat.completions.create(
            model=model,
            messages=[{"role": "user", "content": prompt}]
        )
        return response.choices[0].message.content

    def generate_structured(self, prompt: str, schema: dict) -> dict:
        schema_prompt = f"{prompt}\n\nRespond with valid JSON matching this schema: {schema}"
        response = self.generate_response(schema_prompt)
        return json.loads(response)

# ai_factory.py
import os
from mock_provider import MockAIProvider
from openai_provider import OpenAIProvider

class AIFactory:
    @staticmethod
    def create_provider():
        if os.getenv("TESTING") == "true":
            return MockAIProvider()
        if os.getenv("CI") == "true":
            return MockAIProvider()  # Không tốn tiền trong CI
        api_key = os.getenv("OPENAI_API_KEY")
        if not api_key:
            raise ValueError("OPENAI_API_KEY required for production")
        return OpenAIProvider(api_key)

# test_ai_agents.py
import os
import pytest

@pytest.fixture(autouse=True)
def setup_test_environment():
    os.environ["TESTING"] = "true"
    yield
    os.environ.pop("TESTING", None)

def test_task_prioritization():
    from ai_factory import AIFactory
    ai = AIFactory.create_provider()
    response = ai.generate_structured(
        "Calculate priority score for this task", {"priority_score": "number"}
    )
    assert "priority_score" in response
    assert isinstance(response["priority_score"], (int, str))
    assert response["priority_score"] == 750  # Xác định rõ ràng

def test_team_composition():
    from ai_factory import AIFactory
    ai = AIFactory.create_provider()
    response = ai.generate_structured(
        "Compose a team for this project", {"team": "array"}
    )
    assert "team" in response
    assert len(response["team"]) >= 2

class SmartMockProvider(MockAIProvider):
    def __init__(self):
        super().__init__()
        self.response_history = []

    def generate_response(self, prompt: str, model: str = "gpt-4") -> str:
        response = super().generate_response(prompt, model)
        self.response_history.append((prompt, response))
        return response

    def export_real_responses(self):
        """ Dùng dữ liệu phản hồi thực để cải tiến mock """
        return self.response_history