#!/usr/bin/env python3
import subprocess
import time
import requests
import json
import os
import sys
import argparse
from datetime import datetime
import csv
import random

# 配置常量
LOG_FILE = "llama_perf_logs.csv"
CONTENT_FILE = "llama_content_logs.csv"
MAX_RETRIES = 10
RETRY_DELAY = 15

def init_log_file():
    """初始化日志文件，写入CSV表头（如果文件不存在）"""
    if not os.path.exists(LOG_FILE):
        with open(LOG_FILE, "w") as f:
            f.write("timestamp,model,prompt_tokens,completion_tokens,memory_usage,Decode_TPS,TTFT,Prefill_TPS,Power\n")
        print(f"Initialized log file: {LOG_FILE}")
    if not os.path.exists(CONTENT_FILE):
        with open(CONTENT_FILE, "w") as f:
            f.write("timestamp,model,prompt,content\n")
        print(f"Initialized content file: {CONTENT_FILE}")

def load_test_cases(test_file):
    """从JSON文件加载测试用例"""
    try:
        with open(test_file, 'r', encoding='utf-8') as f:
            data = json.load(f)
        
        meta = data.get("meta_info", {})
        test_cases = data.get("test_cases", [])
        
        print(f"filename: {test_file}")
        print(f"version: {meta.get('version', 'N/A')}")
        print(f"find {len(test_cases)} test cases")
        
        return test_cases
    except FileNotFoundError:
        print(f"error: file not found {test_file}")
        sys.exit(1)
    except json.JSONDecodeError as e:
        print(f"error: JSON decode failed - {e}")
        sys.exit(1)

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-t", "--test_path",
        required=True,
        help="Path to test case JSON file"
    )
    parser.add_argument("--host", default="127.0.0.1", help="Server host")
    parser.add_argument("--port", type=int, default=8080, help="Server port")
    parser.add_argument(
        "--loop-forever",
        action="store_true",
        help="Run test cases repeatedly until interrupted"
    )
    parser.add_argument(
        "--interval",
        type=float,
        default=0.0,
        help="Seconds to sleep between test cases"
    )
    return parser.parse_args()

def start_server(server_path, model_path, case_path, port):
    """启动llama-server进程，返回subprocess.Popen对象"""
    cmd = ["sudo", server_path, "-m", str(model_path), "--case", str(case_path), "--port", str(port), "-np", str(1), "--verbose", "--no-warmup"]
    print(f"Start server: {' '.join(cmd)}")
    process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    time.sleep(1)  # 稍等片刻让进程启动
    return process

def health_check(HEALTH_URL):
    """健康检查，重试直到服务就绪或超时"""
    for i in range(1, MAX_RETRIES + 1):
        try:
            response = requests.get(HEALTH_URL, timeout=15)
            if response.status_code == 200:
                print(f"✔ llama-server is ready (attempt {i}/{MAX_RETRIES})")
                return True
        except requests.exceptions.RequestException as e:
            pass  # 忽略连接错误
        print(f"Waiting for service to be ready (attempt {i}/{MAX_RETRIES},retry in {RETRY_DELAY}s)...")
        time.sleep(RETRY_DELAY)
    print("✘ llama-server startup timed out")
    return False

def extract_stream_parts(chunk):
    choices = chunk.get("choices", [])
    if not choices:
        return "", ""

    c0 = choices[0]
    delta = c0.get("delta", {}) or {}

    think_piece = delta.get("reasoning_content", "") or ""
    answer_piece = delta.get("content", "") or ""

    # 兼容部分服务端在非 delta 里返回
    if not think_piece and not answer_piece:
        msg = c0.get("message", {}) or {}
        think_piece = msg.get("reasoning_content", "") or ""
        answer_piece = msg.get("content", "") or c0.get("text", "") or ""

    return think_piece, answer_piece

def send_request(test_case, API_URL):
    model = "deepseek"
    test_id = test_case.get("id", "unknown")
    prompt = test_case.get("prompt", "")
    if not prompt:
        return {"error": f"prompt {test_id} is null, skip"}

    request_body = {
        "model": model,
        "messages": [{"role": "user", "content": prompt}],
        "stream": True
    }

    print("Prompt:")
    print(prompt)
    print("Response(stream):")

    think_parts = []
    answer_parts = []
    think_opened = False
    answer_opened = False

    usage = {}
    timings = {}

    try:
        with requests.post(API_URL, json=request_body, stream=True, timeout=150) as response:
            response.raise_for_status()
            response.encoding = "utf-8"

            for raw_line in response.iter_lines(decode_unicode=False):
                if not raw_line:
                    continue

                line = raw_line.decode("utf-8", errors="replace")
                if not line.startswith("data:"):
                    continue

                payload = line[5:].strip()
                if not payload:
                    continue
                if payload == "[DONE]":
                    break

                try:
                    chunk = json.loads(payload)
                except json.JSONDecodeError:
                    continue

                think_piece, answer_piece = extract_stream_parts(chunk)

                if think_piece:
                    if not think_opened:
                        print("[THINK] ", end="", flush=True)
                        think_opened = True
                    print(think_piece, end="", flush=True)
                    think_parts.append(think_piece)

                if answer_piece:
                    if not answer_opened:
                        if think_opened:
                            print("\n[ANSWER] ", end="", flush=True)
                        else:
                            print("[ANSWER] ", end="", flush=True)
                        answer_opened = True
                    print(answer_piece, end="", flush=True)
                    answer_parts.append(answer_piece)

                if "usage" in chunk and isinstance(chunk["usage"], dict):
                    usage = chunk["usage"]
                if "timings" in chunk and isinstance(chunk["timings"], dict):
                    timings = chunk["timings"]

        print()

    except requests.exceptions.Timeout:
        return {"error": f"test case {test_id} timed out"}
    except requests.exceptions.RequestException as e:
        return {"error": f"Request failed: {e}"}

    think_content = "".join(think_parts)
    content = "".join(answer_parts)

    prompt_tokens = usage.get("prompt_tokens", 0)
    completion_tokens = usage.get("completion_tokens", 0)
    decode_tps = timings.get("predicted_per_second", 0.0)
    predicted_ms = timings.get("predicted_ms", 0.0)
    prompt_ms = timings.get("prompt_ms", 0.0)
    ttft = predicted_ms + prompt_ms
    prefill_tps = timings.get("prompt_per_second", 0.0)

    return {
        "prompt_tokens": prompt_tokens,
        "completion_tokens": completion_tokens,
        "memory_usage": 0,
        "decode_tps": decode_tps,
        "ttft": ttft,
        "prefill_tps": prefill_tps,
        "Power": 0,
        "prompt": prompt,
        "think_content": think_content,
        "content": content,
        "model": model
    }

