"""
Integration tests for agent CoT rich schema and tool-calling system.

This module tests the complete integration between:
- RichAgentCoT schema
- Tool-calling formatter
- Generator with agent_cot_rich type
"""

from unittest.mock import MagicMock, patch  # noqa: F401

import pytest  # type: ignore # noqa: F401

from deepfabric.config import DeepFabricConfig
from deepfabric.formatters.builtin.tool_calling import ToolCallingFormatter
from deepfabric.generator import DataSetGenerator
from deepfabric.schemas import get_conversation_schema


class TestAgentCoTIntegration:
    """Test integration between agent CoT schema and the generator."""

    def test_agent_cot_tools_schema_available(self):
        """Test that agent_cot_tools schema is available."""
        # Should be able to get the schema
        schema = get_conversation_schema("agent_cot_tools")
        assert schema is not None

        # Should be able to create an instance with agent data
        test_data = {
            "question": "Test question",
            "initial_analysis": "Test analysis",
            "reasoning_steps": ["Step 1", "Step 2"],
            "tool_selection_rationale": "Tool was chosen because...",
            "parameter_reasoning": "Parameters determined by...",
            "result_interpretation": "Result means...",
            "tool_used": "test_tool",
            "tool_input": '{"param": "value"}',
            "tool_output": "Test result",
            "answer": "Test answer",
        }

        instance = schema(**test_data)
        assert instance.question == "Test question"  # type: ignore

    @patch("deepfabric.generator.LLMClient")
    def test_generator_with_agent_cot_tools(self, mock_llm_client):  # noqa: ARG002
        """Test DataSetGenerator with agent_cot_tools conversation type."""
        # Create generator with agent_cot_tools configuration
        generator_config = {
            "generation_system_prompt": "Generate agent reasoning examples",
            "provider": "openai",
            "model_name": "gpt-4",
            "conversation_type": "agent_cot_tools",
            "available_tools": ["get_weather", "calculate"],
            "num_samples": 1,
            "batch_size": 1,
            "temperature": 0.7,
        }

        generator = DataSetGenerator(**generator_config)

        # Should have tool registry initialized
        assert generator.tool_registry is not None
        tool_names = generator.tool_registry.get_tool_names()
        assert "get_weather" in tool_names
        assert "calculate" in tool_names

    def test_tool_calling_formatter_with_rich_sample(self):
        """Test tool calling formatter with a rich agent CoT sample."""
        formatter = ToolCallingFormatter()

        # Rich agent sample (would be generated by the generator)
        rich_sample = {
            "question": "Calculate the area of a circle with radius 5",
            "answer": "The area of a circle with radius 5 is approximately 78.54 square units.",
            "initial_analysis": "This is a geometry problem requiring calculation of circle area using the formula π × r².",
            "reasoning_steps": [
                "Identify that this is a circle area calculation",
                "Recall the formula: Area = π × radius²",
                "Substitute radius = 5 into the formula",
                "Calculate π × 5² = π × 25 ≈ 78.54",
            ],
            "tool_selection_rationale": "Calculator tool is needed for accurate mathematical computation with π.",
            "parameter_reasoning": "Use 'circle_area' operation with radius=5 for precise calculation.",
            "result_interpretation": "Calculator returned 78.54, which matches expected result for π × 25.",
            "tool_used": "calculate",
            "tool_input": '{"operation": "circle_area", "radius": 5}',
            "tool_output": "78.54",
        }

        result = formatter.format_with_metadata([rich_sample])
        assert len(result.samples) == 1

        formatted = result.samples[0]
        messages = formatted["messages"]

        # Should have proper message structure
        assert (
            len(messages) == 4  # noqa: PLR2004
        )  # user, assistant+thinking, tool, assistant+answer  # noqa: PLR2004

        # Check that all rich reasoning is preserved in thinking
        thinking_content = messages[1]["content"]
        assert "Analysis:" in thinking_content
        assert rich_sample["initial_analysis"] in thinking_content
        assert "Step-by-step reasoning:" in thinking_content
        assert "Tool selection:" in thinking_content
        assert "Parameters:" in thinking_content

        # Verify all reasoning steps are included
        for step in rich_sample["reasoning_steps"]:
            assert step in thinking_content

    def test_end_to_end_rich_agent_workflow(self):
        """Test complete workflow from generation to formatting."""
        # This would be a more comprehensive test with actual generation
        # For now, test the data flow

        # 1. Rich agent sample (as would be generated)
        generated_sample = {
            "question": "Find restaurants near Times Square",
            "answer": "I found 3 restaurants near Times Square: Joe's Pizza, The View, and Olive Garden.",
            "initial_analysis": "User wants restaurant recommendations for a specific location in NYC.",
            "reasoning_steps": [
                "Parse the location request for Times Square",
                "Use restaurant search tool for that area",
                "Filter results for quality and ratings",
            ],
            "tool_selection_rationale": "Restaurant search tool can find businesses by location.",
            "parameter_reasoning": "Use 'Times Square, NYC' as location with restaurant category filter.",
            "result_interpretation": "Tool returned list of restaurants with ratings and locations.",
            "tool_used": "search_web",
            "tool_input": '{"location": "Times Square, NYC", "category": "restaurant"}',
            "tool_output": '[{"name": "Joe\'s Pizza", "rating": 4.5}, ...]',
        }

        # 2. Apply tool-calling formatter
        formatter = ToolCallingFormatter(
            {
                "include_tools_in_system": True,
                "system_prompt": "You are a helpful assistant with restaurant search capabilities.",
            }
        )

        result = formatter.format_with_metadata([generated_sample])
        formatted_sample = result.samples[0]

        # 3. Verify final format is suitable for training
        assert "messages" in formatted_sample
        messages = formatted_sample["messages"]

        # Should have embedded execution format
        assistant_thinking = messages[1]["content"]
        assert "<think>" in assistant_thinking and "</think>" in assistant_thinking
        assert "<tool_call>" in assistant_thinking and "</tool_call>" in assistant_thinking

        tool_response = messages[2]["content"]
        assert "<tool_response>" in tool_response and "</tool_response>" in tool_response

        # Verify rich reasoning is preserved
        assert "User wants restaurant recommendations" in assistant_thinking
        assert "Parse the location request" in assistant_thinking
        assert "Restaurant search tool can find" in assistant_thinking

    def test_config_integration_with_yaml_style(self):
        """Test that the config system properly handles agent_cot_rich."""
        # Simulate YAML config structure
        config_dict = {
            "dataset_system_prompt": "Generate agent reasoning examples",
            "topic_tree": {
                "topic_prompt": "Real-world scenarios requiring tools",
                "provider": "openai",
                "model": "gpt-4",
                "depth": 2,
                "degree": 3,
                "save_as": "topics.jsonl",
            },
            "data_engine": {
                "generation_system_prompt": "Create detailed agent reasoning with tools",
                "provider": "openai",
                "model": "gpt-4",
                "conversation_type": "agent_cot_tools",
                "available_tools": ["get_weather", "calculate"],
                "num_samples": 10,
                "save_as": "raw_data.jsonl",
            },
            "dataset": {
                "save_as": "final_dataset.jsonl",
                "creation": {"num_steps": 10, "batch_size": 2, "sys_msg": True},
                "formatters": [
                    {
                        "name": "tool_calling",
                        "template": "builtin://tool_calling",
                        "output": "formatted.jsonl",
                        "config": {"include_tools_in_system": True},
                    }
                ],
            },
        }

        # Should not raise validation errors
        config = DeepFabricConfig(**config_dict)

        # Verify conversation type is accepted
        assert config.data_engine.conversation_type == "agent_cot_tools"
        assert config.data_engine.available_tools == ["get_weather", "calculate"]

        # Verify formatter configuration
        assert len(config.dataset.formatters) == 1
        formatter_config = config.dataset.formatters[0]
        assert formatter_config.template == "builtin://tool_calling"
        assert formatter_config.config["include_tools_in_system"] is True


