#!/usr/bin/env python3
# DGX AI Factory - 14B canary validation training
#
# Goal:
# - Start a small 14B validation LoRA from base 14B model.
# - Do NOT reuse 8B LoRA adapters because 8B adapters are incompatible with 14B.
# - Train only a small canary run.
# - Save under experiments/14b_canary_validation_*.
# - Run smoke test immediately after training.
# - Never overwrite 8B outputs or stable adapters.

import os
import re
import json
import time
import random
import traceback
from pathlib import Path
from typing import Any, Dict, Optional, List

ROOT = Path("/home/harness_user_1/dgx_ai_factory")
EXPERIMENT_ROOT = ROOT / "experiments"
REPORT_DIR = ROOT / "reports"
LOG_DIR = ROOT / "logs"

for d in [EXPERIMENT_ROOT, REPORT_DIR, LOG_DIR]:
    d.mkdir(parents=True, exist_ok=True)

BASE_MODEL = os.environ.get("BASE_MODEL_14B", "Qwen/Qwen2.5-14B-Instruct")
MAX_STEPS = int(os.environ.get("MAX_STEPS", "80"))
LR = float(os.environ.get("LR", "5e-6"))
MAX_LENGTH = int(os.environ.get("MAX_LENGTH", "1024"))
SEED = int(os.environ.get("SEED", "42"))

DATASET_CANDIDATES = [
    Path(os.environ.get("TRAIN_JSONL", "")) if os.environ.get("TRAIN_JSONL") else None,
    ROOT / "data/train_100k_gold/train_canary_1000_gold_clean.jsonl",
    ROOT / "data/train_100k/train_canary_1000.jsonl",
    ROOT / "dataset_quality_2000_candidate_clean.jsonl",
    ROOT / "dataset.jsonl",
]

SYSTEM_PROMPT = (
    "당신은 한국어 존댓말을 사용하는 실무형 개발 비서입니다. "
    "중국어를 절대 쓰지 말고, 질문과 무관한 내용을 반복하지 말고, 명령어와 설명을 분리해서 정확하게 답변합니다."
)

SMOKE_TESTS = [
    {
        "name": "vllm_stop",
        "prompt": "vLLM 서버가 켜져 있는지 확인하고 꺼야 할 때 쓰는 명령어를 알려주세요.",
        "must": ["ps", "grep", "vllm"],
        "any_groups": [["pkill", "kill"], ["8000", "api_server", "vllm.entrypoints"]],
    },
    {
        "name": "systemd_217_user",
        "prompt": "systemd 서비스에서 status=217/USER 오류가 났습니다. 원인과 확인 명령어를 알려주세요.",
        "must": ["217/USER", "User", "systemctl", "journalctl"],
        "any_groups": [["getent passwd", "id "], ["daemon-reload", "restart"]],
    },
    {
        "name": "fastapi_curl",
        "prompt": "FastAPI 헬스체크 엔드포인트를 만들고 curl로 확인하는 예시를 보여주세요.",
        "must": ["FastAPI", "curl", "/health"],
        "any_groups": [],
    },
    {
        "name": "jsonl_validate",
        "prompt": "JSONL 파일이 깨졌는지 파이썬으로 검사하는 짧은 코드를 보여주세요.",
        "must": ["json", "enumerate", "json.loads"],
        "any_groups": [["JSONDecodeError", "Exception"], ["line_no", "line", "줄"]],
    },
    {
        "name": "lora_next_step",
        "prompt": "LoRA 파인튜닝이 끝난 뒤 smoke test와 benchmark를 어떤 순서로 해야 하는지 정리해주세요.",
        "must": ["smoke", "benchmark"],
        "any_groups": [["stable", "승격"], ["rejected", "reject", "폐기", "수리", "repair"]],
    },
]

def has_chinese_ideograph(text: str) -> bool:
    return bool(re.search(r"[\u4e00-\u9fff]", text))

def template_leak(text: str) -> bool:
    bad = [
        "{problem}", "{framework}", "{concept}", "{bug}", "{system}",
        "### Instruction", "### Response", "<|im_start|>", "<|im_end|>",
        "assistant:", "user:"
    ]
    low = text.lower()
    return any(x.lower() in low for x in bad)

def honorific_ok(text: str) -> bool:
    return any(m in text for m in ["습니다", "합니다", "됩니다", "주세요", "하세요", "확인하시면", "좋습니다", "하시면"])

def pick_dataset() -> Path:
    for p in DATASET_CANDIDATES:
        if p and p.exists() and p.stat().st_size > 0:
            return p
    node-7.example.invalid SystemExit("[FAIL] No usable canary dataset found. Expected data/train_100k_gold/train_canary_1000_gold_clean.jsonl")

