"""Run inside the existing upstream container. Read-only; no /prompt requests."""
import asyncio
import json
from unittest.mock import patch

import app
import httpx
from workflow import compile_workflow, load_registry

WORKFLOWS = {
    "minimax_h3_t2v_official": 0,
    "minimax_h3_i2v": 1,
    "minimax_h3_r2v": 1,
    "minimax_h3_hybrid_v0331_t2v": 0,
    "minimax_h3_hybrid_v0331_i2v": 1,
    "minimax_h3_r2v_4step": 1,
    "ltx23_t2v": 0,
    "ltx23_t2v_text_preserve": 0,
    "wan22_simple": 1,
}


async def main():
    report = {"health": await app.health(), "workers": [], "scheduler_checks": {}}
    registry = load_registry()
    async with httpx.AsyncClient(timeout=30) as client:
        for provider in app.providers().values():
            response = await client.get(provider.base_url + "/object_info", headers=app.provider_headers(provider))
            response.raise_for_status()
            info = response.json()
            checks = {}
            for key, refs in WORKFLOWS.items():
                spec = registry[key]
                request = app.JobRequest(provider_id="default", workflow_key=key, prompt="A mountain landscape. No dialogue.", duration=5, width=832, height=480)
                plan = app.request_audio_plan(spec, request)
                if plan.get("supported"):
                    request = request.model_copy(update={"prompt": plan["normalized_prompt"]})
                graph = compile_workflow(spec, request.model_dump(), ["preflight.png"] * refs, "huobao-preflight", preserve_requested_audio=bool(plan.get("supported")))
                missing_nodes = sorted(set(node["class_type"] for node in graph.values()) - set(info))
                missing_choices = []
                for node in graph.values():
                    node_info = info.get(node["class_type"], {}).get("input", {})
                    choices = {**node_info.get("required", {}), **node_info.get("optional", {})}
                    for field, value in node.get("inputs", {}).items():
                        choice = choices.get(field)
                        if isinstance(value, str) and choice and isinstance(choice[0], list) and value not in choice[0] and field != "image":
                            missing_choices.append([node["class_type"], field, value])
                checks[key] = {"missing_nodes": missing_nodes, "missing_choices": missing_choices, "audio_plan": plan.get("status")}
            report["workers"].append({"id": provider.id, "name": provider.name, "queue_depth": await app.provider_queue_depth(provider), "workflows": checks})
    # Isolated in this process: exercise the live scheduler's failover/reservations
    # without changing the running service or reserving real worker capacity.
    async def depth(provider):
        return depths[provider.id]
    with patch.object(app, "provider_queue_depth", depth):
        depths = {"default": 4, "gpu3": 1}
        report["scheduler_checks"]["shorter_gpu3_queue"] = (await app.select_provider("default")).id
        depths = {"default": 1, "gpu3": 4}
        report["scheduler_checks"]["shorter_gpu2_queue"] = (await app.select_provider("default")).id
        depths = {"default": None, "gpu3": 0}
        report["scheduler_checks"]["gpu2_unavailable"] = (await app.select_provider("default")).id
        depths = {"default": 0, "gpu3": 0}
        first = await app.reserve_provider("default")
        second = await app.reserve_provider("default")
        report["scheduler_checks"]["concurrent_reservations"] = [first.id, second.id]
        await app.release_provider(first)
        await app.release_provider(second)
    print(json.dumps(report, ensure_ascii=False))


asyncio.run(main())