#!/usr/bin/env python3
"""
Smartflow Integration Test Suite
=================================
Fill in SMARTFLOW_HOST and VIRTUAL_KEY below (or set them as environment
variables), then run:

    pip install openai anthropic httpx
    python3 smartflow_integration_test.py

All tests are independent — a failure in one does not stop the rest.
"""

import os, sys, time, json

# =============================================================================
# Configuration — replace these or set as env vars
# =============================================================================
SMARTFLOW_HOST = os.getenv("SMARTFLOW_HOST", "https://your-smartflow.example.com")
VIRTUAL_KEY    = os.getenv("VIRTUAL_KEY",    "sk-sf-your-virtual-key")
# =============================================================================

try:
    import httpx
    from openai import OpenAI
    from anthropic import Anthropic
except ImportError as e:
    sys.exit(f"Missing dependency: {e}\n  Run: pip install openai anthropic httpx")

PASS_COUNT = 0
FAIL_COUNT = 0

def ok(label, detail=""):
    global PASS_COUNT
    PASS_COUNT += 1
    suffix = f"  ({detail})" if detail else ""
    print(f"  \033[32mPASS\033[0m  {label}{suffix}")

def fail(label, err):
    global FAIL_COUNT
    FAIL_COUNT += 1
    print(f"  \033[31mFAIL\033[0m  {label}")
    print(f"         {str(err)[:120]}")

def section(title):
    print(f"\n[ {title} ]")


# ── 1. Health ─────────────────────────────────────────────────────────────────
section("Health & Connectivity")
try:
    r = httpx.get(f"{SMARTFLOW_HOST}/health", timeout=10)
    assert r.status_code in (200, 204), f"status {r.status_code}"
    ok("GET /health", r.text[:60].strip())
except Exception as e:
    fail("GET /health", e)


# ── 2. OpenAI-compatible chat ─────────────────────────────────────────────────
section("OpenAI-Compatible  /v1/chat/completions")
openai_client = OpenAI(base_url=f"{SMARTFLOW_HOST}/v1", api_key=VIRTUAL_KEY)

try:
    r = openai_client.chat.completions.create(
        model="gpt-4o",
        messages=[{"role": "user", "content": "Reply with exactly one word: PASS"}],
        max_tokens=16,
    )
    text = r.choices[0].message.content or ""
    ok("GPT-4o chat", text.strip()[:60])
except Exception as e:
    fail("GPT-4o chat", e)

try:
    r = openai_client.chat.completions.create(
        model="claude-sonnet-4-6",
        messages=[{"role": "user", "content": "Reply with exactly one word: PASS"}],
        max_tokens=16,
    )
    text = r.choices[0].message.content or ""
    ok("Claude Sonnet 4.6 via OpenAI route", text.strip()[:60])
except Exception as e:
    fail("Claude Sonnet 4.6 via OpenAI route", e)

try:
    models = list(openai_client.models.list())
    ok("GET /v1/models", f"{len(models)} models")
except Exception as e:
    fail("GET /v1/models", e)


# ── 3. Anthropic native messages ──────────────────────────────────────────────
section("Anthropic Native  /anthropic/v1/messages")
anthropic_client = Anthropic(
    base_url=f"{SMARTFLOW_HOST}/anthropic",
    api_key=VIRTUAL_KEY,
)

try:
    r = anthropic_client.messages.create(
        model="claude-sonnet-4-6",
        max_tokens=32,
        messages=[{"role": "user", "content": "Reply with exactly one word: PASS"}],
    )
    text = r.content[0].text if r.content else ""
    ok("Claude Sonnet 4.6 native message", text.strip()[:60])
except Exception as e:
    fail("Claude Sonnet 4.6 native message", e)


# ── 4. Cache ──────────────────────────────────────────────────────────────────
section("Cache & Stats")
try:
    r = httpx.get(
        f"{SMARTFLOW_HOST}/api/metacache/stats",
        headers={"Authorization": f"Bearer {VIRTUAL_KEY}"},
        timeout=10,
    )
    assert r.status_code == 200, f"status {r.status_code}"
    data = r.json()
    ok("GET /api/metacache/stats", json.dumps(data)[:80])
except Exception as e:
    fail("GET /api/metacache/stats", e)

# Send same prompt twice to exercise semantic cache
prompt = "What is the capital of France?"
try:
    for i in range(2):
        openai_client.chat.completions.create(
            model="claude-sonnet-4-6",
            messages=[{"role": "user", "content": prompt}],
            max_tokens=32,
        )
    ok("Cache warm-up (2× identical prompt)")
except Exception as e:
    fail("Cache warm-up", e)


# ── 5. Auth rejection ─────────────────────────────────────────────────────────
section("Auth Validation")
try:
    bad_client = OpenAI(base_url=f"{SMARTFLOW_HOST}/v1", api_key="invalid-key-xyz")
    bad_client.chat.completions.create(
        model="gpt-4o",
        messages=[{"role": "user", "content": "test"}],
        max_tokens=4,
    )
    fail("Invalid key rejected", "Expected auth error but request succeeded")
except Exception as e:
    msg = str(e).lower()
    if any(w in msg for w in ["401", "403", "unauthorized", "forbidden", "invalid", "authentication"]):
        ok("Invalid key correctly rejected", str(e)[:60])
    else:
        fail("Invalid key rejected", f"Unexpected error type: {e}")


# ── Summary ───────────────────────────────────────────────────────────────────
print(f"\n{'─'*50}")
print(f"  \033[32m{PASS_COUNT} passed\033[0m   \033[31m{FAIL_COUNT} failed\033[0m")
if FAIL_COUNT == 0:
    print("  \033[32mAll tests passed. Deployment is healthy.\033[0m")
else:
    print("  \033[33mSome tests failed — check host URL, virtual key, and running services.\033[0m")
print()