#!/usr/bin/env python3
"""
Verify Moltbot Hybrid Engine API is working.
Tests: health, /v1/models, chat completions (string + array content), optional streaming.

Results:
  - By default: printed to stdout only.
  - With --output PATH: JSON report written to PATH (e.g. scripts/verify_results.json).
  - Path is printed at the end so you know where results are.

Usage:
  python3 scripts/verify_api.py --api-key YOUR_KEY
  python3 scripts/verify_api.py --output scripts/verify_results.json --api-key YOUR_KEY
  python3 scripts/verify_api.py --base-url http://localhost:7860
"""
import argparse
import json
import os
import sys
import time

try:
    import requests
except ImportError:
    print("Install requests: pip install requests")
    sys.exit(1)

DEFAULT_BASE = os.environ.get("MOLTBOT_VERIFY_BASE_URL", "https://deebee7-moltbot-hybrid-engine.hf.space")
DEFAULT_KEY = os.environ.get("MOLTBOT_VERIFY_API_KEY", "")

# Default path for results when using --output without a path
DEFAULT_RESULTS_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "verify_results.json")


def run(name: str, ok: bool, detail: str = "", results: list = None) -> None:
    status = "PASS" if ok else "FAIL"
    print(f"  [{status}] {name}" + (f" — {detail}" if detail else ""))
    if results is not None:
        results.append({"name": name, "passed": ok, "detail": detail or None})


