"""
Qwen3.5 122B-A10B 자동 정밀 튜닝 스크립트
===========================================
각 설정 조합으로 서버를 재시작하고 벤치마크를 자동 수행합니다.
서버 로그에서 순수 eval time (t/s)를 파싱하여 정확한 비교 테이블을 출력합니다.

예상 소요 시간: 약 30-40분 (5개 설정 × ~6-7분/설정)
"""
import subprocess
import time
import json
import urllib.request
import os
import re
import sys
import datetime

try:
    sys.stdout.reconfigure(encoding='utf-8')
except AttributeError:
    pass

BASE_URL = "http://127.0.0.1:8000"
MODEL_PATH = r"models\Q4_K_M\Qwen3.5-122B-A10B-Q4_K_M-00001-of-00003.gguf"
SERVER_EXE = r"llama_bin_run\llama-server.exe"

# ============================================================
# 테스트할 설정 목록
# ============================================================
# 공통 파라미터 (변경하지 않는 것들)
COMMON_ARGS = [
    "--model", MODEL_PATH,
    "-ngl", "999",
    "--cpu-moe",
    "-c", "2048",
    "-np", "1",
    "-fa", "on",
    "--cache-type-k", "q4_0",
    "--cache-type-v", "q4_0",
    "-ub", "256",
    "-b", "1024",
    "--mlock",
    "--port", "8000",
    "--host", "0.0.0.0",
    "--no-warmup",  # 워밍업은 벤치마크 스크립트에서 직접 수행
]

# 변수 파라미터 조합
CONFIGS = [
    {
        "name": "A) --no-mmap -t 8",
        "desc": "서버 권장: mmap 비활성화 (baseline 대비)",
        "extra": ["--no-mmap", "-t", "8", "--prio", "2"],
    },
    {
        "name": "B) --no-mmap -t 6",
        "desc": "스레드 감소 (캐시 경합 회피)",
        "extra": ["--no-mmap", "-t", "6", "--prio", "2"],
    },
    {
        "name": "C) --no-mmap -t 10",
        "desc": "스레드 증가 (RAM 대역폭 포화)",
        "extra": ["--no-mmap", "-t", "10", "--prio", "2"],
    },
    {
        "name": "D) --no-mmap -t 12",
        "desc": "더 많은 스레드",
        "extra": ["--no-mmap", "-t", "12", "--prio", "2"],
    },
    {
        "name": "E) --no-mmap -t 10 --prio 3 --poll 100",
        "desc": "최적 스레드 + 리얼타임 우선순위 + 폴링",
        "extra": ["--no-mmap", "-t", "10", "--prio", "3", "--poll", "100"],
    },
]

# ============================================================
# 유틸리티 함수
# ============================================================

def kill_server():
    """llama-server 프로세스 강제 종료"""
    os.system("taskkill /F /IM llama-server.exe >nul 2>&1")
    time.sleep(3)

def start_server(config, log_path):
    """서버 시작, 로그를 파일로 리다이렉트"""
    cmd = [SERVER_EXE] + COMMON_ARGS + config["extra"]
    log_file = open(log_path, "w", encoding="utf-8")
    proc = subprocess.Popen(
        cmd,
        stdout=log_file,
        stderr=subprocess.STDOUT,
        cwd=os.getcwd()
    )
    return proc, log_file

def wait_for_server(timeout=600):
    """서버가 준비될 때까지 대기"""
    start = time.time()
    while time.time() - start < timeout:
        try:
            req = urllib.request.Request(f"{BASE_URL}/health")
            with urllib.request.urlopen(req, timeout=5) as resp:
                data = json.loads(resp.read())
                if data.get("status") == "ok":
                    return True
        except:
            pass
        time.sleep(5)
    return False

def run_single_benchmark(prompt, max_tokens=200):
    """단일 벤치마크 실행"""
    payload = json.dumps({
        "model": "local-model",
        "messages": [{"role": "user", "content": prompt}],
        "max_tokens": max_tokens,
        "temperature": 0.0
    }).encode("utf-8")

    req = urllib.request.Request(
        f"{BASE_URL}/v1/chat/completions",
        data=payload,
        headers={"Content-Type": "application/json"}
    )

    start = time.time()
    with urllib.request.urlopen(req, timeout=600) as resp:
        result = json.loads(resp.read())
    elapsed = time.time() - start

    usage = result.get("usage", {})
    completion_tokens = usage.get("completion_tokens", 0)
    return completion_tokens, elapsed

