#!/usr/bin/env python3
"""
Real-world functional tests for Tool Synthesis with real API and database.
"""

import asyncio
import json
import os
import sys
import tempfile
from pathlib import Path
from typing import Any

sys.path.insert(0, os.path.dirname(__file__))


def get_provider():
    return dict(
        model=os.getenv("OPENAI_COMPAT_MODEL", "gpt-4o-mini"),
        base_url=os.getenv("OPENAI_COMPAT_BASE_URL", "https://api.example.com/v1"),
        api_key=os.getenv("OPENAI_COMPAT_API_KEY", "test-api-key"),
    )


def reset_runtime():
    from logos_adk.runtime import _default_runtime_holder

    _default_runtime_holder.reset()


# ============================================================================
# TEST 1: Create a real tool and use it with agent
# ============================================================================


async def test_synthesize_and_use_tool():
    """Test: Synthesize a real tool and use it with an agent."""
    print("\n" + "=" * 60)
    print("TEST 1: Synthesize tool and use with agent (real API)")
    print("=" * 60)

    from logos_adk.tool_synthesis import GeneratedToolStore
    from logos_adk.agent import Agent
    from logos_adk.providers.openai import OpenAICompatibleProvider

    reset_runtime()

    with tempfile.TemporaryDirectory() as tmpdir:
        store = GeneratedToolStore(tmpdir)

        # Define a real async tool
        tool_source = '''
import httpx

async def run(url: str, method: str = "GET") -> dict:
    """Make an HTTP request.
    
    Args:
        url: The URL to request
        method: HTTP method (GET, POST, etc.)
    
    Returns:
        Dictionary with status and body
    """
    try:
        async with httpx.AsyncClient() as client:
            response = await client.request(method, url, timeout=10.0)
            return {
                "status_code": response.status_code,
                "body": response.text[:500],
                "headers": dict(response.headers)
            }
    except Exception as e:
        return {"error": str(e)}
'''

        tool_tests = """
def test_run():
    # Basic test placeholder
    pass
"""

        # Save the tool
        artifact = store.save_tool(
            name="http_request",
            description="Make HTTP requests to any URL",
            parameters={
                "type": "object",
                "properties": {
                    "url": {"type": "string", "description": "URL to request"},
                    "method": {"type": "string", "description": "HTTP method", "default": "GET"},
                },
                "required": ["url"],
            },
            source_code=tool_source,
            tests=tool_tests,
            created_by="test_suite",
        )

        print(f"✓ Saved tool: {artifact.name}")

        # Resolve the tool
        tool = store.resolve_tool("http_request")
        assert tool is not None
        print(f"✓ Resolved tool: {tool.name}")

        # Create agent and add tool
        provider = OpenAICompatibleProvider(**get_provider())
        agent = Agent("test_agent", model=provider)
        agent.tools([tool])

        # Test the tool directly (must be async)
        result = await tool(url="https://httpbin.org/get")
        print(f"✓ Tool executed directly: {result.get('status_code', 'error')}")
        assert "status_code" in result or "error" in result

        # Bind tool to agent
        store.bind_tool_to_agent("http_request", "test_agent")
        print(f"✓ Bound tool to agent")

        # Attach tools to agent
        attached = store.attach_tools_for_agent(agent)
        print(f"✓ Attached {len(attached)} tools to agent")

    print("✓ TEST 1 PASSED\n")
    return True


# ============================================================================
# TEST 2: Multi-step tool pipeline
# ============================================================================


