"""
Usage:

python hidden_states_server.py

CUDA graphs use the configured maximum hidden-state mode. Requests may select
that mode or a weaker one without triggering mode-dependent recapture.
"""

import requests
import torch

from sglang.test.test_utils import is_in_ci
from sglang.utils import terminate_process, wait_for_server

if is_in_ci():
    from docs.backend.patch import launch_server_cmd
else:
    from sglang.utils import launch_server_cmd


def main():
    # Launch the server
    server_process, port = launch_server_cmd(
        "python -m sglang.launch_server --model-path "
        "Alibaba-NLP/gte-Qwen2-1.5B-instruct "
        "--return-hidden-states-mode last --host 0.0.0.0"
    )
    wait_for_server(f"http://localhost:{port}", process=server_process)

    prompts = [
        "Hello, my name is",
        "The president of the United States is",
        "The capital of France is",
        "The future of AI is",
    ]

    sampling_params = {
        "temperature": 0.8,
        "top_p": 0.95,
        "max_new_tokens": 10,
    }

    json_data = {
        "text": prompts,
        "sampling_params": sampling_params,
        "return_hidden_states": "last",
    }

    response = requests.post(
        f"http://localhost:{port}/generate",
        json=json_data,
    )

    terminate_process(server_process)

    outputs = response.json()
    for prompt, output in zip(prompts, outputs):
        hidden_state = torch.tensor(
            output["meta_info"]["hidden_states"], dtype=torch.bfloat16
        )
        print("===============================")
        print(
            f"Prompt: {prompt}\n"
            f"Generated text: {output['text']}\n"
            f"Prompt_Tokens: {output['meta_info']['prompt_tokens']}\t"
            f"Completion_tokens: {output['meta_info']['completion_tokens']}"
        )
        print("Last hidden state: ")
        print(hidden_state)
        print()


if __name__ == "__main__":
    main()