"""
Test script for Promptev Python SDK.

Prerequisites:
  1. Backend running: uvicorn app.main:app --host 0.0.0.0 --port 8003 --reload
  2. SDK installed:   pip install -e .

Usage:
  python test/test_promptev.py
"""

import asyncio
from promptev import (
    PromptevClient,
    AgentSession,
    AgentEvent,
    ValidationError,
    NotFoundError,
    AuthenticationError,
)

# ── Config ─────────────────────────────────────────────
# Pick a project_key + chatbot_id from your local DB.

PROJECT_KEY = "pl_sk_HTQ1_3QFr7kREXSyr_vhKANz7eqwLBPn"
CHATBOT_ID = "5b299c4c-7821-48cb-8d14-eb758c880331"  # HR BOT
PROMPT_KEY = "review-user-prompt"                      # change to a prompt in your project
BASE_URL = "http://127.0.0.1:8003"

client = PromptevClient(project_key=PROJECT_KEY, base_url=BASE_URL)

passed = 0
failed = 0
skipped = 0


def divider(title: str):
    print(f"\n{'='*60}")
    print(f"  {title}")
    print(f"{'='*60}\n")


def mark_pass():
    global passed
    passed += 1
    print("  ✅ PASS")


def mark_fail(e):
    global failed
    failed += 1
    print(f"  ❌ FAIL — {type(e).__name__}: {e}")


def mark_skip(reason):
    global skipped
    skipped += 1
    print(f"  ⏭️  SKIP — {reason}")


# ── Test 1: run_prompt (sync) ──────────────────────────
def test_run_prompt():
    divider("Test 1: run_prompt (sync)")
    try:
        output = client.run_prompt(
            PROMPT_KEY,
            query="How to improve LLM reasoning?",
            variables={"user_prompt_text": "How to improve LLM reasoning?"}
        )
        print(f"  Result ({len(output)} chars):")
        print(f"  {output[:200]}...")
        mark_pass()
    except NotFoundError:
        mark_skip("prompt not found (update PROMPT_KEY)")
    except Exception as e:
        mark_fail(e)


# ── Test 2: run_prompt error handling ──────────────────
def test_error_handling():
    divider("Test 2: Error handling — bad API key")
    bad_client = PromptevClient(project_key="invalid_key", base_url=BASE_URL)
    try:
        bad_client.run_prompt("anything", query="test")
        mark_fail(Exception("should have raised an exception"))
    except (NotFoundError, AuthenticationError) as e:
        print(f"  Bad key → {type(e).__name__}: {e}")
        mark_pass()
    except Exception as e:
        mark_fail(e)
    finally:
        bad_client.close()


# ── Test 3: stream_prompt (sync) ──────────────────────
def test_stream_prompt():
    divider("Test 3: stream_prompt (sync)")
    try:
        event_count = 0
        final_output = None
        for event in client.stream_prompt(
            PROMPT_KEY,
            query="How to improve LLM reasoning?",
            variables={"user_prompt_text": "How to improve LLM reasoning?"}
        ):
            event_count += 1
            assert isinstance(event, AgentEvent)
            print(f"  [{event.type}] {event.output[:80]}")
            if event.type == "done":
                final_output = event.output
                break

        if final_output:
            print(f"  Received {event_count} events, final output: {len(final_output)} chars")
            mark_pass()
        else:
            mark_fail(Exception("no 'done' event received"))
    except NotFoundError:
        mark_skip("prompt not found (update PROMPT_KEY)")
    except Exception as e:
        mark_fail(e)


# ── Test 4: Agent start (sync) ─────────────────────────
def test_start_agent():
    divider("Test 4: start_agent (sync)")
    try:
        session = client.start_agent(CHATBOT_ID, visitor="test-user")
        print(f"  session_token: {session.session_token}")
        print(f"  chatbot_id:    {session.chatbot_id}")
        print(f"  name:          {session.name}")
        print(f"  memory:        {session.memory_enabled}")
        print(f"  messages:      {len(session.messages)} previous")
        assert isinstance(session, AgentSession)
        assert session.session_token
        mark_pass()
        return session
    except Exception as e:
        mark_fail(e)
        return None


