import os
import sys
import tempfile
import shutil
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent))

from pygeai.templates.engine import TemplateEngine
from pygeai.lab.managers import AILabManager
from pygeai.chat.clients import ChatClient
from pygeai.chat.settings import LLM_SETTINGS


def test_lab_multi_agent_template():
    """
    Test lab-multi-agent template by:
    1. Generating template files
    2. Creating multiple specialized agents
    3. Testing orchestration workflow
    """
    test_dir = tempfile.mkdtemp(prefix="pygeai_multiagent_test_")
    system_name = f"test_ma_{os.urandom(4).hex()}"

    agent_ids = []

    print(f"Testing in: {test_dir}")

    try:
        manager = AILabManager()
        client = ChatClient()

        print("[STEP 1] Generating lab-multi-agent template...")
        engine = TemplateEngine()

        params = {
            "name": system_name,
            "description": "Test multi-agent system",
            "model": "gpt-4",
            "temperature": "0.7",
        }

        generated_files = engine.generate("lab-multi-agent", Path(test_dir), params)

        assert len(generated_files) == 3, (
            f"Expected 3 files, got {len(generated_files)}"
        )
        print("[PASS] Template generated successfully")

        print("[STEP 2] Creating agents from generated script...")

        agents_to_create = [
            {
                "name": f"{system_name}_researcher",
                "description": "Research agent",
                "instructions": "You are a research assistant. Provide brief research findings.",
            },
            {
                "name": f"{system_name}_writer",
                "description": "Writer agent",
                "instructions": "You are a writer. Create brief content from research.",
            },
            {
                "name": f"{system_name}_reviewer",
                "description": "Reviewer agent",
                "instructions": "You are a reviewer. Provide brief feedback on content.",
            },
        ]

        from pygeai.lab.models import (
            Agent,
            AgentData,
            Prompt,
            LlmConfig,
            Model,
            Sampling,
        )

        for agent_config in agents_to_create:
            agent = Agent(
                name=agent_config["name"],
                description=agent_config["description"],
                access_scope="private",
                agent_data=AgentData(
                    prompt=Prompt(
                        instructions=agent_config["instructions"],
                        inputs=["task_input"],
                        outputs=[],
                    ),
                    llm_config=LlmConfig(
                        max_tokens=150, timeout=60, sampling=Sampling(temperature=0.7)
                    ),
                    models=[Model(name="gpt-4")],
                ),
            )

            result = manager.create_agent(agent=agent, automatic_publish=True)
            agent_ids.append(result.id)
            print(f"[PASS] Created agent: {agent_config['name']}")

        print("[STEP 3] Testing orchestration workflow...")

        task = "Explain Python in one sentence"

        print("  - Testing researcher agent...")
        research_result = client.chat_completion(
            model=f"saia:agent:{system_name}_researcher",
            messages=[{"role": "user", "content": f"Research: {task}"}],
            stream=False,
            **LLM_SETTINGS,
        )

        research_response = research_result["choices"][0]["message"]["content"]
        assert len(research_response) > 0, "Research agent returned empty response"
        print(f"[PASS] Research response: {research_response[:80]}...")

        print("  - Testing writer agent...")
        writer_result = client.chat_completion(
            model=f"saia:agent:{system_name}_writer",
            messages=[
                {
                    "role": "user",
                    "content": f"Write based on: {research_response[:100]}",
                }
            ],
            stream=False,
            **LLM_SETTINGS,
        )

        writer_response = writer_result["choices"][0]["message"]["content"]
        assert len(writer_response) > 0, "Writer agent returned empty response"
        print(f"[PASS] Writer response: {writer_response[:80]}...")

        print("  - Testing reviewer agent...")
        reviewer_result = client.chat_completion(
            model=f"saia:agent:{system_name}_reviewer",
            messages=[{"role": "user", "content": f"Review: {writer_response[:100]}"}],
            stream=False,
            **LLM_SETTINGS,
        )

        reviewer_response = reviewer_result["choices"][0]["message"]["content"]
        assert len(reviewer_response) > 0, "Reviewer agent returned empty response"
        print(f"[PASS] Reviewer response: {reviewer_response[:80]}...")

        print(
            "\n[PASS] ALL TESTS PASSED - lab-multi-agent template is working correctly!"
        )

    except Exception as e:
        print(f"[FAIL] Test failed with error: {e}")
        import traceback

        traceback.print_exc()
        sys.exit(1)

    finally:
        print("\nCleaning up...")

        for agent_id in agent_ids:
            try:
                manager.delete_agent(agent_id=agent_id)
                print(f"Deleted agent: {agent_id}")
            except Exception as e:
                print(f"Warning: Could not delete agent: {e}")

        if os.path.exists(test_dir):
            shutil.rmtree(test_dir)
            print(f"Deleted test directory: {test_dir}")


if __name__ == "__main__":
    test_lab_multi_agent_template()