def print_metrics(round_id, i, total, metrics):
    print("\n" + "=" * 80)
    print(f"[Round {round_id} | {i}/{total}]")
    print(
        f"prompt_tokens={metrics['prompt_tokens']}, "
        f"completion_tokens={metrics['completion_tokens']}, "
        f"decode_tps={metrics['decode_tps']:.2f}, "
        f"ttft={metrics['ttft']:.2f} ms, "
        f"prefill_tps={metrics['prefill_tps']:.2f}"
    )
    print("=" * 80)

def log_data(metrics):
    """记录数据到CSV文件"""
    if metrics is None:
        print("No valid data, skipping logging")
        return

    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")

    with open(LOG_FILE, "a", newline="", encoding="utf-8") as f:
        writer = csv.writer(f)
        writer.writerow([
            timestamp,
            metrics["model"],
            metrics["prompt_tokens"],
            metrics["completion_tokens"],
            f"{metrics['memory_usage']:.2f}",
            f"{metrics['decode_tps']:.2f}",
            f"{metrics['ttft']:.2f}",
            f"{metrics['prefill_tps']:.2f}",
            f"{metrics['Power']:.1f}",
        ])

    with open(CONTENT_FILE, "a", newline="", encoding="utf-8") as f:
        writer = csv.writer(f)
        writer.writerow([
            timestamp,
            metrics["model"],
            metrics["prompt"],
            metrics["content"],
        ])

def main():
    """主函数，执行完整流程"""
    print("=== Starting llama-server test ===")
    server_process = None
    try:
        args = parse_args()

        # 1. 初始化日志文件，加载json测试集
        init_log_file()

        test_cases = load_test_cases(args.test_path)
        if not test_cases:
          print("error:load test_cases failed")
          sys.exit(1)

        HEALTH_URL = f"http://{args.host}:{args.port}/health"
        API_URL = f"http://{args.host}:{args.port}/v1/chat/completions"
 
        # server_process = start_server(args.server_path, MODEL_PATH, CALBIN_PATH, PORT)
        
        # 3. 健康检查
        if not health_check(HEALTH_URL):
            raise RuntimeError("Health check failed")
        
        round_id = 0

        while True:
            round_id += 1
            successful_tests = 0
            # 每轮发送次数，默认=测试集大小；可按需改成固定值
            sample_count = len(test_cases)

            print(f"\n=== Round {round_id} start: {sample_count} samples (with replacement) ===")

            for i in range(1, sample_count + 1):
                test_case = random.choice(test_cases)  # 有放回抽样，允许重复
                print(f"\n[Round {round_id} | Progress: {i}/{sample_count}]")

                metrics = send_request(test_case, API_URL)
                if not metrics or "error" in metrics:
                    err = metrics["error"] if isinstance(metrics, dict) and "error" in metrics else "unknown error"
                    print(f"[Round {round_id} | {i}/{sample_count}] ERROR: {err}")
                else:
                    print_metrics(round_id, i, sample_count, metrics)
                    log_data(metrics)
                    successful_tests += 1

                if args.interval > 0 and i < sample_count:
                    time.sleep(args.interval)

            print(f"=== Round {round_id} done: {successful_tests}/{sample_count} success ===")

            if not args.loop_forever:
                break

    except KeyboardInterrupt:
        print("\nInterrupted by user, stopping stress test...")
    except Exception as e:
        print(f"✘ Test failed: {e}")
        sys.exit(1)
    finally:
        # 6. 停止服务器进程
        if server_process and server_process.poll() is None:  # 检查进程是否仍在运行
            server_process.terminate()
            try:
                server_process.wait(timeout=5)
                print("✔ Test completed")
            except subprocess.TimeoutExpired:
                server_process.kill()
                print("✔ Test completed")
    
    sys.exit(0)

if __name__ == "__main__":
    main()