#!/usr/bin/env python3
"""
推論サーバーAPIのテストスクリプト

Usage:
    # サーバーの状態確認
    python test_server_api.py --status

    # マスク生成のテスト
    python test_server_api.py --video test.mp4 --output /tmp/masks --start 0 --end 10
"""

import argparse
import json
import time
import urllib.request
import urllib.error
from pathlib import Path
import sys


SERVER_URL = "http://127.0.0.1:8181"


def check_status():
    """サーバーの状態を確認"""
    try:
        with urllib.request.urlopen(f"{SERVER_URL}/status", timeout=2) as response:
            data = json.loads(response.read().decode('utf-8'))
            print("✓ サーバーは稼働中です")
            print(f"  Status: {data.get('status')}")
            print(f"  GPU Available: {data.get('gpu_available')}")
            if data.get('gpu_device'):
                print(f"  GPU Device: {data.get('gpu_device')}")
            if data.get('gpu_count'):
                print(f"  GPU Count: {data.get('gpu_count')}")
            if data.get('rocm_version'):
                print(f"  ROCm Version: {data.get('rocm_version')}")
            return True
    except (urllib.error.URLError, ConnectionRefusedError, TimeoutError) as e:
        print("✗ サーバーに接続できません")
        print(f"  エラー: {e}")
        print("\nサーバーを起動してください:")
        print("  ./run_server.sh")
        return False


def submit_task(video_path, output_dir, start_frame, end_frame, conf, iou, mask_scale):
    """マスク生成タスクを送信"""
    data = {
        "video_path": video_path,
        "output_dir": output_dir,
        "start_frame": start_frame,
        "end_frame": end_frame,
        "conf_threshold": conf,
        "iou_threshold": iou,
        "mask_scale": mask_scale,
    }

    req = urllib.request.Request(
        f"{SERVER_URL}/generate",
        data=json.dumps(data).encode('utf-8'),
        headers={'Content-Type': 'application/json'},
        method='POST'
    )

    try:
        with urllib.request.urlopen(req) as response:
            result = json.loads(response.read().decode('utf-8'))
            return result
    except urllib.error.HTTPError as e:
        error_msg = e.read().decode('utf-8')
        raise RuntimeError(f"サーバーエラー: {error_msg}")


def get_task_status(task_id):
    """タスクの状態を取得"""
    try:
        with urllib.request.urlopen(f"{SERVER_URL}/tasks/{task_id}") as response:
            return json.loads(response.read().decode('utf-8'))
    except urllib.error.HTTPError:
        return {"status": "unknown"}


def cancel_task(task_id):
    """タスクをキャンセル"""
    try:
        req = urllib.request.Request(
            f"{SERVER_URL}/tasks/{task_id}/cancel",
            method='POST'
        )
        with urllib.request.urlopen(req):
            pass
        return True
    except urllib.error.HTTPError:
        return False


def monitor_task(task_id, poll_interval=0.5):
    """タスクの進行状況を監視"""
    print(f"\nタスクID: {task_id}")
    print("進行状況を監視中...\n")

    last_progress = -1

    while True:
        status = get_task_status(task_id)
        state = status.get('status')
        progress = status.get('progress', 0)
        total = status.get('total', 0)

        # 進行状況の表示
        if progress != last_progress and total > 0:
            percentage = (progress / total) * 100
            bar_length = 40
            filled = int(bar_length * progress / total)
            bar = '=' * filled + '-' * (bar_length - filled)
            print(f"\r[{bar}] {progress}/{total} ({percentage:.1f}%)", end='', flush=True)
            last_progress = progress

        # 終了状態のチェック
        if state == "completed":
            print("\n\n✓ 処理が完了しました")
            print(f"  出力先: {status.get('result_path')}")
            print(f"  メッセージ: {status.get('message')}")
            return True

        elif state == "failed":
            print("\n\n✗ 処理が失敗しました")
            print(f"  エラー: {status.get('message')}")
            return False

        elif state == "cancelled":
            print("\n\n- 処理がキャンセルされました")
            return False

        time.sleep(poll_interval)


def main():
    parser = argparse.ArgumentParser(
        description="推論サーバーAPIのテストスクリプト"
    )

    # 操作モード
    mode_group = parser.add_mutually_exclusive_group(required=True)
    mode_group.add_argument("--status", action="store_true", help="サーバーの状態を確認")
    mode_group.add_argument("--video", type=str, help="処理する動画ファイル")

    # タスクパラメータ
    parser.add_argument("--output", type=str, default="/tmp/masks", help="マスク出力先ディレクトリ")
    parser.add_argument("--start", type=int, default=0, help="開始フレーム")
    parser.add_argument("--end", type=int, default=10, help="終了フレーム")
    parser.add_argument("--conf", type=float, default=0.5, help="信頼度閾値")
    parser.add_argument("--iou", type=float, default=0.45, help="NMS IoU閾値")
    parser.add_argument("--mask-scale", type=float, default=1.5, help="マスクスケール")

    # その他のオプション
    parser.add_argument("--no-wait", action="store_true", help="タスク送信後、完了を待たない")

    args = parser.parse_args()

    # 状態確認モード
    if args.status:
        check_status()
        return

    # マスク生成モード
    print("=== 推論サーバーAPIテスト ===\n")

    # サーバーの確認
    if not check_status():
        sys.exit(1)

    # 動画ファイルの確認
    if not Path(args.video).exists():
        print(f"\n✗ 動画ファイルが見つかりません: {args.video}")
        sys.exit(1)

    video_path = str(Path(args.video).absolute())
    output_dir = str(Path(args.output).absolute())

    print(f"\n動画: {video_path}")
    print(f"出力先: {output_dir}")
    print(f"フレーム範囲: {args.start} - {args.end}")
    print(f"パラメータ: conf={args.conf}, iou={args.iou}, mask_scale={args.mask_scale}")

    # タスクを送信
    print("\nタスクを送信中...")
    try:
        result = submit_task(
            video_path,
            output_dir,
            args.start,
            args.end,
            args.conf,
            args.iou,
            args.mask_scale
        )
    except Exception as e:
        print(f"\n✗ タスク送信失敗: {e}")
        sys.exit(1)

    task_id = result['id']
    print(f"✓ タスクが送信されました (ID: {task_id})")

    # 完了待機
    if not args.no_wait:
        try:
            success = monitor_task(task_id)
            sys.exit(0 if success else 1)
        except KeyboardInterrupt:
            print("\n\n中断されました")
            print("タスクをキャンセル中...")
            if cancel_task(task_id):
                print("✓ タスクをキャンセルしました")
            sys.exit(130)
    else:
        print("\nタスクの状態を確認するには:")
        print(f"  python test_server_api.py --task-status {task_id}")


if __name__ == "__main__":
    main()