def parse_eval_times(log_path):
    """서버 로그에서 순수 eval time 파싱"""
    try:
        with open(log_path, "r", encoding="utf-8", errors="ignore") as f:
            content = f.read()
    except:
        return []
    
    # "eval time = XXXXX.XX ms / NNN tokens (XXX.XX ms per token, XX.XX tokens per second)"
    pattern = r'^\s+eval time\s*=\s*([\d.]+)\s*ms\s*/\s*(\d+)\s*tokens\s*\(\s*([\d.]+)\s*ms per token,\s*([\d.]+)\s*tokens per second\)'
    matches = re.findall(pattern, content, re.MULTILINE)
    
    results = []
    for m in matches:
        results.append({
            "total_ms": float(m[0]),
            "tokens": int(m[1]),
            "ms_per_token": float(m[2]),
            "tps": float(m[3])
        })
    return results

def parse_prompt_eval_times(log_path):
    """서버 로그에서 prompt eval time 파싱"""
    try:
        with open(log_path, "r", encoding="utf-8", errors="ignore") as f:
            content = f.read()
    except:
        return []
    
    pattern = r'prompt eval time\s*=\s*([\d.]+)\s*ms\s*/\s*(\d+)\s*tokens\s*\(\s*([\d.]+)\s*ms per token,\s*([\d.]+)\s*tokens per second\)'
    matches = re.findall(pattern, content, re.MULTILINE)
    
    results = []
    for m in matches:
        results.append({
            "total_ms": float(m[0]),
            "tokens": int(m[1]),
            "ms_per_token": float(m[2]),
            "tps": float(m[3])
        })
    return results

def parse_vram_usage(log_path):
    """서버 로그에서 CUDA0 모델 버퍼 크기 파싱"""
    try:
        with open(log_path, "r", encoding="utf-8", errors="ignore") as f:
            content = f.read()
    except:
        return "N/A"
    
    match = re.search(r'CUDA0 model buffer size\s*=\s*([\d.]+)\s*MiB', content)
    if match:
        return f"{float(match.group(1)):.0f} MiB"
    return "N/A"

# ============================================================
# 메인 튜닝 루프
# ============================================================

