"""
test_api.py - Verification script for HRNet 2D Garment Keypoint Backend

Usage:
    python test_api.py --image sample_pant.jpg --output test_keypoints_overlay.jpg
    python test_api.py --url http://127.0.0.1:7860/predict --image sample_pant.jpg
"""

import os
import sys
import json
import argparse
import base64
import io
import requests
import cv2
import numpy as np
from PIL import Image

def test_backend(image_path: str, url: str = "http://127.0.0.1:7860/predict", output_overlay: str = "test_keypoints_overlay.jpg"):
    if not os.path.exists(image_path):
        print(f"[Error] Image not found at {image_path}")
        return False

    print(f"\n=======================================================")
    print(f" Testing HRNet Keypoint API on: {image_path}")
    print(f" Target Endpoint: {url}")
    print(f"=======================================================\n")

    res_json = None

    # Try HTTP POST request first
    try:
        with open(image_path, "rb") as f:
            resp = requests.post(
                url,
                files={"image": ("image.jpg", f, "image/jpeg")},
                data={"debug": "true"},
                timeout=30
            )
        if resp.status_code == 200:
            res_json = resp.json()
            print("[Success] Received response from live HTTP endpoint!")
        else:
            print(f"[Warning] HTTP {resp.status_code}: {resp.text}")
    except Exception as e:
        print(f"[Info] Live server not reachable at {url} ({e})")
        print("[Info] Falling back to direct in-process engine verification...")

    # If live server not running, invoke core engine directly
    if res_json is None:
        try:
            cur_dir = os.path.dirname(os.path.abspath(__file__))
            if cur_dir not in sys.path:
                sys.path.insert(0, cur_dir)
            from app import parse_and_predict
            res_json = parse_and_predict(image_path, debug=True)
            print("[Success] In-process inference completed successfully!")
        except Exception as e:
            print(f"[Error] Direct engine execution failed: {e}")
            return False

    # 1. Validate response schema
    print("\n--- API JSON Response ---")
    # Pretty print without the large debug_image string
    print_dict = dict(res_json)
    if "debug_image" in print_dict:
        print_dict["debug_image"] = f"<base64 data URI length={len(print_dict['debug_image'])}>"
    print(json.dumps(print_dict, indent=2))

    if not res_json.get("ok"):
        print(f"\n[FAILED] Endpoint returned ok=False: {res_json.get('error')}")
        return False

    orig_w = res_json.get("w")
    orig_h = res_json.get("h")
    keypoints = res_json.get("keypoints", [])
    inf_ms = res_json.get("inference_ms")

    print(f"\n--- Validation Checks ---")
    print(f"  [+] 'ok' flag: {res_json.get('ok')}")
    print(f"  [+] Original image dimensions: {orig_w}x{orig_h}")
    print(f"  [+] Keypoint count: {len(keypoints)}")
    print(f"  [+] Inference latency: {inf_ms} ms")

    assert orig_w > 0 and orig_h > 0, "Dimensions must be positive"
    assert len(keypoints) == 14, f"Expected 14 keypoints, got {len(keypoints)}"
    assert isinstance(inf_ms, int), "inference_ms must be an integer"

    for kp in keypoints:
        assert "name" in kp and "u" in kp and "v" in kp and "score" in kp
        assert 0.0 <= kp["u"] <= orig_w, f"u={kp['u']} out of bounds [0, {orig_w}]"
        assert 0.0 <= kp["v"] <= orig_h, f"v={kp['v']} out of bounds [0, {orig_h}]"
        assert 0.0 <= kp["score"] <= 1.0, f"score={kp['score']} out of bounds [0, 1]"

    print("  [+] All schema constraints and coordinate bounds PASSED!")

    # 2. Save visual overlay
    # Use debug_image if present or render directly on image
    if "debug_image" in res_json and res_json["debug_image"]:
        b64_str = res_json["debug_image"].split(",")[-1]
        img_bytes = base64.b64decode(b64_str)
        pil_overlay = Image.open(io.BytesIO(img_bytes))
        pil_overlay.save(output_overlay)
    else:
        orig_img = cv2.imread(image_path)
        for kp in keypoints:
            u, v = int(round(kp["u"])), int(round(kp["v"]))
            cv2.circle(orig_img, (u, v), 6, (0, 0, 255), -1)
            cv2.putText(orig_img, kp["name"], (u + 6, v - 4), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (0, 255, 0), 1, cv2.LINE_AA)
        cv2.imwrite(output_overlay, orig_img)

    print(f"\n[Saved] Visual overlay written to: {os.path.abspath(output_overlay)}")
    print("=======================================================\n")
    return True


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Test HRNet Garment Keypoint API")
    parser.add_argument("--image", type=str, default="sample_pant.jpg", help="Path to test image")
    parser.add_argument("--url", type=str, default="http://127.0.0.1:7860/predict", help="Endpoint URL")
    parser.add_argument("--output", type=str, default="test_keypoints_overlay.jpg", help="Output overlay path")
    args = parser.parse_args()

    # Locate sample image relative to script
    img_target = args.image
    if not os.path.exists(img_target):
        alt = os.path.join(os.path.dirname(os.path.abspath(__file__)), args.image)
        if os.path.exists(alt):
            img_target = alt

    test_backend(img_target, url=args.url, output_overlay=args.output)