"""Shared utilities for component tests.

Component tests run inside the Docker image against real GPU models. They
write their output artefacts (MP4s, PNGs, logs) to ``_out/`` so you can
visually inspect results.
"""
from __future__ import annotations

import logging
import os
import sys

import numpy as np


OUT_DIR = os.path.join(os.path.dirname(__file__), "_out")
os.makedirs(OUT_DIR, exist_ok=True)

# A tiny 256x256 portrait PNG lives next to the component tests so tests
# don't need a user-supplied file. If it's missing we synthesise one on
# the fly.
SAMPLE_AVATAR = os.path.join(os.path.dirname(__file__), "sample_avatar.png")


def get_logger(name: str) -> logging.Logger:
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s %(name)s %(levelname)s %(message)s",
        stream=sys.stdout,
    )
    return logging.getLogger(name)


def ensure_sample_avatar() -> str:
    """Guarantee a usable avatar image exists. Returns its path."""
    if os.path.isfile(SAMPLE_AVATAR):
        return SAMPLE_AVATAR
    # Synthesise a simple gradient PNG as a last resort (won't look like a
    # person but is valid input for Wan2.2 so the pipeline doesn't fail).
    try:
        from PIL import Image  # type: ignore[import-not-found]
    except ImportError:
        import imageio.v3 as iio  # type: ignore[import-not-found]
        arr = np.zeros((256, 256, 3), dtype=np.uint8)
        for y in range(256):
            arr[y, :, 0] = y
            arr[y, :, 1] = 255 - y
            arr[y, :, 2] = 128
        iio.imwrite(SAMPLE_AVATAR, arr)
        return SAMPLE_AVATAR

    arr = np.zeros((256, 256, 3), dtype=np.uint8)
    for y in range(256):
        arr[y, :, 0] = y
        arr[y, :, 1] = 255 - y
        arr[y, :, 2] = 128
    Image.fromarray(arr).save(SAMPLE_AVATAR)
    return SAMPLE_AVATAR


def write_bytes(name: str, data: bytes) -> str:
    """Write an artefact to _out/<name> and return the full path."""
    path = os.path.join(OUT_DIR, name)
    with open(path, "wb") as f:
        f.write(data)
    return path


def synth_tone(seconds: float, sample_rate: int = 24000, freq: float = 220.0) -> np.ndarray:
    """Return a float32 sine tone usable as stand-in TTS audio."""
    t = np.arange(int(seconds * sample_rate), dtype=np.float32) / sample_rate
    return (0.2 * np.sin(2 * np.pi * freq * t)).astype(np.float32)