import os
import sys
import time

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

from core.model_executor import ModelExecutor, ExecutionResult
from core.model_pool import SLOT_SUMMARY


def test_connection():
    print("=== Z.AI API Connection Test ===\n")

    api_key = os.environ.get("ZAI_API_KEY")
    if not api_key:
        print("ERROR: ZAI_API_KEY not set in environment")
        return False

    print(f"API Key: {api_key[:20]}...{api_key[-10:]}\n")
    print(SLOT_SUMMARY)

    with ModelExecutor(api_key) as executor:
        print("\nTesting router_small (glm-4.7-flash)...")
        result = executor.execute_simple(
            "router_small",
            "Say 'OK'",
            max_tokens=10,
        )
        print_result(result)
        return result.success


def print_result(result: ExecutionResult):
    if result.success:
        print(f"  [OK] model={result.model_used}")
        print(f"  latency={result.latency_ms}ms")
        print(f"  response: {result.content}")
    else:
        print(f"  [FAIL] {result.error}")


if __name__ == "__main__":
    success = test_connection()
    sys.exit(0 if success else 1)