async def test_tool_pipeline():
    """Test: Create multiple tools that work together."""
    print("\n" + "=" * 60)
    print("TEST 2: Tool pipeline (multiple tools)")
    print("=" * 60)

    from logos_adk.tool_synthesis import GeneratedToolStore

    with tempfile.TemporaryDirectory() as tmpdir:
        store = GeneratedToolStore(tmpdir)

        # Tool 1: Data fetcher (async)
        source1 = '''
async def run(query: str) -> dict:
    """Fetch data based on query."""
    return {"data": f"fetched_data_for_{query}", "source": "api"}
'''

        # Tool 2: Data processor (async)
        source2 = '''
async def run(data: str, operation: str = "uppercase") -> dict:
    """Process data with specified operation."""
    if operation == "uppercase":
        result = data.upper()
    elif operation == "length":
        result = len(data)
    else:
        result = data
    return {"result": result, "operation": operation}
'''

        # Tool 3: Data formatter (async)
        source3 = '''
import json

async def run(content: str, format: str = "json") -> str:
    """Format content as specified format."""
    if format == "json":
        return json.dumps({"content": content})
    elif format == "xml":
        return f"<content>{content}</content>"
    return content
'''

        # Save all tools
        for name, source in [("fetcher", source1), ("processor", source2), ("formatter", source3)]:
            store.save_tool(
                name=name,
                description=f"{name} tool",
                parameters={"type": "object", "properties": {}, "required": []},
                source_code=source,
                tests="def test_run():\n    pass",
            )

        # Resolve all tools
        tools = [store.resolve_tool(name) for name in ["fetcher", "processor", "formatter"]]

        # Execute pipeline
        result1 = await tools[0](query="test")
        print(f"  Step 1 - Fetcher: {result1}")

        result2 = await tools[1](data=result1["data"])
        print(f"  Step 2 - Processor: {result2}")

        result3 = await tools[2](content=result2["result"])
        print(f"  Step 3 - Formatter: {result3}")

        # Verify pipeline (uppercase applied)
        assert "FETCHED_DATA_FOR_TEST" in str(result2)

    print("✓ TEST 2 PASSED\n")
    return True


# ============================================================================
# TEST 3: Tool with agent using real API
# ============================================================================


async def test_tool_with_agent_api():
    """Test: Use synthesized tool with real agent API."""
    print("\n" + "=" * 60)
    print("TEST 3: Agent uses synthesized tool (real Polza AI)")
    print("=" * 60)

    from logos_adk.tool_synthesis import GeneratedToolStore
    from logos_adk.agent import Agent
    from logos_adk.providers.openai import OpenAICompatibleProvider

    reset_runtime()

    with tempfile.TemporaryDirectory() as tmpdir:
        store = GeneratedToolStore(tmpdir)

        # Create a calculator tool (async)
        calc_source = '''
async def run(a: int, b: int, operation: str = "add") -> dict:
    """Perform basic arithmetic operations.
    
    Args:
        a: First number
        b: Second number
        operation: Operation (add, subtract, multiply, divide)
    
    Returns:
        Dictionary with result
    """
    if operation == "add":
        result = a + b
    elif operation == "subtract":
        result = a - b
    elif operation == "multiply":
        result = a * b
    elif operation == "divide":
        if b == 0:
            return {"error": "Cannot divide by zero"}
        result = a / b
    else:
        return {"error": f"Unknown operation: {operation}"}
    return {"result": result, "operation": operation}
'''

        artifact = store.save_tool(
            name="calculator",
            description="Perform basic arithmetic operations",
            parameters={
                "type": "object",
                "properties": {
                    "a": {"type": "integer", "description": "First number"},
                    "b": {"type": "integer", "description": "Second number"},
                    "operation": {"type": "string", "description": "Operation", "default": "add"},
                },
                "required": ["a", "b"],
            },
            source_code=calc_source,
            tests="def test_run():\n    pass",
        )

        # Load tool
        tool = store.resolve_tool("calculator")

        # Create agent with tool
        provider = OpenAICompatibleProvider(**get_provider())
        agent = Agent("math_agent", model=provider)
        agent.tools([tool])

        # Test tool directly first
        direct_result = await tool(a=10, b=5, operation="multiply")
        print(f"  Direct tool call: 10 * 5 = {direct_result}")
        assert direct_result["result"] == 50

        # Run agent with tool - agent should use the tool
        response = await agent.run("What is 25 + 17? Use the calculator if needed.")
        print(f"  Agent response: {response.content[:200]}...")

    print("✓ TEST 3 PASSED\n")
    return True


