"""HPD-Parsing throughput (TPS) benchmark on vLLM.

Batched offline inference over an image folder; reports TPS metrics and dumps
the raw predictions consumed by hpd_to_markdown.py (model is run once).

    MAX_PATCHES_WITH_RESIZE=true python eval/benchmark_tps.py
"""
import os

os.environ.setdefault("MAX_PATCHES_WITH_RESIZE", "true")  # before importing vLLM

import base64
import json
import time
from typing import Any, Dict

from vllm import LLM, SamplingParams
from vllm.sampling_params import RepetitionDetectionParams

PRED_OUT = "batch_512_pred_HPD-Parsing.json"

llm = None
sampling_params = None


def encode_image_to_base64(image_path: str) -> str:
    with open(image_path, "rb") as f:
        return base64.b64encode(f.read()).decode("utf-8")


def batch_predict_vllm(tasks, batch_size=512):
    all_outputs = []
    records = []
    for start in range(0, len(tasks), batch_size):
        batch_tasks = tasks[start:start + batch_size]
        messages_list = [
            [{
                "role": "user",
                "content": [
                    {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{img_b64}"}},
                    {"type": "text", "text": question},
                ],
            }]
            for question, img_b64, _ in batch_tasks
        ]
        outputs = llm.chat(messages=messages_list, sampling_params=sampling_params)
        all_outputs.extend(outputs)
        for j, out in enumerate(outputs):
            records.append({
                "index": start + j,
                "img_path": batch_tasks[j][-1],
                "pred": out.outputs[0].text,
            })
    return all_outputs, records


def parse_ckpt_name(model_path: str) -> str:
    parts = model_path.rstrip('/').split('/')
    return '_'.join(p for p in parts if p and p not in ['checkpoints', 'output', 'rank_0'])


def record_metrics(model_path: str, metrics: Dict[str, Any], output_dir: str = "./records"):
    os.makedirs(output_dir, exist_ok=True)
    output_file = os.path.join(output_dir, f"{parse_ckpt_name(model_path)}.txt")
    with open(output_file, 'w', encoding='utf-8') as f:
        f.write(f"Model Path: {model_path}\n")
        f.write(f"Checkpoint: {parse_ckpt_name(model_path)}\n")
        f.write("=" * 50 + "\n")
        for key, value in metrics.items():
            f.write(f"{key}: {value}\n")
    print(f"Metrics saved to: {output_file}")


if __name__ == '__main__':
    model_path = "PaddlePaddle/HPD-Parsing/"
    model_path_medusa = "PaddlePaddle/HPD-Parsing/P-MTP/"
    root = "OmniDocBench_1_6/images/"
    prompt = "document parsing with fork."
    batch_size = 512

    llm = LLM(
        model=model_path,
        tensor_parallel_size=1,
        gpu_memory_utilization=0.9,
        enable_prefix_caching=True,
        attention_backend="FLASHINFER",
        trust_remote_code=True,
        max_model_len=24000,
        max_num_seqs=512,
        speculative_config={"method": "medusa", "model": model_path_medusa, "num_speculative_tokens": 6},
    )
    sampling_params = SamplingParams(
        temperature=0,
        max_tokens=8000,
        stop=None,
        repetition_detection=RepetitionDetectionParams(
            min_pattern_size=64,
            max_pattern_size=128,
            min_count=10,
        ),
    )

    valid_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.gif', '.webp'}
    image_paths = [
        os.path.join(root, p) for p in os.listdir(root)
        if os.path.splitext(p)[1].lower() in valid_extensions
    ]
    tasks = [(prompt, encode_image_to_base64(p), p) for p in image_paths]
    print(f"Total images: {len(tasks)}")

    start_time = time.perf_counter()
    outputs, records = batch_predict_vllm(tasks, batch_size=batch_size)
    end_time = time.perf_counter()

    with open(PRED_OUT, "w", encoding="utf-8") as fp:
        json.dump(records, fp, indent=4, ensure_ascii=False)
    print(f"Predictions saved to: {PRED_OUT}")

    total_time = end_time - start_time
    total_input_tokens = sum(len(out.prompt_token_ids) for out in outputs)
    total_output_tokens = sum(len(out.outputs[0].token_ids) for out in outputs)

    print(f"Total Time: {total_time:.2f} s")
    print(f"Throughput (Requests/s): {len(tasks) / total_time:.2f}")
    print(f"Input Tokens/s: {total_input_tokens / total_time:.2f}")
    print(f"Output Tokens/s: {total_output_tokens / total_time:.2f}")
    print(f"Total Tokens/s: {(total_input_tokens + total_output_tokens) / total_time:.2f}")
    print(f"Avg Input Tokens per Request: {total_input_tokens / len(tasks):.2f}")
    print(f"Avg Output Tokens per Request: {total_output_tokens / len(tasks):.2f}")

    metrics = {
        "Total Time (s)": f"{total_time:.2f}",
        "Throughput (Requests/s)": f"{len(tasks) / total_time:.2f}",
        "Input Tokens/s": f"{total_input_tokens / total_time:.2f}",
        "Output Tokens/s": f"{total_output_tokens / total_time:.2f}",
        "Total Tokens/s": f"{(total_input_tokens + total_output_tokens) / total_time:.2f}",
        "Avg Input Tokens per Request": f"{total_input_tokens / len(tasks):.2f}",
        "Avg Output Tokens per Request": f"{total_output_tokens / len(tasks):.2f}",
        "Total Requests": len(tasks),
        "Total Input Tokens": total_input_tokens,
        "Total Output Tokens": total_output_tokens,
    }
    record_metrics(model_path, metrics)