import argparse
import json
import os
import sys
from pathlib import Path
from urllib.parse import urlparse


PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))


def require_modal():
    try:
        import modal  # noqa: F401
    except Exception as exc:
        raise SystemExit(
            "Modal Python package is not installed. Install with `python -m pip install -r requirements.txt`, "
            "then run `modal setup` to log in."
        ) from exc


def check_remote_methods(service: str = "all") -> int:
    require_modal()
    from modal_apps.modal_image import CharacterImage, app as image_app
    from modal_apps.modal_llm import PersonaLLM, app as llm_app
    from modal_apps.modal_tts import CharacterTTS, app as tts_app

    all_checks = [
        ("llm", llm_app, lambda: PersonaLLM().health.remote()),
        ("tts", tts_app, lambda: CharacterTTS().health.remote()),
        ("image", image_app, lambda: CharacterImage().health.remote()),
    ]
    checks = [check for check in all_checks if service in {"all", check[0]}]
    ok = True
    for name, modal_app, fn in checks:
        try:
            print(f"[{name}] checking health...")
            with modal_app.run():
                result = fn()
            print(json.dumps(result, ensure_ascii=False, indent=2))
        except Exception as exc:
            ok = False
            print(f"[{name}] failed: {exc}", file=sys.stderr)
    return 0 if ok else 1


def check_deployed_endpoints() -> int:
    try:
        import httpx
    except Exception as exc:
        raise SystemExit("httpx is not installed. Run `python -m pip install -r requirements.txt`.") from exc

    llm_url = os.environ.get("VC_MODAL_LLM_URL")
    tts_url = os.environ.get("VC_MODAL_TTS_URL")
    image_url = os.environ.get("VC_MODAL_IMAGE_URL")
    ok = True

    if llm_url:
        print("[llm:endpoint] checking SSE...")
        with httpx.stream(
            "POST",
            llm_url,
            json={"text": "你好，请简单回复一句。", "character": {"display_name": "星萤"}, "max_new_tokens": 40},
            timeout=120,
        ) as response:
            response.raise_for_status()
            seen = 0
            for line in response.iter_lines():
                if line:
                    print(line)
                    seen += 1
                if seen >= 5:
                    break
    else:
        print("[llm:endpoint] skipped; set VC_MODAL_LLM_URL")

    if tts_url:
        print("[tts:endpoint] checking audio...")
        response = httpx.post(
            _tts_endpoint_url(tts_url),
            json={"text": "你好，我在听。", "voice_id": "default", "emotion": "neutral"},
            timeout=120,
            trust_env=False,
        )
        response.raise_for_status()
        out = Path("modal_tts_check.wav")
        out.write_bytes(response.content)
        print(f"wrote {out} ({len(response.content)} bytes)")
    else:
        print("[tts:endpoint] skipped; set VC_MODAL_TTS_URL")

    if image_url:
        print("[image:endpoint] checking png...")
        response = httpx.post(
            image_url,
            json={"prompt": "original anime virtual character portrait, gentle sci-fi style", "steps": 4, "seed": 7},
            timeout=180,
        )
        response.raise_for_status()
        out = Path("modal_image_check.png")
        out.write_bytes(response.content)
        print(f"wrote {out} ({len(response.content)} bytes)")
    else:
        print("[image:endpoint] skipped; set VC_MODAL_IMAGE_URL")

    return 0 if ok else 1


def _health_targets(url: str) -> list[str]:
    base = url.rstrip("/")
    tail = base.rsplit("/", 1)[-1]
    service_base = base.rsplit("/", 1)[0] if tail in {"tts", "persona_events"} else base
    return list(dict.fromkeys([service_base + "/health", service_base + "/health_http", base + "/health"]))


def _tts_endpoint_url(url: str) -> str:
    base = url.rstrip("/")
    parsed = urlparse(base)
    if not parsed.path or parsed.path == "/":
        return base
    if parsed.path.rstrip("/").rsplit("/", 1)[-1] == "tts":
        return base
    return base + "/tts"


def main() -> int:
    parser = argparse.ArgumentParser(description="Check Modal model connectivity for Virtual Characters.")
    parser.add_argument(
        "--mode",
        choices=["remote-methods", "endpoints"],
        default="remote-methods",
        help="remote-methods checks Modal class methods; endpoints checks deployed web URLs.",
    )
    parser.add_argument(
        "--service",
        choices=["all", "llm", "tts", "image"],
        default="all",
        help="Limit remote-methods checks to one service.",
    )
    args = parser.parse_args()
    if args.mode == "remote-methods":
        return check_remote_methods(args.service)
    if args.service != "all":
        raise SystemExit("--service is only supported with --mode remote-methods")
    return check_deployed_endpoints()


if __name__ == "__main__":
    raise SystemExit(main())