def main():
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    
    print("=" * 70)
    print("  Qwen3.5 122B-A10B 자동 정밀 튜닝")
    print(f"  시작 시간: {datetime.datetime.now().strftime('%H:%M:%S')}")
    print(f"  테스트 설정: {len(CONFIGS)}개")
    print(f"  예상 소요: ~{len(CONFIGS) * 7}분")
    print("=" * 70)
    print()
    print("  기존 Baseline: mmap on, -t 8, --prio 2 → 10.06 t/s (eval)")
    print()
    
    # 결과 저장
    all_results = []
    
    for idx, config in enumerate(CONFIGS):
        config_start = time.time()
        log_path = os.path.join(os.getcwd(), f"tune_log_{idx}.txt")
        
        print(f"\n{'='*70}")
        print(f"  [{idx+1}/{len(CONFIGS)}] {config['name']}")
        print(f"  {config['desc']}")
        print(f"  시작: {datetime.datetime.now().strftime('%H:%M:%S')}")
        print(f"{'='*70}")
        
        # 1. 기존 서버 종료
        print("  [1/4] 서버 종료 중...")
        kill_server()
        
        # 2. 새 서버 시작
        print(f"  [2/4] 서버 시작 중... (모델 로딩 ~3-5분)")
        proc, log_file = start_server(config, log_path)
        
        # 3. 서버 준비 대기
        if not wait_for_server(timeout=600):
            print("  ❌ 서버 시작 실패! 다음 설정으로 넘어갑니다.")
            kill_server()
            log_file.close()
            all_results.append({
                "config": config["name"],
                "status": "FAILED",
                "eval_tps": [],
                "prompt_tps": [],
                "vram": "N/A"
            })
            continue
        
        load_time = time.time() - config_start
        print(f"  [3/4] 서버 준비 완료! (로딩 {load_time:.0f}초)")
        
        # 4. 벤치마크 실행 (워밍업 1회 + 본 테스트 3회)
        print("  [4/4] 벤치마크 실행 중...")
        
        # 워밍업
        try:
            run_single_benchmark("Say hello.", max_tokens=20)
            print("    워밍업 완료")
        except Exception as e:
            print(f"    워밍업 실패: {e}")
        
        # 본 테스트 3회
        prompts = [
            "Write a detailed explanation of how neural networks learn through backpropagation and gradient descent.",
            "Explain the complete process of photosynthesis including light and dark reactions in detail.",
            "Describe the differences between SQL and NoSQL databases with examples and performance characteristics.",
        ]
        
        for i, prompt in enumerate(prompts):
            try:
                tokens, elapsed = run_single_benchmark(prompt, max_tokens=200)
                approx_tps = tokens / elapsed if elapsed > 0 else 0
                print(f"    Run {i+1}/3: {tokens} tokens, {elapsed:.1f}s, ~{approx_tps:.2f} t/s (approx)")
            except Exception as e:
                print(f"    Run {i+1}/3: ERROR - {e}")
        
        # 서버 종료 전에 로그 플러시를 위해 잠시 대기
        time.sleep(2)
        
        # 서버 종료
        kill_server()
        log_file.close()
        time.sleep(2)
        
        # 로그 파싱
        eval_times = parse_eval_times(log_path)
        prompt_times = parse_prompt_eval_times(log_path)
        vram = parse_vram_usage(log_path)
        
        # 워밍업 제외 (첫 번째 결과)
        if len(eval_times) > 1:
            bench_evals = eval_times[1:]  # 워밍업 제외
        else:
            bench_evals = eval_times
        
        if len(prompt_times) > 1:
            bench_prompts = prompt_times[1:]
        else:
            bench_prompts = prompt_times
        
        eval_speeds = [e["tps"] for e in bench_evals]
        prompt_speeds = [p["tps"] for p in bench_prompts]
        
        result = {
            "config": config["name"],
            "status": "OK",
            "eval_tps": eval_speeds,
            "prompt_tps": prompt_speeds,
            "vram": vram,
        }
        all_results.append(result)
        
        config_elapsed = time.time() - config_start
        print(f"\n  완료! 소요: {config_elapsed:.0f}초")
        
        if eval_speeds:
            avg_eval = sum(eval_speeds) / len(eval_speeds)
            max_eval = max(eval_speeds)
            print(f"  📊 Eval TPS: avg={avg_eval:.2f}, max={max_eval:.2f}")
    
    # ============================================================
    # 최종 결과 비교 테이블
    # ============================================================
    print("\n")
    print("=" * 80)
    print("  🏆 최종 결과 비교 테이블")
    print("=" * 80)
    print()
    
    # 기존 baseline 추가
    print(f"  {'설정':<45} {'Eval t/s':>10} {'최대':>8} {'Prompt t/s':>12} {'VRAM':>12}")
    print(f"  {'-'*45} {'-'*10} {'-'*8} {'-'*12} {'-'*12}")
    
    # Baseline (이전 결과)
    print(f"  {'[기준] mmap on, -t 8, --prio 2':<45} {'10.02':>10} {'10.06':>8} {'29.52':>12} {'5392 MiB':>12}")
    
    best_avg = 0
    best_config = ""
    
    for r in all_results:
        if r["status"] != "OK" or not r["eval_tps"]:
            print(f"  {r['config']:<45} {'FAILED':>10} {'':>8} {'':>12} {r['vram']:>12}")
            continue
        
        avg_e = sum(r["eval_tps"]) / len(r["eval_tps"])
        max_e = max(r["eval_tps"])
        avg_p = sum(r["prompt_tps"]) / len(r["prompt_tps"]) if r["prompt_tps"] else 0
        
        if avg_e > best_avg:
            best_avg = avg_e
            best_config = r["config"]
        
        marker = " ⭐" if avg_e > 10.06 else ""
        print(f"  {r['config']:<45} {avg_e:>10.2f} {max_e:>8.2f} {avg_p:>12.2f} {r['vram']:>12}{marker}")
    
    print()
    if best_avg > 0:
        improvement = ((best_avg - 10.02) / 10.02) * 100
        print(f"  🏆 최고 성능: {best_config}")
        print(f"     → {best_avg:.2f} t/s (기준 10.02 t/s 대비 {improvement:+.1f}%)")
    
    print()
    print(f"  완료 시간: {datetime.datetime.now().strftime('%H:%M:%S')}")
    print("=" * 80)
    
    # 결과를 파일로도 저장
    result_path = os.path.join(os.getcwd(), f"tune_results_{timestamp}.txt")
    with open(result_path, "w", encoding="utf-8") as f:
        f.write("Qwen3.5 122B-A10B Fine Tuning Results\n")
        f.write(f"Date: {timestamp}\n\n")
        for r in all_results:
            f.write(f"{r['config']}: {r['eval_tps']} (VRAM: {r['vram']})\n")
    print(f"  결과 저장: {result_path}")

if __name__ == "__main__":
    main()