# ── Test 5: Agent stream (sync) ────────────────────────
def test_stream_agent(session: AgentSession):
    divider("Test 5: stream_agent (sync)")
    if not session:
        mark_skip("no session from test 4")
        return
    try:
        event_count = 0
        for event in client.stream_agent(
            session.chatbot_id,
            session_token=session.session_token,
            query="Hello, who are you?"
        ):
            event_count += 1
            assert isinstance(event, AgentEvent)
            print(f"  [{event.type}] {event.output[:100]}")
            if event.type == "done":
                break

        print(f"  Received {event_count} events")
        mark_pass()
    except Exception as e:
        mark_fail(e)


# ── Test 6: Multi-turn conversation ───────────────────
def test_multi_turn(session: AgentSession):
    divider("Test 6: Multi-turn conversation")
    if not session:
        mark_skip("no session from test 4")
        return
    try:
        final = None
        for event in client.stream_agent(
            session.chatbot_id,
            session_token=session.session_token,
            query="Tell me more about that."
        ):
            if event.type == "done":
                final = event.output
                break

        if final:
            print(f"  Follow-up response ({len(final)} chars):")
            print(f"  {final[:200]}...")
            mark_pass()
        else:
            mark_fail(Exception("no 'done' event received"))
    except Exception as e:
        mark_fail(e)


# ── Test 7: arun_prompt (async) ───────────────────────
async def test_arun_prompt():
    divider("Test 7: arun_prompt (async)")
    try:
        output = await client.arun_prompt(
            PROMPT_KEY,
            query="How to improve LLM reasoning?",
            variables={"user_prompt_text": "How to improve LLM reasoning?"}
        )
        print(f"  Result ({len(output)} chars):")
        print(f"  {output[:200]}...")
        mark_pass()
    except NotFoundError:
        mark_skip("prompt not found (update PROMPT_KEY)")
    except Exception as e:
        mark_fail(e)


# ── Test 8: astream_prompt (async) ────────────────────
async def test_astream_prompt():
    divider("Test 8: astream_prompt (async)")
    try:
        event_count = 0
        final_output = None
        async for event in client.astream_prompt(
            PROMPT_KEY,
            query="How to improve LLM reasoning?",
            variables={"user_prompt_text": "How to improve LLM reasoning?"}
        ):
            event_count += 1
            assert isinstance(event, AgentEvent)
            print(f"  [{event.type}] {event.output[:80]}")
            if event.type == "done":
                final_output = event.output
                break

        if final_output:
            print(f"  Received {event_count} events, final output: {len(final_output)} chars")
            mark_pass()
        else:
            mark_fail(Exception("no 'done' event received"))
    except NotFoundError:
        mark_skip("prompt not found (update PROMPT_KEY)")
    except Exception as e:
        mark_fail(e)


# ── Test 9: Async agent flow ─────────────────────────
async def test_async_agent():
    divider("Test 9: astart_agent + astream_agent (async)")
    try:
        session = await client.astart_agent(CHATBOT_ID, visitor="async-test")
        print(f"  session_token: {session.session_token}")

        event_count = 0
        async for event in client.astream_agent(
            session.chatbot_id,
            session_token=session.session_token,
            query="What can you help me with?"
        ):
            event_count += 1
            print(f"  [{event.type}] {event.output[:100]}")
            if event.type == "done":
                break

        print(f"  Received {event_count} events")
        mark_pass()
    except Exception as e:
        mark_fail(e)


# ── Run all tests ─────────────────────────────────────
if __name__ == "__main__":
    print("Promptev SDK Test Suite")
    print(f"Backend: {BASE_URL}")
    print(f"Project: {PROJECT_KEY[:20]}...")
    print(f"Chatbot: {CHATBOT_ID}")
    print(f"Prompt:  {PROMPT_KEY}")

    # Sync tests
    test_run_prompt()
    test_error_handling()
    test_stream_prompt()
    session = test_start_agent()
    test_stream_agent(session)
    test_multi_turn(session)

    # Async tests — single asyncio.run() to share the async client
    async def run_async_tests():
        await test_arun_prompt()
        await test_astream_prompt()
        await test_async_agent()

    asyncio.run(run_async_tests())

    # Summary
    divider("Results")
    total = passed + failed + skipped
    print(f"  Total:   {total}")
    print(f"  Passed:  {passed}")
    print(f"  Failed:  {failed}")
    print(f"  Skipped: {skipped}")
    print()

    if failed > 0:
        print("  ⚠️  Some tests failed!")
    else:
        print("  🎉 All tests passed!")

    client.close()