# ============================================================================
# TEST 4: Team with synthesized tools
# ============================================================================


async def test_team_with_tools():
    """Test: Team members share synthesized tools."""
    print("\n" + "=" * 60)
    print("TEST 4: Team with synthesized tools")
    print("=" * 60)

    from logos_adk.tool_synthesis import GeneratedToolStore
    from logos_adk.team import Team
    from logos_adk.agent import Agent
    from logos_adk.providers.openai import OpenAICompatibleProvider

    reset_runtime()

    with tempfile.TemporaryDirectory() as tmpdir:
        store = GeneratedToolStore(tmpdir)

        # Create a shared tool (async)
        tool_source = '''
async def run(message: str) -> dict:
    """Process a message."""
    return {"processed": message.upper(), "length": len(message)}
'''

        store.save_tool(
            name="message_processor",
            description="Process messages",
            parameters={
                "type": "object",
                "properties": {"message": {"type": "string"}},
                "required": ["message"],
            },
            source_code=tool_source,
            tests="def test_run():\n    pass",
        )

        # Bind to team
        store.bind_tool_to_team("message_processor", "my_team")

        # Create team
        provider = OpenAICompatibleProvider(**get_provider())

        agent1 = Agent("researcher", model=provider)
        agent2 = Agent("writer", model=provider)

        team = Team(
            name="my_team",
            agents=[agent1, agent2],
        )

        # Attach tools to team
        attached = store.attach_tools_for_team(team)
        print(f"  Attached {len(attached)} tools to team")

        # Check both agents have the tool
        tool = store.resolve_tool("message_processor")
        assert tool is not None

        # Test tool
        result = await tool(message="hello world")
        assert result["processed"] == "HELLO WORLD"
        print(f"  Tool result: {result}")

    print("✓ TEST 4 PASSED\n")
    return True


# ============================================================================
# TEST 5: Tool with PostgreSQL store (metadata)
# ============================================================================


async def test_tool_store_with_postgres():
    """Test: Tool store with PostgreSQL database."""
    print("\n" + "=" * 60)
    print("TEST 5: Tool store metadata in PostgreSQL")
    print("=" * 60)

    from logos_adk.tool_synthesis import GeneratedToolStore

    # Get PostgreSQL connection
    dsn = os.getenv("LOGOS_ADK_DATABASE_URL", "postgresql://crm:crm@localhost:5432/crm")

    try:
        import psycopg

        conn = psycopg.connect(dsn)
        print(f"  Connected to PostgreSQL")

        # Create a table to store tool metadata
        conn.execute("""
            CREATE TABLE IF NOT EXISTS synthesized_tool_metadata (
                tool_name VARCHAR(255) PRIMARY KEY,
                description TEXT,
                parameters JSONB,
                module_path TEXT,
                created_at TIMESTAMP,
                created_by VARCHAR(255)
            )
        """)
        conn.commit()
        print(f"  Created metadata table")

        # Store tool metadata
        tool_name = "pg_test_tool"
        conn.execute(
            "INSERT INTO synthesized_tool_metadata (tool_name, description, parameters, created_by) VALUES (%s, %s, %s, %s) ON CONFLICT (tool_name) DO NOTHING",
            (tool_name, "Test tool stored in PG", json.dumps({"type": "object"}), "test_suite"),
        )
        conn.commit()

        # Read back
        result = conn.execute(
            "SELECT tool_name, description FROM synthesized_tool_metadata WHERE tool_name = %s",
            (tool_name,),
        ).fetchone()
        print(f"  Stored and retrieved: {result}")

        assert result is not None
        assert result[0] == tool_name

        # Cleanup
        conn.execute("DROP TABLE IF EXISTS synthesized_tool_metadata")
        conn.commit()
        conn.close()

    except Exception as e:
        print(f"  PostgreSQL test skipped: {e}")
        return True

    print("✓ TEST 5 PASSED\n")
    return True