def normalize_row(obj: Any) -> Optional[Dict[str, str]]:
    if not isinstance(obj, dict):
        return None

    msgs = obj.get("messages")
    if isinstance(msgs, list):
        user_parts, assistant_parts = [], []
        for m in msgs:
            if not isinstance(m, dict):
                continue
            role = str(m.get("role", "")).lower()
            content = str(m.get("content", "") or "").strip()
            if not content:
                continue
            if role == "user":
                user_parts.append(content)
            elif role == "assistant":
                assistant_parts.append(content)
        if user_parts and assistant_parts:
            return {
                "category": str(obj.get("category", "canary")),
                "instruction": "\n\n".join(user_parts).strip(),
                "input": "",
                "output": "\n\n".join(assistant_parts).strip(),
            }

    instruction = obj.get("instruction") or obj.get("prompt") or obj.get("question") or obj.get("input_text")
    output = obj.get("output") or obj.get("response") or obj.get("answer") or obj.get("completion")

    if instruction and output:
        return {
            "category": str(obj.get("category", "canary")),
            "instruction": str(instruction).strip(),
            "input": str(obj.get("input", "") or "").strip(),
            "output": str(output).strip(),
        }

    return None

def load_rows(path: Path, limit: int = 1000) -> List[Dict[str, str]]:
    rows = []
    bad_json = 0
    skipped = 0

    with open(path, "r", encoding="utf-8") as f:
        for line_no, line in enumerate(f, 1):
            line = line.strip()
            if not line:
                continue
            try:
                obj = json.loads(line)
            except Exception:
                bad_json += 1
                continue

            row = normalize_row(obj)
            if not row:
                skipped += 1
                continue

            if len(row["instruction"]) < 4 or len(row["output"]) < 20:
                skipped += 1
                continue

            if has_chinese_ideograph(row["output"]):
                skipped += 1
                continue

            rows.append(row)
            if len(rows) >= limit:
                break

    print("==== DATASET LOAD ====")
    print("dataset:", path)
    print("loaded_rows:", len(rows))
    print("bad_json:", bad_json)
    print("skipped:", skipped)

    if len(rows) < 100:
        node-7.example.invalid SystemExit("[FAIL] Too few usable canary rows. Need at least 100.")

    return rows

def format_chat(tokenizer, row: Dict[str, str]) -> str:
    user = row["instruction"]
    if row.get("input"):
        user += "\n\n입력:\n" + row["input"]

    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": user},
        {"role": "assistant", "content": row["output"]},
    ]

    return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False)

def build_dataset(tokenizer, rows: List[Dict[str, str]], max_length: int):
    from datasets import Dataset

    ds = Dataset.from_list(rows)

    def tok(row):
        text = format_chat(tokenizer, row)
        enc = tokenizer(
            text,
            truncation=True,
            max_length=max_length,
            padding=False,
        )
        enc["labels"] = enc["input_ids"].copy()
        return enc

    return ds.map(tok, remove_columns=ds.column_names)

def generate_one(tokenizer, model, prompt: str, max_new_tokens: int = 320) -> str:
    import torch

    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": prompt},
    ]

    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer(text, return_tensors="pt").to(model.device)

    with torch.no_grad():
        out = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=False,
            repetition_penalty=1.08,
            pad_token_id=tokenizer.eos_token_id,
        )

    gen = out[0][inputs["input_ids"].shape[-1]:]
    return tokenizer.decode(gen, skip_special_tokens=True).strip()

def score_smoke_output(output: str, test: Dict[str, Any]) -> Dict[str, Any]:
    score = 100
    flags = []
    low = output.lower()

    if has_chinese_ideograph(output):
        score -= 60
        flags.append("chinese_cjk_leakage")

    if template_leak(output):
        score -= 35
        flags.append("template_or_role_leakage")

    if not honorific_ok(output):
        score -= 15
        flags.append("weak_korean_honorific")

    if len(output.strip()) < 60:
        score -= 15
        flags.append("too_short")

    for item in test["must"]:
        if item.lower() not in low:
            score -= 18
            flags.append("missing:" + item)

    for group in test.get("any_groups", []):
        if not any(x.lower() in low for x in group):
            score -= 10
            flags.append("missing_any:" + "|".join(group))

    score = max(0, score)
    return {
        "score": score,
        "status": "PASS" if score >= 70 else "FAIL",
        "flags": flags,
    }

