"""SFT warmstart: roll a rule-based "good triager" through the OpsGuard env,
record (prompt, completion) pairs, then SFT a LoRA adapter for one epoch.

The resulting LoRA can be loaded as the starting point for GRPO — this avoids
GRPO collapse on cold-start (model emits gibberish for many steps before any
reward signal). The expert policy uses `truth_action` from the env-side
IssueRow when accessible (via direct env import), otherwise falls back to a
heuristic on title/body content.

Usage:
    # 1. Generate dataset (rolls episodes through scenarios E0/E1)
    python scripts/sft_warmstart.py --gen-only --episodes 64 \\
        --out-jsonl data/sft_warmstart.jsonl

    # 2. SFT 1 epoch on the recorded pairs
    python scripts/sft_warmstart.py --train-only \\
        --jsonl data/sft_warmstart.jsonl \\
        --model unsloth/Qwen2.5-7B-Instruct-bnb-4bit \\
        --output-dir checkpoints/opsguard-sft

    # 3. Or do both back-to-back
    python scripts/sft_warmstart.py --episodes 64 --model <m> --output-dir <d>
"""
from __future__ import annotations

import argparse
import json
import os
import sys
from pathlib import Path
from typing import Any

# Make project root importable when run as `python scripts/sft_warmstart.py`
_PROJECT_ROOT = Path(__file__).resolve().parent.parent
if str(_PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(_PROJECT_ROOT))

from models import ActionType, OpsguardAction, OpsguardObservation  # noqa: E402
try:
    from scripts.system_prompt import SYSTEM_PROMPT, format_observation  # noqa: E402
except ImportError:
    # Fallback when scripts/ has no __init__.py and namespace-package
    # discovery doesn't pick it up (rare, but possible on certain layouts).
    sys.path.insert(0, str(_PROJECT_ROOT / "scripts"))
    from system_prompt import SYSTEM_PROMPT, format_observation  # type: ignore  # noqa: E402


# ---------------------------------------------------------------------------
# Expert policy
# ---------------------------------------------------------------------------

# Heuristic mapping from common label tokens to the canonical action.
_SPAM_KEYWORDS = (
    "urgent", "p0", "production down", "critical", "blocking",
    "stop everything", "+1", "everyone", "bumping", "still waiting",
    "doesn't work", "broken pls", "unusable",
)
_INFO_KEYWORDS = ("not sure", "maybe", "doesn't reproduce", "no repro",
                  "could you", "can you provide")


def _heuristic_action(obs: OpsguardObservation) -> OpsguardAction:
    """Fallback when no truth_action is available."""
    if obs.current_issue is None:
        return OpsguardAction(action_type=ActionType.WAIT, reasoning="no current issue")
    ci = obs.current_issue
    text = f"{ci.title}\n{ci.body}".lower()

    # Spam heuristics: brand new author + urgency phrasing
    is_likely_spam = (
        ci.author_pr_count == 0
        and ci.author_account_age_days < 90
        and any(k in text for k in _SPAM_KEYWORDS)
    )
    if is_likely_spam:
        return OpsguardAction(
            action_type=ActionType.CLOSE_SPAM,
            target_issue_id=ci.issue_id,
            reasoning="urgency-language + new account, no prior contributions",
        )

    if ci.is_pr:
        return OpsguardAction(
            action_type=ActionType.MERGE_PR,
            target_issue_id=ci.issue_id,
            reasoning="PR — defer to maintainer review queue",
        )

    if any(k in text for k in _INFO_KEYWORDS) or len(ci.body or "") < 60:
        return OpsguardAction(
            action_type=ActionType.REQUEST_INFO,
            target_issue_id=ci.issue_id,
            comment_body="Could you share a minimal reproduction (Python version, library version, full traceback)?",
            reasoning="too little context to act",
        )

    # Default: label using the most plausible label from available pool
    label = "bug"
    if ci.available_labels:
        # Pick a label whose name appears in the issue text, else 'bug'
        for cand in ci.available_labels:
            if cand.lower() in text:
                label = cand
                break
        else:
            # try common fallbacks
            for cand in ("bug", "enhancement", "documentation", "question"):
                if cand in [l.lower() for l in ci.available_labels]:
                    label = cand
                    break

    return OpsguardAction(
        action_type=ActionType.LABEL,
        target_issue_id=ci.issue_id,
        label=label,
        reasoning=f"applying topical label '{label}'",
    )


