#!/usr/bin/env python3
"""
Test script — connects via WebSocket, runs a diagnostic, fixes a few issues.
Works against local server or live HF Space.

Usage:
    python3 test_live.py                           # HF Space
    python3 test_live.py http://localhost:8000      # local
"""

import asyncio
import json
import sys

try:
    import websockets
except ImportError:
    import subprocess
    subprocess.check_call([sys.executable, "-m", "pip", "install", "websockets"])
    import websockets

BASE = sys.argv[1] if len(sys.argv) > 1 else "https://devaatmik-shopify-store-audit.hf.space"
WS_URL = BASE.replace("https://", "wss://").replace("http://", "ws://") + "/ws"


def parse_ws(raw: dict) -> dict:
    data = raw.get("data", {})
    obs = data.get("observation", {})
    obs["reward"] = data.get("reward", 0)
    obs["done"] = data.get("done", False)
    return obs


async def main():
    print(f"Connecting to {WS_URL} ...\n")
    async with websockets.connect(WS_URL) as ws:
        await ws.send(json.dumps({"type": "reset", "data": {"task": "product_listing_qa"}}))
        obs = parse_ws(json.loads(await ws.recv()))
        print(f"  RESET | Task: {obs.get('task_name')} | Issues: {obs.get('issues_remaining')}/{obs.get('total_issues')}")
        print(f"       | Products: {obs.get('data', {}).get('product_count', '?')}")
        print()

        actions = [
            {"command": "query_store_health", "params": {}},
            {"command": "query_products", "params": {"limit": 5}},
        ]

        for i, action in enumerate(actions, 1):
            await ws.send(json.dumps({"type": "step", "data": action}))
            obs = parse_ws(json.loads(await ws.recv()))
            fixed = obs.get("issues_fixed", 0)
            total = obs.get("total_issues", 0)
            health = obs.get("store_health_score", 0)
            reward = obs.get("reward", 0)
            print(f"  Step {i:2d} | {action['command']:30s} | health {health:5.1%} | fixed {fixed}/{total} | reward {reward:+.3f}")

            if action["command"] == "query_store_health":
                issues = obs.get("data", {}).get("issues", [])
                if issues:
                    print(f"         | Issues found: {len(issues)}")
                    for iss in issues[:3]:
                        desc = iss.get("description", "")[:80]
                        print(f"         |   - {desc}")
                cats = obs.get("data", {}).get("issues_by_category", {})
                if cats:
                    print(f"         | Categories: {cats}")

        await ws.send(json.dumps({"type": "state"}))
        state = json.loads(await ws.recv())
        print(f"\n  State: {state.get('data', {})}")
        print("\n  Test complete.")


if __name__ == "__main__":
    asyncio.run(main())