# pylint: disable=too-many-lines,line-too-long,useless-suppression
# ------------------------------------
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.
# ------------------------------------
# cSpell:disable

"""
Multi-Tool Tests: File Search + Function Tool

Tests various scenarios using an agent with File Search and Function Tool.
All tests use the same tool combination but different inputs and workflows.
"""

import json
from io import BytesIO
from test_base import TestBase, servicePreparer
from devtools_testutils import recorded_by_proxy, RecordedTransport
from azure.ai.projects.models import PromptAgentDefinition, FileSearchTool, FunctionTool
from openai.types.responses.response_input_param import FunctionCallOutput, ResponseInputParam


class TestAgentFileSearchAndFunction(TestBase):
    """Tests for agents using File Search + Function Tool combination."""

    @servicePreparer()
    @recorded_by_proxy(RecordedTransport.AZURE_CORE, RecordedTransport.HTTPX)
    def test_data_analysis_workflow(self, **kwargs):
        """
        Test data analysis workflow: upload data, search, save results.
        """

        model = kwargs.get("azure_ai_model_deployment_name")

        # Setup
        project_client = self.create_client(operation_group="agents", **kwargs)
        openai_client = project_client.get_openai_client()

        # Create sample data file
        txt_content = """Sales Data Q1-Q3:

Product: Widget A
Q1: $15,000
Q2: $18,000
Q3: $21,000
Total: $54,000

Product: Widget B
Q1: $22,000
Q2: $25,000
Q3: $28,000
Total: $75,000

Overall Total Revenue: $129,000
"""

        # Create vector store and upload
        vector_store = openai_client.vector_stores.create(name="SalesDataStore")
        print(f"Vector store created (id: {vector_store.id})")

        txt_file = BytesIO(txt_content.encode("utf-8"))
        txt_file.name = "sales_data.txt"

        file = openai_client.vector_stores.files.upload_and_poll(
            vector_store_id=vector_store.id,
            file=txt_file,
        )
        print(f"File uploaded (id: {file.id})")
        assert file.status == "completed"

        # Define function tool
        save_results_function = FunctionTool(
            name="save_analysis_results",
            description="Save the analysis results",
            parameters={
                "type": "object",
                "properties": {
                    "summary": {
                        "type": "string",
                        "description": "Summary of the analysis",
                    },
                },
                "required": ["summary"],
                "additionalProperties": False,
            },
            strict=True,
        )

        # Create agent
        agent = project_client.agents.create_version(
            agent_name="file-search-function-agent",
            definition=PromptAgentDefinition(
                model=model,
                instructions="You are a data analyst. Use file search to find data and save_analysis_results function to save your findings.",
                tools=[
                    FileSearchTool(vector_store_ids=[vector_store.id]),
                    save_results_function,
                ],
            ),
            description="Agent with File Search and Function Tool.",
        )
        print(f"Agent created (id: {agent.id}, name: {agent.name}, version: {agent.version})")

        # Request analysis
        print("\nAsking agent to analyze the sales data...")

        response = openai_client.responses.create(
            input="Analyze the sales data and calculate the total revenue for each product. Then save the results.",
            extra_body={"agent": {"name": agent.name, "type": "agent_reference"}},
        )

        print(f"Initial response completed (id: {response.id})")

        # Check if function was called
        function_calls_found = 0
        input_list: ResponseInputParam = []

        for item in response.output:
            if item.type == "function_call":
                function_calls_found += 1
                print(f"Function call detected (id: {item.call_id}, name: {item.name})")

                assert item.name == "save_analysis_results"

                arguments = json.loads(item.arguments)
                print(f"Function arguments: {arguments}")

                assert "summary" in arguments
                assert len(arguments["summary"]) > 20

                input_list.append(
                    FunctionCallOutput(
                        type="function_call_output",
                        call_id=item.call_id,
                        output=json.dumps({"status": "success", "saved": True}),
                    )
                )

        assert function_calls_found > 0, "Expected save_analysis_results function to be called"

        # Send function results back
        if input_list:
            response = openai_client.responses.create(
                input=input_list,
                previous_response_id=response.id,
                extra_body={"agent": {"name": agent.name, "type": "agent_reference"}},
            )
            print(f"Final response: {response.output_text[:200]}...")

        print("\n[PASS] Workflow completed successfully")

        # Teardown
        project_client.agents.delete_version(agent_name=agent.name, agent_version=agent.version)
        openai_client.vector_stores.delete(vector_store.id)
        print("Cleanup completed")

    @servicePreparer()
    @recorded_by_proxy(RecordedTransport.AZURE_CORE, RecordedTransport.HTTPX)
    def test_empty_vector_store_handling(self, **kwargs):
        """
        Test how agent handles empty vector store (no files uploaded).
        """

        model = kwargs.get("azure_ai_model_deployment_name")

        # Setup
        project_client = self.create_client(operation_group="agents", **kwargs)
        openai_client = project_client.get_openai_client()

        # Create EMPTY vector store
        vector_store = openai_client.vector_stores.create(name="EmptyStore")
        print(f"Empty vector store created (id: {vector_store.id})")

        # Define function tool
        error_function = FunctionTool(
            name="report_error",
            description="Report when data is not available",
            parameters={
                "type": "object",
                "properties": {
                    "error_message": {
                        "type": "string",
                        "description": "Description of the error",
                    },
                },
                "required": ["error_message"],
                "additionalProperties": False,
            },
            strict=True,
        )

        # Create agent
        agent = project_client.agents.create_version(
            agent_name="file-search-function-error-agent",
            definition=PromptAgentDefinition(
                model=model,
                instructions="Search for data. If you can't find it, use report_error function to explain what happened.",
                tools=[
                    FileSearchTool(vector_store_ids=[vector_store.id]),
                    error_function,
                ],
            ),
            description="Agent for testing error handling.",
        )
        print(f"Agent created (id: {agent.id})")

        # Request analysis of non-existent file
        print("\nAsking agent to find non-existent data...")

        response = openai_client.responses.create(
            input="Find and analyze the quarterly sales report.",
            extra_body={"agent": {"name": agent.name, "type": "agent_reference"}},
        )

        response_text = response.output_text
        print(f"Response: '{response_text[:200] if response_text else '(empty)'}...'")

        # Verify agent didn't crash
        self.validate_response(response)

        # If there's text, it should be meaningful
        if response_text:
            assert len(response_text) > 10, "Non-empty response should be meaningful"

        print("\n[PASS] Agent handled missing data gracefully")

        # Teardown
        project_client.agents.delete_version(agent_name=agent.name, agent_version=agent.version)
        openai_client.vector_stores.delete(vector_store.id)
        print("Cleanup completed")

    @servicePreparer()
    @recorded_by_proxy(RecordedTransport.AZURE_CORE, RecordedTransport.HTTPX)
    def test_python_code_file_search(self, **kwargs):
        """
        Test searching for Python code files and saving findings.

        This test verifies that both File Search and Function Tool are used:
        1. File Search: Agent searches vector store for Python code
        2. Function Tool: Agent saves the code review findings
        """

        model = kwargs.get("azure_ai_model_deployment_name")

        # Setup
        project_client = self.create_client(operation_group="agents", **kwargs)
        openai_client = project_client.get_openai_client()

        # Create sample Python code
        python_code = """# Sample Python code for analysis
def calculate_sum(numbers):
    return sum(numbers)

# Test data
data = [1, 2, 3, 4, 5]
result = calculate_sum(data)
print(f"Sum: {result}")
"""

        # Create vector store and upload
        vector_store = openai_client.vector_stores.create(name="CodeStore")

        code_file = BytesIO(python_code.encode("utf-8"))
        code_file.name = "sample_code.py"

        file = openai_client.vector_stores.files.upload_and_poll(
            vector_store_id=vector_store.id,
            file=code_file,
        )
        print(f"Python file uploaded (id: {file.id})")

        # Define function tool
        save_function = FunctionTool(
            name="save_code_review",
            description="Save code review findings. Must be called to persist the review.",
            parameters={
                "type": "object",
                "properties": {
                    "findings": {
                        "type": "string",
                        "description": "Code review findings",
                    },
                },
                "required": ["findings"],
                "additionalProperties": False,
            },
            strict=True,
        )

        # Create agent with explicit instructions to use both tools
        agent = project_client.agents.create_version(
            agent_name="file-search-function-code-agent",
            definition=PromptAgentDefinition(
                model=model,
                instructions="You are a code reviewer. Search for code files, analyze them, and ALWAYS save your findings using the save_code_review function.",
                tools=[
                    FileSearchTool(vector_store_ids=[vector_store.id]),
                    save_function,
                ],
            ),
            description="Agent for testing file search with code.",
        )
        print(f"Agent created (id: {agent.id})")

        # Request code analysis with explicit save instruction
        print("\nAsking agent to find, analyze, and save the code review...")

        response = openai_client.responses.create(
            input="Find the Python code file, explain what the calculate_sum function does, and save your code review findings.",
            extra_body={"agent": {"name": agent.name, "type": "agent_reference"}},
        )

        # Verify function call is made (both tools should be used)
        function_call_found = False
        input_list: ResponseInputParam = []

        for item in response.output:
            if item.type == "function_call":
                function_call_found = True
                print(f"Function call detected (name: {item.name})")
                assert item.name == "save_code_review", f"Expected save_code_review, got {item.name}"

                arguments = json.loads(item.arguments)
                print(f"Function arguments: {arguments}")
                assert "findings" in arguments
                assert len(arguments["findings"]) > 20, "Expected meaningful findings"

                # Verify findings discuss the code (proves File Search was used)
                findings_lower = arguments["findings"].lower()
                assert any(
                    keyword in findings_lower
                    for keyword in ["sum", "calculate", "function", "numbers", "list", "return"]
                ), f"Expected findings to discuss the code content. Got: {arguments['findings'][:100]}"

                input_list.append(
                    FunctionCallOutput(
                        type="function_call_output",
                        call_id=item.call_id,
                        output=json.dumps({"status": "saved", "review_id": "review_001"}),
                    )
                )

        assert function_call_found, "Expected save_code_review function to be called"

        # Send function results back to get final response
        response = openai_client.responses.create(
            input=input_list,
            previous_response_id=response.id,
            extra_body={"agent": {"name": agent.name, "type": "agent_reference"}},
        )
        response_text = response.output_text
        print(f"Final response: {response_text[:200] if response_text else '(confirmation)'}...")

        print("\n[PASS] Agent successfully used both File Search and Function Tool")

        # Teardown
        project_client.agents.delete_version(agent_name=agent.name, agent_version=agent.version)
        openai_client.vector_stores.delete(vector_store.id)
        print("Cleanup completed")

    @servicePreparer()
    @recorded_by_proxy(RecordedTransport.AZURE_CORE, RecordedTransport.HTTPX)
    def test_multi_turn_search_and_save_workflow(self, **kwargs):
        """
        Test multi-turn workflow: search documents, ask follow-ups, save findings.

        This tests:
        - File Search across multiple turns
        - Function calls interspersed with searches
        - Context retention across searches and function calls
        """

        model = kwargs.get("azure_ai_model_deployment_name")

        # Setup
        project_client = self.create_client(operation_group="agents", **kwargs)
        openai_client = project_client.get_openai_client()

        # Create research documents
        doc1_content = """Research Paper: Machine Learning in Healthcare

Abstract:
Machine learning algorithms are revolutionizing healthcare diagnostics.
Recent studies show 95% accuracy in disease detection using neural networks.
Key applications include medical imaging, patient risk prediction, and drug discovery.

Conclusion:
ML techniques offer promising solutions for early disease detection and personalized treatment.
"""

        doc2_content = """Research Paper: AI Ethics and Governance

Abstract:
As AI systems become more prevalent, ethical considerations are paramount.
Issues include bias in algorithms, data privacy, and accountability.
Regulatory frameworks are being developed globally to address these concerns.

Conclusion:
Responsible AI development requires multistakeholder collaboration and transparent governance.
"""

        # Create vector store and upload documents
        vector_store = openai_client.vector_stores.create(name="ResearchStore")
        print(f"Vector store created: {vector_store.id}")

        file1 = BytesIO(doc1_content.encode("utf-8"))
        file1.name = "ml_healthcare.txt"
        file2 = BytesIO(doc2_content.encode("utf-8"))
        file2.name = "ai_ethics.txt"

        uploaded1 = openai_client.vector_stores.files.upload_and_poll(
            vector_store_id=vector_store.id,
            file=file1,
        )
        uploaded2 = openai_client.vector_stores.files.upload_and_poll(
            vector_store_id=vector_store.id,
            file=file2,
        )
        print(f"Documents uploaded: {uploaded1.id}, {uploaded2.id}")

        # Define save function
        save_finding = FunctionTool(
            name="save_finding",
            description="Save research finding",
            parameters={
                "type": "object",
                "properties": {
                    "topic": {"type": "string", "description": "Research topic"},
                    "finding": {"type": "string", "description": "Key finding"},
                },
                "required": ["topic", "finding"],
                "additionalProperties": False,
            },
            strict=True,
        )

        # Create agent
        agent = project_client.agents.create_version(
            agent_name="research-assistant-multi-turn",
            definition=PromptAgentDefinition(
                model=model,
                instructions="You are a research assistant. Search documents and save important findings.",
                tools=[
                    FileSearchTool(vector_store_ids=[vector_store.id]),
                    save_finding,
                ],
            ),
            description="Research assistant for multi-turn testing.",
        )
        print(f"Agent created: {agent.id}")

        # Turn 1: Search for ML in healthcare
        print("\n--- Turn 1: Initial search query ---")
        response_1 = openai_client.responses.create(
            input="What does the research say about machine learning in healthcare?",
            extra_body={"agent": {"name": agent.name, "type": "agent_reference"}},
        )

        response_1_text = response_1.output_text
        print(f"Response 1: {response_1_text[:200]}...")
        assert "95" in response_1_text or "accuracy" in response_1_text.lower()

        # Turn 2: Follow-up for specifics
        print("\n--- Turn 2: Follow-up for details ---")
        response_2 = openai_client.responses.create(
            input="What specific applications are mentioned?",
            previous_response_id=response_1.id,
            extra_body={"agent": {"name": agent.name, "type": "agent_reference"}},
        )

        response_2_text = response_2.output_text
        print(f"Response 2: {response_2_text[:200]}...")
        response_2_lower = response_2_text.lower()
        assert any(keyword in response_2_lower for keyword in ["imaging", "drug", "risk", "prediction"])

        # Turn 3: Save the finding
        print("\n--- Turn 3: Save finding ---")
        response_3 = openai_client.responses.create(
            input="Please save that finding about ML accuracy.",
            previous_response_id=response_2.id,
            extra_body={"agent": {"name": agent.name, "type": "agent_reference"}},
        )

        # Handle function call
        input_list: ResponseInputParam = []
        function_called = False
        for item in response_3.output:
            if item.type == "function_call":
                function_called = True
                print(f"Function called: {item.name} with args: {item.arguments}")
                assert item.name == "save_finding"

                args = json.loads(item.arguments)
                assert "topic" in args and "finding" in args
                print(f"  Topic: {args['topic']}")
                print(f"  Finding: {args['finding'][:100]}...")

                input_list.append(
                    FunctionCallOutput(
                        type="function_call_output",
                        call_id=item.call_id,
                        output=json.dumps({"status": "saved", "id": "finding_001"}),
                    )
                )

        assert function_called, "Expected save_finding to be called"

        # Send function result
        response_3 = openai_client.responses.create(
            input=input_list,
            previous_response_id=response_3.id,
            extra_body={"agent": {"name": agent.name, "type": "agent_reference"}},
        )
        print(f"Response 3: {response_3.output_text[:150]}...")

        # Turn 4: Switch to different topic (AI ethics)
        print("\n--- Turn 4: New search topic ---")
        response_4 = openai_client.responses.create(
            input="Now tell me about AI ethics concerns mentioned in the research.",
            previous_response_id=response_3.id,
            extra_body={"agent": {"name": agent.name, "type": "agent_reference"}},
        )

        response_4_text = response_4.output_text
        print(f"Response 4: {response_4_text[:200]}...")
        response_4_lower = response_4_text.lower()
        assert any(keyword in response_4_lower for keyword in ["bias", "privacy", "ethics", "accountability"])

        print("\n[PASS] Multi-turn File Search + Function workflow successful!")
        print("  - Multiple searches across different documents")
        print("  - Function called after context-building searches")
        print("  - Topic switching works correctly")

        # Cleanup
        project_client.agents.delete_version(agent_name=agent.name, agent_version=agent.version)
        openai_client.vector_stores.delete(vector_store.id)
        print("Cleanup completed")