class TestToolSystemIntegration:
    """Test integration between tool system and agent CoT generation."""

    def test_tool_registry_creation_for_agent_cot(self):
        """Test that tool registry is properly created for agent_cot_tools."""
        with patch("deepfabric.generator.LLMClient"):
            generator = DataSetGenerator(
                generation_system_prompt="Test prompt",
                conversation_type="agent_cot_tools",
                available_tools=["get_weather", "calculate"],
                provider="openai",
                model_name="gpt-4",
            )

        # Should have initialized tool registry
        assert generator.tool_registry is not None
        assert len(generator.tool_registry.tools) > 0

        # Should contain requested tools
        tool_names = generator.tool_registry.get_tool_names()
        assert "get_weather" in tool_names
        assert "calculate" in tool_names

    def test_tool_metadata_injection(self):
        """Test that available tools are properly injected into samples."""
        with patch("deepfabric.generator.LLMClient"):
            generator = DataSetGenerator(
                generation_system_prompt="Test prompt",
                conversation_type="agent_cot_tools",
                available_tools=["get_weather"],
                provider="openai",
                model_name="gpt-4",
            )

            # Should have tool registry with correct tools
            assert generator.tool_registry is not None
            tool_names = generator.tool_registry.get_tool_names()
            assert "get_weather" in tool_names

            # Should have tool definitions available
            weather_tool = generator.tool_registry.get_tool("get_weather")
            assert weather_tool is not None
            assert weather_tool.name == "get_weather"
            assert len(weather_tool.parameters) > 0
