#!/usr/bin/env python3
import subprocess
import time
import requests
import json
import os
import sys
import argparse
from datetime import datetime
from typing import List, Tuple, Optional

# 配置常量
LOG_FILE = "llama_test_logs.csv"
CONTENT_FILE = "llama_content_logs.csv"
MAX_RETRIES = 10
RETRY_DELAY = 15
SERVER_PATH = "./build/bin/llama-server"

# 固定的测试prompt
FIXED_PROMPT = "who are you?"

def init_log_file():
    """初始化日志文件，写入CSV表头（如果文件不存在）"""
    if not os.path.exists(LOG_FILE):
        with open(LOG_FILE, "w") as f:
            f.write("timestamp,model_path,case_path,case_name,prompt_tokens,completion_tokens,memory_usage,Decode_TPS,TTFT,Prefill_TPS,Power,result\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_path,case_path,case_name,prompt,content\n")
        print(f"Initialized content file: {CONTENT_FILE}")

def parse_args():
    parser = argparse.ArgumentParser(
        description='Automatically test model (.gguf) files in case directories with fixed prompt.'
    )
    parser.add_argument(
        "-c", "--case_path",
        required=True,
        help="Root directory to scan for case folders containing .gguf model files."
    )
    parser.add_argument(
        "--port",
        type=int,
        default=8080,
        help="llama-server port"
    )
    return parser.parse_args()

def discover_gguf_models(root_dir: str) -> List[Tuple[str, str, str]]:
    """
    遍历根目录，查找所有包含.gguf文件的文件夹。
    返回一个列表，元素为 (model_path, case_path, case_name) 元组。
    case_path 是包含.gguf文件的文件夹路径，case_name 是文件夹名称。
    """
    test_cases = []
    
    if not os.path.isdir(root_dir):
        print(f"Error: Provided root_dir '{root_dir}' is not a valid directory.")
        return test_cases
    
    print(f"Scanning root directory: {root_dir}")
    
    for item in os.scandir(root_dir):
        if item.is_dir():
            case_path = item.path
            case_name = item.name
            gguf_files = []
            
            # 遍历文件夹内所有文件，查找.gguf文件
            for file_item in os.scandir(case_path):
                if file_item.is_file() and file_item.name.endswith('.gguf'):
                    gguf_files.append(file_item.path)
            
            if gguf_files:
                # 如果有多个.gguf文件，可以选择第一个或让用户选择
                # 这里我们选择第一个找到的.gguf文件
                model_path = gguf_files[0]
                if len(gguf_files) > 1:
                    print(f"  Warning: Found {len(gguf_files)} .gguf files in '{case_name}', using: {os.path.basename(model_path)}")
                test_cases.append((model_path, case_path, case_name))
                print(f"  Found: Case='{case_name}', Model='{os.path.basename(model_path)}'")
            else:
                print(f"  Skipping folder '{case_name}': No .gguf files found.")
    
    print(f"Total discovered test cases: {len(test_cases)}\n")
    return test_cases

def start_server(model_path: str, case_path: str, port: int) -> Optional[subprocess.Popen]:
    """启动llama-server进程，返回subprocess.Popen对象。如果启动失败，返回None。"""
    cmd = ["sudo", SERVER_PATH, "-m", str(model_path), "--case", str(case_path), "--port", str(port), "-np", str(1), "--verbose", "--no-warmup"]
    print(f"Starting server with: Model='{model_path}', Case='{case_path}'")
    print(f"Command: {' '.join(cmd)}")
    try:
        process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        time.sleep(1)  # 稍等片刻让进程启动
        return process
    except Exception as e:
        print(f"Failed to start server: {e}")
        return None

def health_check(health_url: str) -> bool:
    """健康检查，重试直到服务就绪或超时"""
    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 send_request(prompt: str, model_path: str, api_url: str) -> Optional[dict]:
    """发送对话请求，提取并返回性能数据。如果失败，返回None。"""
    request_body = {
        "model": model_path,
        "messages": [{"role": "user", "content": prompt}],
        "stream": False
    }
    print(f"  Sending request with prompt: '{prompt[:50]}...'" if len(prompt) > 50 else f"  Sending request: '{prompt}'")
    try:
        response = requests.post(api_url, json=request_body, timeout=150)
        response.raise_for_status()
        data = response.json()
    except requests.exceptions.Timeout:
        print(f"  ✗ Request timed out.")
        return None
    except requests.exceptions.RequestException as e:
        print(f"  Request failed: {e}")
        return None
    except json.JSONDecodeError as e:
        print(f"  JSON decode failed: {e}")
        return None

    # 提取数据
    usage = data.get("usage", {})
    timings = data.get("timings", {})
    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)
    message = data.get("choices", [{}])[0].get("message", {})
    content = message.get("reasoning_content") or message.get("content", "")

    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,
        "content": content,
        "model_path": model_path
    }

