#!/usr/bin/env python3
"""
Smoke test for video pricing table billing.

Required environment:
  NEW_API_BASE_URL=http://127.0.0.1:3000
  NEW_API_ADMIN_USERNAME=root
  NEW_API_ADMIN_PASSWORD=123456
  NEW_API_TOKEN=sk-...

Optional model overrides:
  VIDEO_PRICING_HAILUO_MODEL=hailuo-video
  VIDEO_PRICING_SEEDANCE_MODEL=seedance-2

The server must already have usable channels for the selected models. For the
per_1m_tokens case, point the Seedance/Doubao channel at a mock upstream that
returns a completed task with usage.total_tokens.
"""

from __future__ import annotations

import json
import os
import sys
import time
import urllib.error
import urllib.parse
import urllib.request
from http.cookiejar import CookieJar


BASE_URL = os.getenv("NEW_API_BASE_URL", "http://127.0.0.1:3000").rstrip("/")
ADMIN_USERNAME = os.getenv("NEW_API_ADMIN_USERNAME", "root")
ADMIN_PASSWORD = os.getenv("NEW_API_ADMIN_PASSWORD", "123456")
API_TOKEN = os.getenv("NEW_API_TOKEN", "")
HAILUO_MODEL = os.getenv("VIDEO_PRICING_HAILUO_MODEL", "hailuo-video")
SEEDANCE_MODEL = os.getenv("VIDEO_PRICING_SEEDANCE_MODEL", "seedance-2")

QUOTA_PER_UNIT = int(os.getenv("VIDEO_PRICING_QUOTA_PER_UNIT", "500000"))
GROUP_RATIO = float(os.getenv("VIDEO_PRICING_GROUP_RATIO", "1"))


class Client:
    def __init__(self) -> None:
        self.cookies = CookieJar()
        self.opener = urllib.request.build_opener(urllib.request.HTTPCookieProcessor(self.cookies))

    def request(self, method: str, path: str, body=None, token: str | None = None, expected: int | None = None):
        url = f"{BASE_URL}{path}"
        data = None
        headers = {"Content-Type": "application/json"}
        if body is not None:
            data = json.dumps(body).encode("utf-8")
        if token:
            headers["Authorization"] = f"Bearer {token}"
        req = urllib.request.Request(url, data=data, headers=headers, method=method)
        try:
            with self.opener.open(req, timeout=30) as resp:
                payload = read_json(resp)
                if expected is not None and resp.status != expected:
                    raise AssertionError(f"{method} {path}: expected {expected}, got {resp.status}: {payload}")
                return resp.status, payload
        except urllib.error.HTTPError as err:
            payload = read_json(err)
            if expected is not None and err.code == expected:
                return err.code, payload
            raise AssertionError(f"{method} {path}: HTTP {err.code}: {payload}") from err


def read_json(resp):
    raw = resp.read().decode("utf-8", errors="replace")
    if not raw:
        return None
    try:
        return json.loads(raw)
    except json.JSONDecodeError:
        return raw


def assert_success(payload, context: str):
    if not isinstance(payload, dict) or payload.get("success") is not True:
        raise AssertionError(f"{context} failed: {payload}")
    return payload.get("data")


def usage(client: Client) -> dict:
    _, payload = client.request("GET", "/api/usage/token/", token=API_TOKEN, expected=200)
    if not payload.get("code"):
        raise AssertionError(f"usage failed: {payload}")
    return payload["data"]


def used_quota(client: Client) -> int:
    return int(usage(client)["total_used"])


def login(client: Client) -> None:
    _, payload = client.request(
        "POST",
        "/api/user/login",
        {"username": ADMIN_USERNAME, "password": ADMIN_PASSWORD},
        expected=200,
    )
    assert_success(payload, "admin login")


def put_model_pricing(client: Client, model: str, config: dict) -> None:
    _, payload = client.request("PUT", f"/api/option/model_pricing/{urllib.parse.quote(model, safe='')}", config, expected=200)
    assert_success(payload, f"put pricing {model}")


def delete_model_pricing(client: Client, model: str) -> None:
    _, payload = client.request("DELETE", f"/api/option/model_pricing/{urllib.parse.quote(model, safe='')}", expected=200)
    assert_success(payload, f"delete pricing {model}")