def main() -> int:
    p = argparse.ArgumentParser(description="Verify Moltbot Hybrid Engine API")
    p.add_argument("--base-url", default=DEFAULT_BASE, help="API base URL (no trailing slash)")
    p.add_argument("--api-key", default=DEFAULT_KEY, help="API key for /v1/chat/completions (Bearer). Default: MOLTBOT_VERIFY_API_KEY")
    p.add_argument("--no-stream-test", action="store_true", help="Skip streaming chat test")
    p.add_argument("--timeout", type=int, default=60, help="Request timeout seconds")
    p.add_argument("--output", "-o", default="", help="Write JSON results to this path (e.g. scripts/verify_results.json)")
    args = p.parse_args()

    results = []  # list of {name, passed, detail}
    def r(name: str, ok: bool, detail: str = "") -> None:
        run(name, ok, detail, results)

    base = args.base_url.rstrip("/")
    headers = {
        "Authorization": f"Bearer {args.api_key}",
        "Content-Type": "application/json",
    }
    timeout = args.timeout
    fails = 0

    print("Moltbot Hybrid Engine — API verification")
    print(f"Base URL: {base}")
    print()

    # --- GET / ---
    print("1. GET / (health)")
    try:
        resp = requests.get(f"{base}/", timeout=10)
        r("GET /", resp.status_code == 200, f"status={resp.status_code}" if resp.status_code != 200 else "")
        if resp.status_code != 200:
            fails += 1
        else:
            try:
                d = resp.json()
                print(f"     status={d.get('status')}, version={d.get('version')}")
            except Exception:
                pass
    except requests.RequestException as e:
        r("GET /", False, str(e))
        fails += 1
    print()

    # --- GET /health ---
    print("2. GET /health")
    try:
        resp = requests.get(f"{base}/health", timeout=10)
        r("GET /health", resp.status_code == 200, f"status={resp.status_code}" if resp.status_code != 200 else "")
        if resp.status_code != 200:
            fails += 1
        else:
            try:
                d = resp.json()
                llm = d.get("llm_backends", {})
                print(f"     ollama={llm.get('ollama', {}).get('running')}, hf_available={llm.get('hf_inference_api', {}).get('available')}")
            except Exception:
                pass
    except requests.RequestException as e:
        r("GET /health", False, str(e))
        fails += 1
    print()

    # --- GET /v1/models ---
    print("3. GET /v1/models")
    try:
        resp = requests.get(f"{base}/v1/models", timeout=10)
        r("GET /v1/models", resp.status_code == 200, f"status={resp.status_code}" if resp.status_code != 200 else "")
        if resp.status_code != 200:
            fails += 1
        else:
            try:
                d = resp.json()
                ids = [m.get("id") for m in d.get("data", [])]
                print(f"     models={ids}")
            except Exception:
                pass
    except requests.RequestException as e:
        r("GET /v1/models", False, str(e))
        fails += 1
    print()

    # --- POST /v1/chat/completions (string content) ---
    print("4. POST /v1/chat/completions (string content, non-streaming)")
    try:
        body = {
            "model": "moltbot-legal",
            "messages": [{"role": "user", "content": "Reply with exactly: OK"}],
            "stream": False,
            "max_tokens": 50,
        }
        resp = requests.post(f"{base}/v1/chat/completions", headers=headers, json=body, timeout=timeout)
        r("Chat (string)", resp.status_code == 200, f"status={resp.status_code}" if resp.status_code != 200 else "")
        if resp.status_code != 200:
            fails += 1
            if resp.text:
                print(f"     body: {resp.text[:300]}")
        else:
            try:
                d = resp.json()
                content = (d.get("choices") or [{}])[0].get("message", {}).get("content", "")
                print(f"     response length={len(content)}, has 'OK'={('OK' in content)}")
            except Exception as e:
                print(f"     parse: {e}")
    except requests.RequestException as e:
        r("Chat (string)", False, str(e))
        fails += 1
    print()

    # --- POST /v1/chat/completions (array content — 422 fix) ---
    print("5. POST /v1/chat/completions (array content, R1/vision-style)")
    try:
        body = {
            "model": "moltbot-legal",
            "messages": [
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": "Reply with exactly: ARRAY_OK"},
                    ],
                }
            ],
            "stream": False,
            "max_tokens": 50,
        }
        resp = requests.post(f"{base}/v1/chat/completions", headers=headers, json=body, timeout=timeout)
        r("Chat (array content)", resp.status_code == 200, f"status={resp.status_code}" if resp.status_code != 200 else "")
        if resp.status_code != 200:
            fails += 1
            if resp.text:
                print(f"     body: {resp.text[:300]}")
        else:
            try:
                d = resp.json()
                content = (d.get("choices") or [{}])[0].get("message", {}).get("content", "")
                print(f"     response length={len(content)}, has 'ARRAY_OK'={('ARRAY_OK' in content)}")
            except Exception as e:
                print(f"     parse: {e}")
    except requests.RequestException as e:
        r("Chat (array content)", False, str(e))
        fails += 1
    print()

    # --- POST /v1/chat/completions (array with text + image placeholder) ---
    print("6. POST /v1/chat/completions (array: text + image part)")
    try:
        body = {
            "model": "moltbot-legal",
            "messages": [
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": "Say only: IMAGE_OK"},
                        {"type": "image_url", "image_url": {"url": "https://example.com/fake.png"}},
                    ],
                }
            ],
            "stream": False,
            "max_tokens": 50,
        }
        resp = requests.post(f"{base}/v1/chat/completions", headers=headers, json=body, timeout=timeout)
        r("Chat (text+image parts)", resp.status_code == 200, f"status={resp.status_code}" if resp.status_code != 200 else "")
        if resp.status_code != 200:
            fails += 1
            if resp.text:
                print(f"     body: {resp.text[:200]}")
        else:
            try:
                d = resp.json()
                content = (d.get("choices") or [{}])[0].get("message", {}).get("content", "")
                print(f"     response length={len(content)}")
            except Exception as e:
                print(f"     parse: {e}")
    except requests.RequestException as e:
        r("Chat (text+image parts)", False, str(e))
        fails += 1
    print()

    # --- Streaming (optional) ---
    if not args.no_stream_test:
        print("7. POST /v1/chat/completions (streaming)")
        try:
            body = {
                "model": "moltbot-legal",
                "messages": [{"role": "user", "content": "Say hello in one word."}],
                "stream": True,
                "max_tokens": 20,
            }
            resp = requests.post(f"{base}/v1/chat/completions", headers=headers, json=body, timeout=timeout, stream=True)
            r("Chat (streaming)", resp.status_code == 200, f"status={resp.status_code}" if resp.status_code != 200 else "")
            if resp.status_code != 200:
                fails += 1
            else:
                chunks = [line for line in resp.iter_lines(decode_unicode=True) if line and line.startswith("data: ") and line != "data: [DONE]"]
                print(f"     chunks received={len(chunks)}")
        except requests.RequestException as e:
            r("Chat (streaming)", False, str(e))
            fails += 1
        print()
    else:
        print("7. POST /v1/chat/completions (streaming) — skipped (--no-stream-test)")
        print()

    # --- Summary ---
    print("---")
    if fails == 0:
        print("All checks passed.")
    else:
        print(f"Failed: {fails} check(s).")

    # --- Write results to file if requested ---
    out_path = (args.output or "").strip()
    if out_path:
        report = {
            "base_url": base,
            "timestamp_iso": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
            "all_passed": fails == 0,
            "failed_count": fails,
            "checks": results,
        }
        try:
            with open(out_path, "w") as f:
                json.dump(report, f, indent=2)
            print(f"Results written to: {os.path.abspath(out_path)}")
        except Exception as e:
            print(f"Could not write results to {out_path}: {e}")

    return 0 if fails == 0 else 1


if __name__ == "__main__":
    sys.exit(main())