import unittest

from agent_helpers import (
    build_user_content,
    classify_attachment,
    cleanup_exact_answer,
    is_youtube_question,
    is_youtube_visual_question,
)


class AgentHelperTests(unittest.TestCase):
    def test_build_user_content_exposes_task_id_to_tool_agent(self):
        content = build_user_content("What is in the attached file?", "abc-123")

        self.assertIn("What is in the attached file?", content)
        self.assertIn("abc-123", content)
        self.assertIn("download_task_file", content)

    def test_build_user_content_without_task_id_is_plain_question(self):
        self.assertEqual(build_user_content("What is 2 + 2?", None), "What is 2 + 2?")

    def test_classify_attachment_prefers_file_suffix(self):
        self.assertEqual(classify_attachment("What does it say?", ".mp3"), "audio")
        self.assertEqual(classify_attachment("Analyze the table", ".xlsx"), "spreadsheet")
        self.assertEqual(classify_attachment("What is shown?", ".png"), "image")
        self.assertEqual(classify_attachment("What is the output?", ".py"), "python")

    def test_classify_attachment_uses_question_when_suffix_is_missing(self):
        self.assertEqual(classify_attachment("Use the attached spreadsheet", ""), "spreadsheet")
        self.assertEqual(classify_attachment("What is said in the recording?", ""), "audio")
        self.assertEqual(classify_attachment("Review the chess position image", ""), "image")
        self.assertEqual(classify_attachment("Run the attached Python code", ""), "python")
        self.assertEqual(classify_attachment("Read the attached text file", ""), "text")

    def test_youtube_detection(self):
        self.assertTrue(is_youtube_question("Watch https://www.youtube.com/watch?v=L1vXCYZAYYM"))
        self.assertTrue(is_youtube_question("See https://youtu.be/L1vXCYZAYYM"))
        self.assertFalse(is_youtube_question("This mentions video but has no URL"))

    def test_youtube_visual_question_detection(self):
        self.assertTrue(is_youtube_visual_question("What is the highest number of bird species on camera? https://youtu.be/x"))
        self.assertFalse(is_youtube_visual_question("What does Teal'c say in response? https://youtu.be/x"))

    def test_cleanup_exact_answer_removes_common_wrappers(self):
        self.assertEqual(cleanup_exact_answer("FINAL ANSWER: 519"), "519")
        self.assertEqual(cleanup_exact_answer("`b, e`"), "b, e")
        self.assertEqual(cleanup_exact_answer("The answer is: Right."), "Right")


if __name__ == "__main__":
    unittest.main()