def submit_video(client: Client, body: dict, expected: int = 200):
    _, payload = client.request("POST", "/v1/videos", body, token=API_TOKEN, expected=expected)
    return payload


def fetch_video(client: Client, task_id: str):
    _, payload = client.request("GET", f"/v1/videos/{urllib.parse.quote(task_id, safe='')}", token=API_TOKEN, expected=200)
    return payload


def hailuo_config() -> dict:
    return {
        "schema_version": 1,
        "scope": "model",
        "billing_unit": "per_call",
        "preconsume_strategy": "exact",
        "dimensions": [
            {"key": "resolution", "source": "request.resolution", "type": "string", "default": None},
            {"key": "duration", "source": "request.duration", "type": "number", "default": None},
        ],
        "table": [
            {"resolution": "768P", "duration": 10, "price": 0.04, "source": "manual"},
        ],
        "fallback": {"strategy": "reject"},
    }


def seedance_config() -> dict:
    return {
        "schema_version": 1,
        "scope": "model",
        "billing_unit": "per_1m_tokens",
        "preconsume_strategy": "minimum",
        "dimensions": [
            {"key": "resolution", "source": "request.resolution", "type": "string", "default": None},
        ],
        "table": [
            {"resolution": "720P", "price": 0.5, "source": "manual"},
        ],
        "fallback": {"strategy": "reject"},
    }


def expected_quota(price_usd: float, tokens: int | None = None) -> int:
    units = tokens if tokens is not None else 1
    divisor = 1_000_000 if tokens is not None else 1
    quota = int(float(units) * price_usd / divisor * QUOTA_PER_UNIT * GROUP_RATIO)
    if tokens and quota <= 0 and price_usd > 0 and GROUP_RATIO > 0:
        return 1
    return quota


def run() -> None:
    if not API_TOKEN:
        raise RuntimeError("NEW_API_TOKEN is required")

    client = Client()
    login(client)

    try:
        put_model_pricing(client, HAILUO_MODEL, hailuo_config())
        before = used_quota(client)
        payload = submit_video(
            client,
            {"model": HAILUO_MODEL, "prompt": "pricing smoke", "resolution": "768P", "duration": 10},
        )
        after = used_quota(client)
        assert after - before == expected_quota(0.04), f"Hailuo quota mismatch: before={before}, after={after}, response={payload}"

        before_reject = used_quota(client)
        rejected = submit_video(
            client,
            {"model": HAILUO_MODEL, "prompt": "pricing smoke", "resolution": "1080P", "duration": 10},
            expected=400,
        )
        after_reject = used_quota(client)
        assert after_reject == before_reject, f"reject changed quota: {before_reject} -> {after_reject}, response={rejected}"

        put_model_pricing(client, SEEDANCE_MODEL, seedance_config())
        before_usage = used_quota(client)
        seedance_submit = submit_video(
            client,
            {"model": SEEDANCE_MODEL, "prompt": "pricing smoke", "resolution": "720P", "duration": 5},
        )
        task_id = seedance_submit.get("id") or seedance_submit.get("task_id") or seedance_submit.get("data", {}).get("id")
        if not task_id:
            raise AssertionError(f"Seedance submit did not return task id: {seedance_submit}")

        total_tokens = int(os.getenv("VIDEO_PRICING_EXPECTED_TOTAL_TOKENS", "1000"))
        for _ in range(12):
            fetched = fetch_video(client, task_id)
            status = fetched.get("status") or fetched.get("data", {}).get("status")
            if status in {"succeeded", "completed", "success"}:
                break
            time.sleep(1)
        after_usage = used_quota(client)
        delta = after_usage - before_usage
        assert delta == expected_quota(0.5, total_tokens), f"Seedance usage quota mismatch: delta={delta}, response={fetched}"
    finally:
        delete_model_pricing(client, HAILUO_MODEL)
        delete_model_pricing(client, SEEDANCE_MODEL)


if __name__ == "__main__":
    try:
        run()
    except Exception as exc:
        print(f"FAIL: {exc}", file=sys.stderr)
        sys.exit(1)
    print("PASS: video pricing smoke completed")