from unittest import TestCase
import unittest
import uuid
from pygeai.lab.managers import AILabManager
from pygeai.lab.models import Agent, AgentData, Prompt, LlmConfig, Model, Sampling, PromptExample, PromptOutput
from pydantic import ValidationError


class TestAILabCreateAgentIntegration(TestCase):    
    def setUp(self):
        """
        Set up the test environment.
        """
        self.ai_lab_manager = AILabManager(alias="beta")
        self.project_id = "be4889df-cacc-4e6f-b3bb-153c4ac0d168"
        self.new_agent = self.__load_agent()
        self.created_agent: Agent = None


    def tearDown(self):
        """
        Clean up after each test if necessary.
        This can be used to delete the created agent or reset the state.
        """
        if isinstance(self.created_agent, Agent):
            self.ai_lab_manager.delete_agent(self.project_id, self.created_agent.id)


    def __load_agent(self):
        random_str = str(uuid.uuid4())
        return Agent(
            name=random_str,
            access_scope="public",
            public_name=f"public_{random_str}",
            job_description="Translator",
            avatar_image="https://www.shareicon.net/data/128x128/2016/11/09/851442_logo_512x512.png",
            description="Agent that translates from any language to english.",
            agent_data=AgentData(
                prompt=Prompt(
                    instructions="the user will provide a text, you must return the same text translated to english",
                    inputs=["text", "avoid slang indicator"],
                    outputs=[
                        PromptOutput(key="translated_text", description="translated text, with slang or not depending on the indication. in plain text.")
                    ],
                    examples=[
                        PromptExample(input_data="hola mundo [no-slang]", output='{"translated_text":"hello world"}'),
                        PromptExample(input_data="esto es una prueba pincheguey [keep-slang]", output='{"translated_text":"this is a test pal"}')
                    ]
                ),
                llm_config=LlmConfig(
                    max_tokens=1800,
                    timeout=0,
                    sampling=Sampling(temperature=0.3, top_k=0, top_p=0)
                ),
                models=[Model(name="gpt-4-turbo-preview")]
            )
        )


    def __create_agent(self, agent=None, automatic_publish=False):
        """
        Helper to create an agent with the current project_id and ai_lab_manager.
        If automatic_publish is None, do not pass it (useful for tests that omit it).
        """
        return self.ai_lab_manager.create_agent(
            project_id=self.project_id,
            agent=self.new_agent if agent is None else agent,
            automatic_publish=automatic_publish
        )


    def test_create_agent_full_data(self):
        self.created_agent = self.__create_agent()
        created_agent = self.created_agent
        agent = self.new_agent

        self.assertTrue(isinstance(created_agent, Agent), "Expected a created agent")

        # Assert the main fields of the created agent
        self.assertIsNotNone(created_agent.id)
        self.assertEqual(created_agent.name, agent.name)
        self.assertEqual(created_agent.access_scope, agent.access_scope)
        self.assertEqual(created_agent.public_name, agent.public_name)
        self.assertEqual(created_agent.avatar_image, agent.avatar_image)
        self.assertEqual(created_agent.description, agent.description)
        self.assertEqual(created_agent.job_description, agent.job_description)

        # Assert agentData fields
        agent_data = created_agent.agent_data
        self.assertIsNotNone(agent_data)
        self.assertEqual(agent_data.llm_config.max_tokens, agent.agent_data.llm_config.max_tokens)
        self.assertEqual(agent_data.llm_config.timeout, agent.agent_data.llm_config.timeout)
        self.assertEqual(agent_data.llm_config.sampling.temperature, agent.agent_data.llm_config.sampling.temperature)
        self.assertEqual(agent_data.llm_config.sampling.top_k, agent.agent_data.llm_config.sampling.top_k)
        self.assertEqual(agent_data.llm_config.sampling.top_p, agent.agent_data.llm_config.sampling.top_p)
        self.assertEqual(agent_data.models[0].name, agent.agent_data.models[0].name)

        # Assert prompt fields
        prompt = agent_data.prompt
        self.assertEqual(prompt.instructions, agent.agent_data.prompt.instructions)
        self.assertEqual(prompt.inputs, agent.agent_data.prompt.inputs)
        self.assertEqual(prompt.outputs[0].key, agent.agent_data.prompt.outputs[0].key)
        self.assertEqual(prompt.outputs[0].description, agent.agent_data.prompt.outputs[0].description)

        # Assert prompt examples
        self.assertEqual(prompt.examples[0].input_data, agent.agent_data.prompt.examples[0].input_data)
        self.assertEqual(prompt.examples[0].output, agent.agent_data.prompt.examples[0].output)


    def test_create_agent_minimum_required_data(self):
        self.new_agent = Agent(
            name=str(uuid.uuid4()),
            description="Agent that translates from any language to english.",
            agent_data=AgentData(
                prompt=Prompt(
                    instructions="the user will provide a text, you must return the same text translated to english",
                    inputs=["text", "avoid slang indicator"]
                ),
                llm_config=LlmConfig(
                    max_tokens=1800,
                    timeout=0
                ),
                models=[Model(name="gpt-4-turbo-preview")]
            )
        )
        self.created_agent = self.__create_agent()
        agent = self.new_agent

        self.assertIsNotNone(self.created_agent.id)
        self.assertEqual(self.created_agent.name, agent.name)
        self.assertEqual(self.created_agent.description, agent.description)

         # Assert agentData fields
        agent_data = self.created_agent.agent_data
        self.assertIsNotNone(agent_data)
        self.assertEqual(agent_data.llm_config.max_tokens, agent.agent_data.llm_config.max_tokens)
        self.assertEqual(agent_data.llm_config.timeout, agent.agent_data.llm_config.timeout)
        self.assertEqual(agent_data.models[0].name, agent.agent_data.models[0].name)

        # Assert prompt fields
        prompt = agent_data.prompt
        self.assertEqual(prompt.instructions, agent.agent_data.prompt.instructions)
        self.assertEqual(prompt.inputs, agent.agent_data.prompt.inputs)


    def test_create_agent_without_required_data(self):
        test_params = [ True, False ]

        for auto_publish in test_params:
            with self.subTest(input=auto_publish):
                self.new_agent = Agent(
                    name=str(uuid.uuid4())
                )
                self.created_agent = self.__create_agent(automatic_publish=auto_publish)
        
                #TODO: Change validation error to a more specific one
                self.assertEqual(self.created_agent.id, 400)
                self.assertEqual(self.created_agent.description, "Bad Request")


    def test_create_agent_no_name(self):
        test_params = [ True, False ]

        for auto_publish in test_params:
            with self.subTest(input=auto_publish):
                self.new_agent.name = ""
                self.created_agent = self.__create_agent(automatic_publish=auto_publish)

                self.assertEqual(
                    self.created_agent.errors[0].description,
                    "Agent name cannot be empty.",
                    f"Expected an error about the missing agent name with autopublish {'enabled' if auto_publish else 'disabled'}"
                )


    def test_create_agent_duplicated_name(self):
        test_params = [ True, False ]

        for auto_publish in test_params:

            with self.subTest(input=auto_publish):
                self.new_agent.name = "AritmeticaExpert"

                self.created_agent = self.__create_agent(automatic_publish=auto_publish)

                self.assertEqual(
                    self.created_agent.errors[0].description,
                    "Agent cannot be created as it already exists [name=AritmeticaExpert].",
                    f"Expected an error about duplicated agent name with autopublish {'enabled' if auto_publish else 'disabled'}"
                )


    def test_create_agent_invalid_name(self):
        test_params = [ True, False ]

        for auto_publish in test_params:

            with self.subTest(input=auto_publish):
                new_agent = self.__load_agent()
                new_agent2 = self.__load_agent()
        
                new_agent.name = f"{new_agent.name}:invalid"
                result = self.__create_agent(agent=new_agent, automatic_publish=auto_publish)

                new_agent2.name = f"{new_agent2.name}/invalid"
                result2 = self.__create_agent(agent=new_agent2, automatic_publish=auto_publish)

                self.assertEqual(
                    result.errors[0].description,
                    "Invalid character in name (: is not allowed).",
                    f"Expected an error about invalid character (:) in agent name with autopublish {'enabled' if auto_publish else 'disabled'}"
                )

                self.assertEqual(
                    result2.errors[0].description,
                    "Invalid character in name (/ is not allowed).",
                    f"Expected an error about invalid character (/) in agent name with autopublish {'enabled' if auto_publish else 'disabled'}"
                )
        

    def test_create_agent_invalid_scope(self):
        self.new_agent.access_scope = "project" 
        with self.assertRaises(ValueError) as exc:
            self.__create_agent()
            self.assertEqual(
                str(exc.exception),
                "Access scope must be one of public, private",
                "Expected a ValueError exception for invalid access scope"
            )


    def test_create_agent_default_scope(self):
        self.new_agent.access_scope = None
        self.created_agent = self.__create_agent()

        self.assertEqual(self.created_agent.access_scope, "private", "Expected the default access scope to be 'private' when not specified")


    def test_create_agent_no_public_name(self):
        test_params = [ True, False ]

        for auto_publish in test_params:

            with self.subTest(input=auto_publish):
                self.new_agent.public_name = None
                self.created_agent = self.__create_agent(automatic_publish=auto_publish)

                self.assertEqual(
                    self.created_agent.errors[0].description,
                    "Agent publicName is required for agents with accessScope=public.",
                    f"Expected an error about missing publicName for public access scope with autopublish {'enabled' if auto_publish else 'disabled'}"
                )

    
    def test_create_agent_invalid_public_name(self):
        test_params = [ True, False ]

        for auto_publish in test_params:
            with self.subTest(input=auto_publish): 
                self.new_agent.public_name = self.new_agent.public_name.replace("_", "#")  # Add invalid character to public name
                self.created_agent = self.__create_agent(automatic_publish=auto_publish)

                self.assertEqual(
                    self.created_agent.errors[0].description,
                    "Invalid public name, it can only contain lowercase letters, numbers, periods (.), dashes (-), and underscores (_). Please remove any other characters.",
                    f"The expected error about invalid publicName was not returned when autopublish is {'enabled' if auto_publish else 'disabled'}"
                )

    
    def test_create_agent_duplicated_public_name(self):
        test_params = [ True, False ]
        self.created_agent = self.__create_agent()

        for auto_publish in test_params:
            with self.subTest(input=auto_publish):
                duplicated_pn_agent = self.__load_agent()
                duplicated_pn_agent.public_name = self.created_agent.public_name
                result = self.__create_agent(agent=duplicated_pn_agent, automatic_publish=auto_publish)

                self.assertEqual(
                    result.errors[0].description,
                    f"Agent already exists [publicName={self.created_agent.public_name}].",
                    f"Expected an error about the duplicated public name when autopublish is {'enabled' if auto_publish else 'disabled'}"
                )

    
    @unittest.skip("Agent is getting created regardless of the prompt instructions being empty")
    def test_create_agent_no_prompt_instructions(self):
        self.new_agent.agent_data.prompt.instructions = ""
        with self.assertRaises(ValidationError) as exception:
            self.__create_agent()
        
            self.assertIn(
                """instructions
                    Input should be a valid string [type=string_type, input_value=None, input_type=NoneType]""",
                str(exception.exception),
                "Expected a validation error about allowed values for instructions"
            )

    
    def test_create_agent_autopublish(self): 
        self.created_agent = self.__create_agent(automatic_publish=True)
        self.assertFalse(self.created_agent.is_draft, "Expected the agent to be published automatically")
    
    
    def test_create_agent_autopublish_private_scope(self):
        self.new_agent.access_scope = "private"

        self.created_agent = self.__create_agent(automatic_publish=True)
        self.assertFalse(self.created_agent.is_draft, "Expected the agent to be published automatically even with private scope")