from __future__ import annotations

import sys
from pathlib import Path

SCRIPT_DIR = Path(__file__).resolve().parent
sys.path.insert(0, str(SCRIPT_DIR / "vllm"))

from vllm_caption_runtime import build_vllm_request, load_llm, load_processor, sampling_params


# =============================================================================
# Paths
# =============================================================================

VIDEO_PATH = "/workspace/test7.mp4"
MODEL_DIR = str(SCRIPT_DIR / "model")


# =============================================================================
# Prompt Settings
# =============================================================================

PROMPT_OVERRIDE = ""
CAPTION_LENGTH = "very large"
INCLUDE_WATERMARK_INFO = False
HAS_THINKING = True

VULGARITY = "low"
UNCERTAINTY = "low"
CHARACTER_NAMES = "none"
FLUFF = "none"
HAS_REPETITION = False
SPECULATION = "low"
TEMPORAL_DETAIL = "medium"
VISUAL_SPECIFICITY = "moderate"
CAMERA_DETAIL = "medium"
CAPTION_STYLE = "plain"


# =============================================================================
# vLLM / Generation Hyperparameters
# =============================================================================

NUM_FRAMES = 12
SAMPLING_RATE = 16_000
MAX_MODEL_LEN = 4096
MAX_NUM_SEQS = 1
GPU_MEMORY_UTILIZATION = 0.88
DTYPE = "bfloat16"
# Online vLLM FP8 for Gemma linear weights. Parakeet is left unchanged.
USE_FP8 = True
FP8_QUANTIZATION = "fp8_per_tensor"
ENFORCE_EAGER = False
ENABLE_PREFIX_CACHING = False
TRUST_REMOTE_CODE = False

MAX_TOKENS = 1200
TEMPERATURE = 0.0
TOP_P = 0.9
REPETITION_PENALTY = 1.1


def prompt_settings() -> dict[str, object]:
    return {
        "caption_length": CAPTION_LENGTH,
        "include_watermark_info": INCLUDE_WATERMARK_INFO,
        "has_thinking": HAS_THINKING,
        "vulgarity": VULGARITY,
        "uncertainty": UNCERTAINTY,
        "character_names": CHARACTER_NAMES,
        "fluff": FLUFF,
        "has_repetition": HAS_REPETITION,
        "speculation": SPECULATION,
        "temporal_detail": TEMPORAL_DETAIL,
        "visual_specificity": VISUAL_SPECIFICITY,
        "camera_detail": CAMERA_DETAIL,
        "caption_style": CAPTION_STYLE,
    }


def main() -> None:
    quantization = FP8_QUANTIZATION if USE_FP8 else None
    processor = load_processor(MODEL_DIR)
    request = build_vllm_request(
        processor=processor,
        model_dir=MODEL_DIR,
        video_path=VIDEO_PATH,
        num_frames=NUM_FRAMES,
        sampling_rate=SAMPLING_RATE,
        prompt_override=PROMPT_OVERRIDE,
        prompt_settings=prompt_settings(),
    )

    llm = load_llm(
        model_dir=MODEL_DIR,
        max_model_len=MAX_MODEL_LEN,
        max_num_seqs=MAX_NUM_SEQS,
        gpu_memory_utilization=GPU_MEMORY_UTILIZATION,
        dtype=DTYPE,
        quantization=quantization,
        enforce_eager=ENFORCE_EAGER,
        enable_prefix_caching=ENABLE_PREFIX_CACHING,
        trust_remote_code=TRUST_REMOTE_CODE,
    )
    outputs = llm.generate(
        [request],
        sampling_params=sampling_params(
            temperature=TEMPERATURE,
            max_tokens=MAX_TOKENS,
            top_p=TOP_P,
            repetition_penalty=REPETITION_PENALTY,
        ),
    )
    print(outputs[0].outputs[0].text)


if __name__ == "__main__":
    main()