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 TestAILabUpdateAgentIntegration(TestCase):    
    def setUp(self):
        """
        Set up the test environment.
        """
        self.ai_lab_manager = AILabManager(alias="beta")
        self.project_id = "be4889df-cacc-4e6f-b3bb-153c4ac0d168"

        load_agent = self.__load_agent()
        self.agent_to_update = load_agent["agent"]
        self.random_str = load_agent["random_str"]

    def __load_agent(self):
        random_str = str(uuid.uuid4())
        agent = Agent(
            id="b4b09935-2ad2-42c0-bd55-1ee6fa4b6034",
            name=f"UpdatedAgent{random_str}",
            access_scope="public",
            public_name=f"public_{random_str}",
            job_description=f"SummarizerAgent{random_str}",
            avatar_image="https://www.shareicon.net/data/128x128/2016/11/09/851442_logo_512x512.png",
            description=f"Agent that summarized documents. {random_str}",
            agent_data=AgentData(
                prompt=Prompt(
                    instructions="the user will provide a document, you must return a summary of the document.",
                    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="openai/gpt-4-turbo-preview")]
            )
        )

        return { 
            "agent": agent,
            "random_str": random_str 
        }


    def __update_agent(self, agent=None, automatic_publish=False, upsert=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.update_agent(
            project_id=self.project_id,
            agent=self.agent_to_update if agent is None else agent,
            automatic_publish=automatic_publish, 
            upsert=upsert
        )
    

    def test_update_agent(self):
        self.agent_to_update.agent_data.models[0].name = "gemini/gemini-1.5-flash-8b-exp-0827"
        self.agent_to_update.access_scope = "private"
        updated_agent = self.__update_agent()

        self.assertTrue(isinstance(updated_agent, Agent), "Expected a created agent")
       
        self.assertEqual(updated_agent.name, f"UpdatedAgent{self.random_str}")
        self.assertEqual(updated_agent.access_scope, "private")
        self.assertEqual(updated_agent.public_name, f"public_{self.random_str}")
        self.assertEqual(updated_agent.job_description, f"SummarizerAgent{self.random_str}",)
        self.assertEqual(updated_agent.description, f"Agent that summarized documents. {self.random_str}")

        self.assertEqual(updated_agent.agent_data.models[0].name, "gemini/gemini-1.5-flash-8b-exp-0827")
        self.assertTrue(updated_agent.is_draft, "gemini/gemini-1.5-flash-8b-exp-0827")

        self.agent_to_update.agent_data.models[0].name = "openai/gpt-4-turbo-preview"
        updated_agent = self.__update_agent()

    
    def test_update_agent_no_name(self):
        test_params = [ True, False ]
        self.agent_to_update.name = ""
        for auto_publish in test_params:
            with self.subTest(input=auto_publish):
        
                updated_agent = self.__update_agent(automatic_publish=auto_publish)
                self.assertEqual(
                    updated_agent.errors[0].description,
                    "Agent name cannot be empty.",
                    f"Expected error about missing agent name when autopublish is {'enabled' if auto_publish else 'disabled'} was not returned."
                )


    def test_update_agent_duplicated_name(self):
        test_params = [ True, False ]
        self.agent_to_update.name = "AritmeticaExpert"

        for auto_publish in test_params:
            with self.subTest(input=auto_publish):
                updated_agent = self.__update_agent(automatic_publish=auto_publish)
                self.assertEqual(
                    updated_agent.errors[0].description,
                    "Agent already exists [name=AritmeticaExpert]",
                    f"Expected error about duplicated agent name when autopublish is {'enabled' if auto_publish else 'disabled'} was not returned."
                )

    
    def test_update_agent_invalidName_name(self):
        test_params = [ True, False ]
        self.agent_to_update.name = f"{self.agent_to_update.name}:/"

        for auto_publish in test_params:
            with self.subTest(input=auto_publish):
        
                updated_agent = self.__update_agent(automatic_publish=auto_publish)

                self.assertEqual(
                    updated_agent.errors[0].description,
                    "Invalid character in name (: is not allowed).",
                    f"Expected error about invalid character (:) in name when autopublish is {'enabled' if auto_publish else 'disabled'} was not returned."
                )
                self.assertEqual(
                    updated_agent.errors[1].description, 
                    "Invalid character in name (/ is not allowed).",
                    f"Expected error about invalid character (/) in name when autopublish is {'enabled' if auto_publish else 'disabled'} was not returned."
                )


    def test_update_agent_no_public_name(self):        
        test_params = [ True, False ]
        self.agent_to_update.public_name = ""

        for auto_publish in test_params:
            with self.subTest(input=auto_publish):
                updated_agent = self.__update_agent(automatic_publish=auto_publish)                
                self.assertEqual(
                    updated_agent.errors[0].description, 
                    "Agent publicName is required for agents with accessScope=public.",
                    f"Expected error about missing public name when autopublish is {'enabled' if auto_publish else 'disabled'} was not returned."
                )


    def test_update_agent_duplicated_public_name(self):               
        test_params = [ True, False ]
        self.agent_to_update.public_name = "com.testing.geai.googlesummarizer"

        for auto_publish in test_params:
            with self.subTest(input=auto_publish):                
                updated_agent = self.__update_agent(automatic_publish=auto_publish)

                self.assertEqual(
                    updated_agent.errors[0].description, 
                    "Agent already exists [publicName=com.testing.geai.googlesummarizer].",
                    f"Expected error about duplicated public name when autopublish is {'enabled' if auto_publish else 'disabled'} was not returned."
                )

    
    def test_update_agent_invalid_public_name(self):        
        test_params = [ True, False ]
        self.agent_to_update.public_name = self.agent_to_update.public_name.replace("_", "/")

        for auto_publish in test_params:
            with self.subTest(input=auto_publish): 
                updated_agent = self.__update_agent(automatic_publish=auto_publish)
                self.assertEqual(
                    updated_agent.errors[0].description, 
                    "Invalid public name, it can only contain lowercase letters, numbers, periods (.), dashes (-), and underscores (_). Please remove any other characters.",
                    f"Expected error about invalid public name when autopublish is {'enabled' if auto_publish else 'disabled'} was not returned."
                )


    def test_update_agent_no_prompt_instructions(self):
        test_params = [ True, False ]
        self.agent_to_update.agent_data.prompt.instructions = ""

        for auto_publish in test_params:
            with self.subTest(input=auto_publish): 
                with self.assertRaises(ValidationError) as exception:
                    self.__update_agent(automatic_publish=auto_publish)
                
                    self.assertIn(
                        """instructions
                            Input should be a valid string [type=string_type, input_value=None, input_type=NoneType]""",
                        str(exception.exception),
                        f"Expected a validation error about allowed values for instructions when autopublish is {'enabled' if auto_publish else 'disabled'}"
                    )

    
    @unittest.skip("Mismatched error message")
    def test_update_agent_no_model(self):        
        test_params = [ True, False ]
        self.agent_to_update.agent_data.models[0].name = ""
        for auto_publish in test_params:
            with self.subTest(input=auto_publish): 
                     
                with self.assertRaises(ValidationError) as exception:
                    self.__update_agent(automatic_publish=auto_publish)
                
                    self.assertAlmostEqual(
                        """name
                           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 model name" 
                    )

    def test_update_agent_autopublish(self):        
        updated_agent = self.__update_agent(automatic_publish=True)

        self.assertEqual(
            updated_agent.name, 
            f"UpdatedAgent{self.random_str}",
            "Expected agent name to remain unchanged after update."
        )
        self.assertFalse(
            updated_agent.is_draft,
            "Expected agent to be published when autopublish is True, but it is still a draft."
        )


    def test_update_agent_autopublish_private_scope(self):
        self.agent_to_update.access_scope = "private"

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