def run_smoke(tokenizer, model, adapter_dir: Path, exp_dir: Path, dataset_path: Path):
    print("\n" + "=" * 80)
    print("14B CANARY SMOKE TEST")
    print("=" * 80)

    results = []
    pass_count = 0

    model.eval()

    for test in SMOKE_TESTS:
        print("\n" + "-" * 80)
        print("TEST:", test["name"])
        print("-" * 80)

        output = generate_one(tokenizer, model, test["prompt"])
        scored = score_smoke_output(output, test)
        if scored["status"] == "PASS":
            pass_count += 1

        print(output)
        print("\nSCORE:", scored["score"])
        print("STATUS:", scored["status"])
        if scored["flags"]:
            print("FLAGS:", scored["flags"])

        results.append({
            "name": test["name"],
            "prompt": test["prompt"],
            "output": output,
            **scored,
        })

    avg = sum(r["score"] for r in results) / len(results)
    decision = "PASS_14B_CANARY_CANDIDATE" if pass_count == len(results) else ("REVIEW_14B_CANARY" if pass_count >= 4 else "REJECT_OR_REPAIR_NEEDED")

    report = {
        "timestamp": time.strftime("%Y%m%d_%H%M%S"),
        "base_model": BASE_MODEL,
        "adapter_dir": str(adapter_dir),
        "experiment_dir": str(exp_dir),
        "dataset_path": str(dataset_path),
        "average_score": round(avg, 2),
        "smoke_pass": pass_count,
        "smoke_total": len(results),
        "decision": decision,
        "results": results,
    }

    report_path = exp_dir / "canary_smoke_report.json"
    report_path.write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")

    global_report_path = REPORT_DIR / f"14b_canary_smoke_{report['timestamp']}.json"
    global_report_path.write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")

    print("\n" + "=" * 80)
    print("FINAL 14B CANARY SUMMARY")
    print("=" * 80)
    print("base_model:", BASE_MODEL)
    print("adapter_dir:", adapter_dir)
    print("average_score:", round(avg, 2))
    print("smoke_pass:", f"{pass_count}/{len(results)}")
    print("decision:", decision)
    print("report_path:", report_path)
    print("global_report_path:", global_report_path)

def main():
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer, DataCollatorForLanguageModeling
    from peft import LoraConfig, get_peft_model, TaskType

    random.seed(SEED)
    torch.manual_seed(SEED)

    dataset_path = pick_dataset()
    ts = time.strftime("%Y%m%d_%H%M%S")
    exp_dir = EXPERIMENT_ROOT / f"14b_canary_validation_{ts}"
    adapter_dir = exp_dir / "adapter"
    exp_dir.mkdir(parents=True, exist_ok=True)

    config = {
        "base_model": BASE_MODEL,
        "dataset_path": str(dataset_path),
        "max_steps": MAX_STEPS,
        "learning_rate": LR,
        "max_length": MAX_LENGTH,
        "seed": SEED,
        "lora_r": 16,
        "lora_alpha": 32,
        "lora_dropout": 0.05,
        "target_modules": ["q_proj", "v_proj"],
        "note": "14B canary validation. 8B adapters are not reused.",
    }
    (exp_dir / "config.json").write_text(json.dumps(config, ensure_ascii=False, indent=2), encoding="utf-8")

    print("==== 14B CANARY VALIDATION START ====")
    print("base_model:", BASE_MODEL)
    print("dataset_path:", dataset_path)
    print("experiment_dir:", exp_dir)
    print("adapter_dir:", adapter_dir)
    print("max_steps:", MAX_STEPS)
    print("lr:", LR)
    print("max_length:", MAX_LENGTH)

    rows = load_rows(dataset_path, limit=1000)

    print("\n==== LOAD TOKENIZER ====")
    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = <REDACTED_SECRET>

    tokenized = build_dataset(tokenizer, rows, MAX_LENGTH)

    print("\n==== LOAD 14B BASE MODEL ====")
    dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32

    model = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL,
        torch_dtype=dtype,
        device_map="auto",
        trust_remote_code=True,
    )

    model.config.use_cache = False

    if hasattr(model, "gradient_checkpointing_enable"):
        model.gradient_checkpointing_enable()

    print("\n==== APPLY LORA ====")
    lora_cfg = LoraConfig(
        task_type=TaskType.CAUSAL_LM,
        r=16,
        lora_alpha=32,
        lora_dropout=0.05,
        target_modules=["q_proj", "v_proj"],
        bias="none",
    )
    model = get_peft_model(model, lora_cfg)
    model.print_trainable_parameters()

    collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)

    args = TrainingArguments(
        output_dir=str(adapter_dir),
        per_device_train_batch_size=1,
        gradient_accumulation_steps=8,
        learning_rate=LR,
        max_steps=MAX_STEPS,
        logging_steps=5,
        save_steps=MAX_STEPS,
        save_total_limit=1,
        bf16=torch.cuda.is_available(),
        fp16=False,
        gradient_checkpointing=True,
        report_to=[],
        remove_unused_columns=False,
        optim="adamw_torch",
    )

    trainer = Trainer(
        model=model,
        args=args,
        train_dataset=tokenized,
        data_collator=collator,
    )

    print("\n==== TRAIN 14B CANARY ====")
    train_result = trainer.train()

    metrics = train_result.metrics or {}
    (exp_dir / "train_metrics.json").write_text(json.dumps(metrics, ensure_ascii=False, indent=2), encoding="utf-8")

    print("\n==== SAVE 14B CANARY ADAPTER ====")
    model.save_pretrained(str(adapter_dir))
    tokenizer.save_pretrained(str(adapter_dir))

    run_smoke(tokenizer, model, adapter_dir, exp_dir, dataset_path)

if __name__ == "__main__":
    try:
        main()
    except Exception:
        print("\n[FATAL ERROR]")
        traceback.print_exc()
        node-7.example.invalid