def _truth_action(obs: OpsguardObservation, env: Any) -> OpsguardAction | None:
    """Try to read the ground-truth action from the env's internal queue.

    This works only when we have a direct in-process env handle (not the
    HTTP client). When unavailable returns None.
    """
    ep = getattr(env, "_episode", None)
    if ep is None or obs.current_issue is None:
        return None
    pos = getattr(ep, "pos", None)
    queue = getattr(ep, "queue", None)
    if pos is None or queue is None or pos >= len(queue):
        return None
    cur = queue[pos]
    if cur.issue_id != obs.current_issue.issue_id:
        return None

    truth = cur.truth_action
    if truth == "label":
        label = (cur.truth_labels[0] if cur.truth_labels else "bug")
        return OpsguardAction(
            action_type=ActionType.LABEL,
            target_issue_id=cur.issue_id,
            label=label,
            reasoning=f"truth label={label}",
        )
    if truth == "close_spam":
        return OpsguardAction(
            action_type=ActionType.CLOSE_SPAM,
            target_issue_id=cur.issue_id,
            reasoning="truth=close_spam",
        )
    if truth == "request_info":
        return OpsguardAction(
            action_type=ActionType.REQUEST_INFO,
            target_issue_id=cur.issue_id,
            comment_body="Please share a minimal reproduction.",
            reasoning="truth=request_info",
        )
    if truth == "link_duplicate":
        return OpsguardAction(
            action_type=ActionType.LINK_DUPLICATE,
            target_issue_id=cur.issue_id,
            duplicate_of_id=None,  # unknown from row-level metadata
            reasoning="truth=link_duplicate",
        )
    if truth == "assign":
        return OpsguardAction(
            action_type=ActionType.ASSIGN,
            target_issue_id=cur.issue_id,
            assignee_login=cur.truth_assignee or "maintainer",
            reasoning="truth=assign",
        )
    if truth == "comment":
        return OpsguardAction(
            action_type=ActionType.COMMENT,
            target_issue_id=cur.issue_id,
            comment_body="Thanks for the report — taking a look.",
            reasoning="truth=comment",
        )
    if truth == "merge_pr":
        return OpsguardAction(
            action_type=ActionType.MERGE_PR,
            target_issue_id=cur.issue_id,
            reasoning="truth=merge_pr",
        )
    return None


def expert_action(obs: OpsguardObservation, env: Any) -> OpsguardAction:
    truth = _truth_action(obs, env)
    if truth is not None:
        return truth
    return _heuristic_action(obs)


# ---------------------------------------------------------------------------
# Dataset generation by rolling expert through the env
# ---------------------------------------------------------------------------

def _make_inproc_env():
    """Construct the env directly (no HTTP server) for fast SFT data gen."""
    # Add server/ to path so its relative imports resolve
    server_dir = _PROJECT_ROOT / "server"
    if str(server_dir) not in sys.path:
        sys.path.insert(0, str(server_dir))
    from server.opsguard_environment import OpsguardEnvironment  # type: ignore
    return OpsguardEnvironment()