# ============================================================================
# TEST 6: Full workflow - Agent generates tool dynamically
# ============================================================================


async def test_dynamic_tool_generation():
    """Test: Agent dynamically generates a new tool."""
    print("\n" + "=" * 60)
    print("TEST 6: Dynamic tool generation by agent")
    print("=" * 60)

    from logos_adk.tool_synthesis import GeneratedToolStore
    from logos_adk.agent import Agent
    from logos_adk.providers.openai import OpenAICompatibleProvider

    reset_runtime()

    with tempfile.TemporaryDirectory() as tmpdir:
        store = GeneratedToolStore(tmpdir)

        # Create an agent that can develop tools
        provider = OpenAICompatibleProvider(**get_provider())
        agent = Agent("tool_builder", model=provider)

        # In real scenario, agent would call develop_tool function
        # Here we simulate what the agent would produce
        new_tool_source = '''
async def run(text: str, mode: str = "word_count") -> dict:
    """Analyze text in different modes.
    
    Args:
        text: Text to analyze
        mode: Analysis mode (word_count, char_count, reverse)
    
    Returns:
        Analysis result
    """
    if mode == "word_count":
        result = len(text.split())
    elif mode == "char_count":
        result = len(text)
    elif mode == "reverse":
        result = text[::-1]
    else:
        result = text
    return {"result": result, "mode": mode}
'''

        # Save the dynamically generated tool
        artifact = store.save_tool(
            name="text_analyzer",
            description="Analyze text in various ways",
            parameters={
                "type": "object",
                "properties": {
                    "text": {"type": "string", "description": "Text to analyze"},
                    "mode": {
                        "type": "string",
                        "description": "Analysis mode",
                        "default": "word_count",
                    },
                },
                "required": ["text"],
            },
            source_code=new_tool_source,
            tests="def test_run():\n    pass",
            created_by="tool_builder",
        )

        print(f"  Saved dynamically generated tool: {artifact.name}")

        # Load and test
        tool = store.resolve_tool("text_analyzer")

        # Test various modes
        result1 = await tool(text="hello world", mode="word_count")
        print(f"  word_count: {result1}")

        result2 = await tool(text="hello world", mode="char_count")
        print(f"  char_count: {result2}")

        result3 = await tool(text="abc", mode="reverse")
        print(f"  reverse: {result3}")

        assert result1["result"] == 2
        assert result2["result"] == 11
        assert result3["result"] == "cba"

    print("✓ TEST 6 PASSED\n")
    return True


# ============================================================================
# MAIN
# ============================================================================


async def main():
    print("=" * 60)
    print("TOOL SYNTHESIS - REAL API & DATABASE TESTS")
    print("=" * 60)
    print(f"API: {get_provider()['base_url']}")
    print(f"Model: {get_provider()['model']}")
    print("=" * 60)

    results = []

    tests = [
        ("Synthesize & Use Tool", test_synthesize_and_use_tool),
        ("Tool Pipeline", test_tool_pipeline),
        ("Tool with Agent API", test_tool_with_agent_api),
        ("Team with Tools", test_team_with_tools),
        ("PostgreSQL Metadata", test_tool_store_with_postgres),
        ("Dynamic Tool Generation", test_dynamic_tool_generation),
    ]

    for name, test_func in tests:
        try:
            result = await test_func()
            results.append((name, result, None))
        except Exception as e:
            print(f"✗ {name}: {e}")
            import traceback

            traceback.print_exc()
            results.append((name, False, str(e)))

    print("=" * 60)
    print("RESULTS")
    print("=" * 60)

    passed = sum(1 for _, ok, _ in results if ok)
    total = len(results)

    for name, ok, err in results:
        status = "✓ PASS" if ok else "✗ FAIL"
        print(f"{status}: {name}")
        if err:
            print(f"    Error: {err}")

    print(f"\n{passed}/{total} tests passed")
    return passed == total


if __name__ == "__main__":
    success = asyncio.run(main())
    sys.exit(0 if success else 1)
