import os
import time
import urllib

import modal
from suno_utils.worker.modal_base import get_modal_base_image

vllm_image = (
    get_modal_base_image()
    .pip_install(
        "vllm==0.7.2",
        "huggingface_hub[hf_transfer]==0.26.2",
        "flashinfer-python==0.2.0.post2",  # pinning, very unstable
        extra_index_url="https://flashinfer.ai/whl/cu124/torch2.5",
    )
    .env({"HF_HUB_ENABLE_HF_TRANSFER": "1"})  # faster model transfers
    .add_local_python_source("suno_utils", copy=False)
)


MODEL_DIR = "llamas"

hf_cache_vol = modal.Volume.from_name("huggingface-cache", create_if_missing=True)
vllm_cache_vol = modal.Volume.from_name("vllm-cache", create_if_missing=True)

llamas_vol = modal.Volume.from_name(MODEL_DIR, create_if_missing=False)
app = modal.App("lyrics-gen")
ADAPTER_PATH = f"/root/{MODEL_DIR}/llama3-1-8b-lora-dpo-beta0p01-gradient-acc-lr3/epoch_2"
N_GPU = 1  # tip: for best results, first upgrade to more powerful GPUs, and only then increase GPU count
MAX_LORA_RANK = 32
MINUTES = 60  # seconds

VLLM_PORT = 8000

BASE_MODEL_NAME = "meta-llama/Llama-3.1-8B-Instruct"
FT_MODEL_NAME = "my-remi-8b-1"
SERVER_API_KEY = "sunosunosuno"


@app.function(
    image=vllm_image,
    gpu=f"H100:{N_GPU}",
    # how long should we stay up with no requests?
    scaledown_window=15 * MINUTES,
    volumes={
        "/root/.cache/huggingface": hf_cache_vol,
        "/root/.cache/vllm": vllm_cache_vol,
        "/root/llamas": llamas_vol,
    },
    secrets=[modal.Secret.from_name("huggingface-secret-suno")],
)
@modal.concurrent(max_inputs=100)
@modal.web_server(port=VLLM_PORT, startup_timeout=5 * MINUTES)
def serve():
    import subprocess

    cmd = [
        "vllm",
        "serve",
        BASE_MODEL_NAME,
        "--enable-lora",
        "--lora-modules",
        f"{FT_MODEL_NAME}={ADAPTER_PATH}",
        "--api-key",
        SERVER_API_KEY,
        "--max-lora-rank",
        str(MAX_LORA_RANK),
        "--host",
        "0.0.0.0",
        "--port",
        str(VLLM_PORT),
    ]
    subprocess.Popen(" ".join(cmd), shell=True)


def perform_health_check(url, test_timeout=1 * MINUTES):
    print(f"Running health check for server at {url}")
    up, start, delay = False, time.time(), 10
    full_url = str(os.path.join(url, "health"))
    print("hitting:", full_url)
    while not up:
        try:
            with urllib.request.urlopen(full_url) as response:
                if response.getcode() == 200:
                    up = True
        except Exception as e:
            print(f"Got exception: {e} while performing server health check")
            if time.time() - start > test_timeout:
                break
            time.sleep(delay)

    assert up, f"Failed health check for server at {url}"

    print(f"Successful health check for server at {url}")


@app.local_entrypoint()
def test(test_timeout=5 * MINUTES):
    import json
    import urllib

    import openai

    perform_health_check(serve.web_url)

    messages = [
        {
            "role": "system",
            "content": "You are an expert songwriter.  Write a song based on the following prompt:",
        },
        {"role": "user", "content": "a heart-rending country song about a literal banana"},
    ]
    print(f"Sending a sample message to {serve.web_url}", *messages, sep="\n")

    headers = {
        "Authorization": f"Bearer {SERVER_API_KEY}",
        "Content-Type": "application/json",
    }
    payload = json.dumps({"messages": messages, "model": FT_MODEL_NAME})
    req = urllib.request.Request(
        serve.web_url + "/v1/chat/completions",
        data=payload.encode("utf-8"),
        headers=headers,
        method="POST",
    )
    with urllib.request.urlopen(req) as response:
        payload = json.loads(response.read().decode())
        lyrics = payload["choices"][0]["message"]["content"]
        print(payload)
        print(lyrics)

    client = openai.OpenAI(base_url=str(serve.web_url) + "/v1", api_key=SERVER_API_KEY)
    resp = client.chat.completions.create(
        model=FT_MODEL_NAME,
        messages=[
            {"role": "system", "content": "You are a cool assistant"},
            {"role": "user", "content": "Hi, is this thing on?"},
        ],
    )
    print(resp)