def generate_dataset(
    *,
    n_episodes: int,
    scenarios: list[str],
    max_steps: int,
    out_path: Path,
) -> int:
    """Roll the expert and write JSONL of {prompt, completion} pairs."""
    env = _make_inproc_env()
    out_path.parent.mkdir(parents=True, exist_ok=True)
    n_pairs = 0
    with out_path.open("w", encoding="utf-8") as fh:
        for ep_idx in range(n_episodes):
            scen = scenarios[ep_idx % len(scenarios)]
            obs = env.reset(scenario_id=scen, seed=ep_idx)
            for step_i in range(max_steps):
                if obs.done or obs.current_issue is None:
                    break
                action = expert_action(obs, env)

                # Build the (prompt, completion) pair
                user_msg = format_observation(obs)
                completion_obj: dict[str, Any] = {
                    "action_type": action.action_type.value,
                    "target_issue_id": action.target_issue_id,
                }
                for k in ("label", "duplicate_of_id", "assignee_login",
                          "comment_body", "query", "reasoning"):
                    v = getattr(action, k, None)
                    if v is not None:
                        completion_obj[k] = v
                completion = json.dumps(completion_obj, ensure_ascii=False)

                fh.write(json.dumps({
                    "messages": [
                        {"role": "system", "content": SYSTEM_PROMPT},
                        {"role": "user", "content": user_msg},
                        {"role": "assistant", "content": completion},
                    ],
                    "scenario": scen,
                    "step": step_i,
                }, ensure_ascii=False) + "\n")
                n_pairs += 1

                obs = env.step(action)
            print(f"  episode {ep_idx + 1}/{n_episodes} ({scen}): wrote {n_pairs} pairs total", flush=True)

    print(f"DONE: wrote {n_pairs} (prompt, completion) pairs to {out_path}", flush=True)
    return n_pairs


# ---------------------------------------------------------------------------
# SFT training
# ---------------------------------------------------------------------------

def _safe_kwargs(cls, kwargs: dict) -> dict:
    """Drop kwargs not accepted by `cls.__init__` to survive TRL API drift."""
    import inspect
    try:
        sig = inspect.signature(cls.__init__)
        valid = set(sig.parameters.keys())
    except (TypeError, ValueError):
        return kwargs
    return {k: v for k, v in kwargs.items() if k in valid or "kwargs" in valid}