def log_data(metrics: dict, case_path: str, case_name: str, result: str):
    """记录数据到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") as f:
        row = f"{timestamp},{metrics['model_path']},{case_path},{case_name},{metrics['prompt_tokens']},{metrics['completion_tokens']},{metrics['memory_usage']:.2f},{metrics['decode_tps']:.2f},{metrics['ttft']:.2f},{metrics['prefill_tps']:.2f},{metrics['Power']:.1f},{result}\n"
        f.write(row)
    with open(CONTENT_FILE, "a") as f:
        # 对内容中的换行符和逗号进行处理，用空格替换换行，用占位符替换逗号以确保CSV格式正确
        DELIMITER = '\t'
        # safe_content = metrics['content'].replace('\n', ' ').replace(',', ';')
        # safe_prompt = metrics['prompt'].replace('\n', ' ').replace(',', ';')
        row = f"{timestamp},{metrics['model_path']}{DELIMITER}{case_path}{DELIMITER}{case_name}{DELIMITER}{metrics['prompt']}{DELIMITER}{metrics['content']}\n"
        f.write(row)

def main():
    """主函数，执行完整流程"""
    print("=== Starting Automated llama-server Test Suite (Fixed Prompt) ===")
    
    # 1. 初始化日志文件
    init_log_file()

    # 2. 解析参数
    args = parse_args()
    
    # 设置固定prompt
    #fixed_prompt = args.prompt
    print(f"Using fixed prompt: '{FIXED_PROMPT}'")
    
    # 发现所有包含.gguf文件的文件夹
    test_cases = discover_gguf_models(args.case_path)
    
    if not test_cases:
        print("No folders with .gguf files found. Exiting.")
        sys.exit(1)
    
    overall_results = []

    # 3. 遍历每个测试组合
    for model_path, case_path, case_name in test_cases:
        server_process = None
        combination_result = "FAIL"  # 默认失败
        print(f"\n{'='*60}")
        print(f"Testing Case:")
        print(f"  Case Name: {case_name}")
        print(f"  Case Path: {case_path}")
        print(f"  Model:     {model_path}")
        print(f"{'='*60}")
        
        try:
            # 启动服务器
            server_process = start_server(model_path, case_path, args.port)
            if not server_process:
                raise RuntimeError(f"Failed to start server for model '{model_path}' and case '{case_path}'")
            
            # 健康检查
            health_url = f"http://localhost:{args.port}/health"
            api_url = f"http://localhost:{args.port}/v1/chat/completions"
            
            if not health_check(health_url):
                raise RuntimeError("Health check failed. Server may not have started correctly.")
            
            # 使用固定prompt发送请求
            print(f"\n  Sending test request with fixed prompt...")
            
            metrics = send_request(FIXED_PROMPT, model_path, api_url)
            
            if metrics:
                # 记录请求的详细日志
                log_data(metrics, case_path, case_name, "REQUEST_OK")
                combination_result = "PASS"
                print(f"\n  ✔ Case '{case_name}' TEST PASSED: Request succeeded.")
            else:
                combination_result = "FAIL"
                print(f"\n  ✗ Case '{case_name}' TEST FAILED: Request failed.")
            
            # 记录该组合的最终结果（汇总行）
            timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            # with open(LOG_FILE, "a") as f:
            #     summary_row = f"{timestamp},{model_path},{case_path},{case_name},,,,,,,,{combination_result}\n"
            #     f.write(summary_row)
            
            overall_results.append((model_path, case_path, case_name, combination_result))
            
        except Exception as e:
            print(f"\n  ✗ An error occurred during testing of case '{case_name}': {e}")
            combination_result = "ERROR"
            overall_results.append((model_path, case_path, case_name, f"ERROR: {e}"))
        finally:
            # 停止当前组合的服务器进程
            if server_process and server_process.poll() is None:
                print(f"  Stopping server for case '{case_name}'...")
                server_process.terminate()
                try:
                    server_process.wait(timeout=5)
                except subprocess.TimeoutExpired:
                    server_process.kill()
                # 短暂等待，避免端口占用影响下一个测试
                time.sleep(2)
    
    # 4. 打印最终汇总报告
    print(f"\n{'='*60}")
    print("TEST SUITE SUMMARY:")
    print(f"{'='*60}")
    for model_path, case_path, case_name, result in overall_results:
        status_symbol = "✔" if result == "PASS" else "✗"
        print(f"{status_symbol} Case: {case_name:50} | Result: {result}")
    
    # 判断整个测试套件是否全部通过
    all_passed = all(result == "PASS" for _, _, _, result in overall_results)
    if all_passed:
        print("\n✔ All test cases PASSED.")
        sys.exit(0)
    else:
        print("\n✗ Some test cases FAILED or encountered ERRORS.")
        sys.exit(1)

if __name__ == "__main__":
    main()