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.lab.models import Agent, AgentData, Prompt, LlmConfig, Model, Sampling


def test_chat_agent_streaming_template():
    """
    Test chat-agent-streaming template by:
    1. Creating a Lab agent
    2. Generating chat interface
    3. Testing programmatic chat (simulating user input)
    """
    test_dir = tempfile.mkdtemp(prefix="pygeai_chat_agent_test_")
    agent_name = f"test_chat_agent_{os.urandom(4).hex()}"
    agent_id = None

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

    try:
        # Step 1: Create Lab agent
        print("[STEP 1] Creating Lab agent...")
        manager = AILabManager()

        agent = Agent(
            name=agent_name,
            description="Test agent for chat",
            access_scope="private",
            agent_data=AgentData(
                prompt=Prompt(
                    instructions="You are a helpful test assistant. Answer briefly.",
                    inputs=["user_input"],
                    outputs=[],
                ),
                llm_config=LlmConfig(
                    max_tokens=100, timeout=60, sampling=Sampling(temperature=0.5)
                ),
                models=[Model(name="gpt-4")],
            ),
        )

        result = manager.create_agent(agent=agent, automatic_publish=True)

        if not result or not result.id:
            print("[FAIL] Could not create Lab agent")
            sys.exit(1)

        agent_id = result.id
        print(f"[PASS] Lab agent created: {agent_id}")

        # Step 2: Generate chat template
        print("[STEP 2] Generating chat-agent-streaming template...")
        engine = TemplateEngine()

        params = {
            "name": agent_name,
            "description": "Test streaming chat with agent",
        }

        generated_files = engine.generate(
            "chat-agent-streaming", Path(test_dir), params
        )

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

        chat_script = Path(test_dir) / f"chat_{agent_name}.py"
        readme = Path(test_dir) / "README.md"

        assert chat_script.exists(), "chat script not found"
        assert readme.exists(), "README not found"

        print("[PASS] All expected files exist")

        # Step 3: Test the chat script programmatically
        print("[STEP 3] Testing chat functionality...")

        from pygeai.chat.clients import ChatClient
        from pygeai.chat.settings import LLM_SETTINGS

        client = ChatClient()

        messages = [{"role": "user", "content": "Hello, test!"}]

        response_stream = client.chat_completion(
            model=f"saia:agent:{agent_name}",
            messages=messages,
            stream=True,
            **LLM_SETTINGS,
        )

        full_response = ""
        for chunk in response_stream:
            full_response += chunk

        if not full_response:
            print("[FAIL] Chat returned empty response")
            sys.exit(1)

        print(f"[PASS] Chat response received (length: {len(full_response)})")
        print(f"Response preview: {full_response[:100]}...")

        print(
            "\n[PASS] ALL TESTS PASSED - chat-agent-streaming 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...")
        if agent_id:
            try:
                manager = AILabManager()
                manager.delete_agent(agent_id=agent_id)
                print(f"Deleted test 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_chat_agent_streaming_template()