def train_sft(
    *,
    model_name: str,
    jsonl_path: Path,
    output_dir: Path,
    epochs: int = 1,
    batch_size: int = 2,
    grad_accum: int = 8,
    lr: float = 2e-4,
    max_seq_length: int = 4096,
    push_to_hub: str | None = None,
):
    """Run SFT with Unsloth (preferred) or plain transformers + peft + bnb (fallback)."""
    from datasets import load_dataset  # type: ignore
    from trl import SFTConfig, SFTTrainer  # type: ignore

    try:
        from unsloth import FastLanguageModel  # type: ignore
        print(f"loading model {model_name} via Unsloth (4-bit)...", flush=True)
        model, tokenizer = FastLanguageModel.from_pretrained(
            model_name=model_name, max_seq_length=max_seq_length,
            load_in_4bit=True, dtype=None,
        )
        model = FastLanguageModel.get_peft_model(
            model, r=32, lora_alpha=64, lora_dropout=0.0, bias="none",
            target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                            "gate_proj", "up_proj", "down_proj"],
            use_gradient_checkpointing="unsloth",
        )
    except ImportError:
        import torch
        from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
        from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
        print(f"loading model {model_name} via transformers + bnb 4-bit (no unsloth)...", flush=True)
        bnb = BitsAndBytesConfig(
            load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16,
            bnb_4bit_use_double_quant=True, bnb_4bit_quant_type="nf4",
        )
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        if tokenizer.pad_token_id is None:
            tokenizer.pad_token_id = tokenizer.eos_token_id
        model = AutoModelForCausalLM.from_pretrained(
            model_name, quantization_config=bnb, torch_dtype=torch.bfloat16,
            device_map="auto", trust_remote_code=True,
        )
        model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=True)
        lc = LoraConfig(r=32, lora_alpha=64, lora_dropout=0.0, bias="none",
                        target_modules=["q_proj","k_proj","v_proj","o_proj",
                                        "gate_proj","up_proj","down_proj"],
                        task_type="CAUSAL_LM")
        model = get_peft_model(model, lc)

    print(f"loading dataset {jsonl_path}...", flush=True)
    ds = load_dataset("json", data_files=str(jsonl_path), split="train")

    # Apply chat template to render the messages column to text
    def _apply_template(example):
        return {
            "text": tokenizer.apply_chat_template(
                example["messages"],
                tokenize=False,
                add_generation_prompt=False,
            )
        }
    ds = ds.map(_apply_template, remove_columns=[c for c in ds.column_names if c != "messages"])

    sft_kwargs = dict(
        output_dir=str(output_dir),
        per_device_train_batch_size=batch_size,
        gradient_accumulation_steps=grad_accum,
        num_train_epochs=epochs,
        learning_rate=lr,
        warmup_ratio=0.03,
        logging_steps=5,
        save_strategy="epoch",
        save_total_limit=1,
        bf16=True,
        max_seq_length=max_seq_length,
        dataset_text_field="text",
        report_to="none",
        push_to_hub=bool(push_to_hub),
        hub_model_id=push_to_hub,
    )
    sft_config = SFTConfig(**_safe_kwargs(SFTConfig, sft_kwargs))

    trainer_kwargs = dict(
        model=model,
        tokenizer=tokenizer,
        args=sft_config,
        train_dataset=ds,
    )
    trainer = SFTTrainer(**_safe_kwargs(SFTTrainer, trainer_kwargs))
    trainer.train()

    print(f"saving LoRA adapter to {output_dir}...", flush=True)
    model.save_pretrained(str(output_dir))
    tokenizer.save_pretrained(str(output_dir))
    if push_to_hub:
        try:
            model.push_to_hub(push_to_hub)
            tokenizer.push_to_hub(push_to_hub)
            print(f"pushed to https://huggingface.co/{push_to_hub}", flush=True)
        except Exception as e:
            print(f"WARN: push_to_hub failed: {e}", flush=True)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--episodes", type=int, default=64,
                    help="how many episodes to roll for SFT data")
    ap.add_argument("--scenarios", nargs="+",
                    default=["E0_quiet_day", "E1_release_week"],
                    help="which scenarios to roll")
    ap.add_argument("--max-steps-per-episode", type=int, default=30)
    ap.add_argument("--out-jsonl", type=str,
                    default=str(_PROJECT_ROOT / "data" / "sft_warmstart.jsonl"))
    ap.add_argument("--jsonl", type=str, default=None,
                    help="If set with --train-only, train on this file")
    ap.add_argument("--model", type=str,
                    default="unsloth/Qwen2.5-7B-Instruct-bnb-4bit")
    ap.add_argument("--output-dir", type=str,
                    default=str(_PROJECT_ROOT / "checkpoints" / "opsguard-sft"))
    ap.add_argument("--epochs", type=int, default=1)
    ap.add_argument("--batch-size", type=int, default=2)
    ap.add_argument("--grad-accum", type=int, default=8)
    ap.add_argument("--lr", type=float, default=2e-4)
    ap.add_argument("--max-seq-length", type=int, default=4096)
    ap.add_argument("--push-to-hub", type=str, default=None,
                    help="repo id like 'me/opsguard-sft' to push the LoRA")
    ap.add_argument("--gen-only", action="store_true")
    ap.add_argument("--train-only", action="store_true")
    args = ap.parse_args()

    out_jsonl = Path(args.jsonl) if (args.train_only and args.jsonl) else Path(args.out_jsonl)

    if not args.train_only:
        print(f"=== SFT data generation ===", flush=True)
        generate_dataset(
            n_episodes=args.episodes,
            scenarios=args.scenarios,
            max_steps=args.max_steps_per_episode,
            out_path=out_jsonl,
        )

    if args.gen_only:
        return

    print(f"=== SFT training ===", flush=True)
    train_sft(
        model_name=args.model,
        jsonl_path=out_jsonl,
        output_dir=Path(args.output_dir),
        epochs=args.epochs,
        batch_size=args.batch_size,
        grad_accum=args.grad_accum,
        lr=args.lr,
        max_seq_length=args.max_seq_length,
        push_to_hub=args.push_to_hub,
    )


if __name__ == "__main__":
    main()