import os
import json
import math
import time
import csv
import logging
import argparse
from typing import List, Dict, Any, Tuple, Optional

import torch
import torch.nn as nn
from tqdm import tqdm

from transformers import AutoTokenizer, AutoModel
from peft import LoraConfig, get_peft_model

# Flash SDP (Optional)
torch.backends.cuda.enable_flash_sdp(True)
torch.backends.cuda.enable_mem_efficient_sdp(True)
torch.backends.cuda.enable_math_sdp(False)

# =========================
# logging
# =========================
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s: %(message)s")
logger = logging.getLogger("infer_worker")

# =========================
# System Prompt
# =========================
SYSTEM_PROMPT = """
Act as an expert in reticular chemistry. You will receive reaction conditions as a JSON object with the fields: metal_precursor, 
organic_linker, modulator, solvent, metal_concentration_mM, M_L_ratio, temperature_C, and time_h. Based on these inputs, output exactly one uppercase label:
'P' if the conditions are likely to yield a crystalline metal-organic framework under experimental conditions, or 'N' if not..
"""

# =========================
# Data Loading
# =========================
def load_messages_from_jsonl(path: str) -> List[List[Dict[str, Any]]]:
    all_messages = []
    with open(path, "r", encoding="utf-8-sig") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            data = json.loads(line)
            if "messages" in data:
                all_messages.append(data["messages"])
    return all_messages

def standardize_json_input(json_string: str) -> str:
    try:
        data = json.loads(json_string)
        keep = {
            "metal_precursor": data.get("metal_precursor"),
            "organic_linker": data.get("organic_linker"),
            "modulator": data.get("modulator"),
            "solvent": data.get("solvent"),
            "metal_concentration_mM": data.get("metal_concentration_mM"),
            "M_L_ratio": data.get("M_L_ratio"),
            "temperature_C": data.get("temperature_C"),
            "time_h": data.get("time_h"),
        }
        return json.dumps(keep, ensure_ascii=False)
    except Exception:
        return json_string

def char_norm(s: Optional[str]) -> Optional[str]:
    s = (s or "").strip().upper()
    return s if s in ("P", "N") else None

def standardize_messages(messages: List[Dict[str, Any]]) -> Tuple[List[Dict[str, str]], Optional[str]]:
    user_content = None
    assistant_content = None
    for m in messages:
        if m.get("role") == "user":
            user_content = m.get("content")
        elif m.get("role") == "assistant":
            assistant_content = m.get("content")

    user_content = standardize_json_input(user_content or "")
    standardized = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": user_content},
    ]
    gold = char_norm(assistant_content)
    return standardized, gold

def get_user_before_last_assistant(messages):
    if not messages:
        return None, None
    last_asst_idx = None
    for i in range(len(messages) - 1, -1, -1):
        if messages[i].get("role") == "assistant":
            last_asst_idx = i
            break
    if last_asst_idx is None:
        return None, None
    gold_char = char_norm(messages[last_asst_idx].get("content"))
    user_raw = None
    for j in range(last_asst_idx - 1, -1, -1):
        if messages[j].get("role") == "user":
            user_raw = messages[j].get("content")
            break
    return user_raw, gold_char

# =========================
# Model
# =========================
def build_model_with_lora(base_model_name: str, cache_dir: str, device: torch.device):
    lora_config = LoraConfig(
        r=8,
        lora_alpha=16,
        target_modules=["q_proj", "v_proj", "o_proj"],
        lora_dropout=0.01,
        bias="none",
        task_type="SEQ_CLS"
    )

    backbone = AutoModel.from_pretrained(
        base_model_name,
        torch_dtype=torch.bfloat16,
        device_map=None,  # 强制单卡
        cache_dir=cache_dir,
    ).to(device)

    backbone = get_peft_model(backbone, lora_config)

    class QwenWithLoRAForBinaryClassification(nn.Module):
        def __init__(self, backbone):
            super().__init__()
            self.backbone = backbone
            self.config = backbone.config
            hidden_size = backbone.config.hidden_size
            self.classifier = nn.Linear(hidden_size, 2)

        def forward(self, input_ids, attention_mask=None, labels=None):
            outputs = self.backbone(
                input_ids=input_ids,
                attention_mask=attention_mask,
                return_dict=True,
            )
            mask = attention_mask.unsqueeze(-1)
            pooled = (outputs.last_hidden_state * mask).sum(1) / mask.sum(1)

            if (self.classifier.weight.device != pooled.device) or (self.classifier.weight.dtype != pooled.dtype):
                self.classifier = self.classifier.to(device=pooled.device, dtype=pooled.dtype)

            logits = self.classifier(pooled)
            loss = None
            if labels is not None:
                loss = nn.CrossEntropyLoss()(logits, labels)
            return {"loss": loss, "logits": logits}

    return QwenWithLoRAForBinaryClassification(backbone).to(device)

def load_trained_weights(model: nn.Module, ckpt_dir: str):
    candidates = [
        os.path.join(ckpt_dir, "pytorch_model.bin"),
        os.path.join(ckpt_dir, "model.safetensors"),
        os.path.join(ckpt_dir, "adapter_model.bin"),
    ]
    weight_path = None
    for p in candidates:
        if os.path.isfile(p):
            weight_path = p
            break
    if weight_path is None:
        raise FileNotFoundError(f"在 {ckpt_dir} 找不到权重文件")

    logger.info(f"Loading weights from: {weight_path}")
    if weight_path.endswith(".safetensors"):
        from safetensors.torch import load_file
        state = load_file(weight_path)
    else:
        state = torch.load(weight_path, map_location="cpu")

    missing, unexpected = model.load_state_dict(state, strict=False)
    logger.info(f"load_state_dict done. missing={len(missing)} unexpected={len(unexpected)}")

# =========================
# Inference (batch)
# =========================
@torch.inference_mode()
def predict_batch(
    model,
    tokenizer,
    device: torch.device,
    input_texts: List[str],
    max_length: int,
):
    enc = tokenizer(
        input_texts,
        truncation=True,
        max_length=max_length,
        padding="longest",
        return_tensors="pt",
    )
    enc = {k: v.pin_memory().to(device, non_blocking=True) for k, v in enc.items()}

    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
        out = model(**enc)

    logits = out["logits"] if isinstance(out, dict) else out.logits
    probs = torch.softmax(logits, dim=-1)
    prob_p = probs[:, 1].detach().float().cpu().tolist()
    pred_int = probs.argmax(dim=-1).detach().cpu().tolist()
    return pred_int, prob_p

# =========================
# progress (resume)
# =========================
def load_progress(progress_path: str) -> int:
    if not os.path.exists(progress_path):
        return -1
    try:
        with open(progress_path, "r", encoding="utf-8") as f:
            obj = json.load(f)
        return int(obj.get("last_done_index", -1))
    except Exception:
        return -1

def save_progress(progress_path: str, last_done_index: int):
    tmp = progress_path + ".tmp"
    with open(tmp, "w", encoding="utf-8") as f:
        json.dump({"last_done_index": int(last_done_index)}, f, ensure_ascii=False)
        f.flush()
        os.fsync(f.fileno())
    os.replace(tmp, progress_path)

# =========================
# worker main
# =========================
def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--base_model_name", default="openai")
    ap.add_argument("--cache_dir", default=".")
    ap.add_argument("--ckpt_dir", required=True)
    ap.add_argument("--data_jsonl", required=True)
    ap.add_argument("--out_dir", required=True)

    ap.add_argument("--batch_size", type=int, default=64)
    ap.add_argument("--max_length", type=int, default=512)

    ap.add_argument("--num_shards", type=int, required=True)
    ap.add_argument("--shard_id", type=int, required=True)

    ap.add_argument("--eval_every_steps", type=int, default=100)  # Compatible with launcher
    ap.add_argument("--flush_every_steps", type=int, default=1)  # 1=flush every step; can be larger for speed
    args = ap.parse_args()

    os.makedirs(args.out_dir, exist_ok=True)

    if not torch.cuda.is_available():
        raise RuntimeError("CUDA not available: multi-GPU inference requires GPU")
    device = torch.device("cuda")

    # tokenizer
    tokenizer = AutoTokenizer.from_pretrained(
        args.base_model_name,
        trust_remote_code=True,
        cache_dir=args.cache_dir,
        local_files_only=False,
        enable_thinking=False,
    )
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    def apply_chat_template(messages: List[Dict[str, str]]) -> str:
        return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False)

    # load data
    logger.info(f"[shard {args.shard_id}/{args.num_shards}] Loading data...")
    messages_all = load_messages_from_jsonl(args.data_jsonl)
    n = len(messages_all)
    logger.info(f"[shard {args.shard_id}] total samples={n}")

    # outputs per shard
    shard_tag = f"shard{args.shard_id:03d}_of_{args.num_shards:03d}"
    pred_all_path = os.path.join(args.out_dir, f"preds_all_{shard_tag}.jsonl")
    pred_err_path = os.path.join(args.out_dir, f"preds_errors_{shard_tag}.jsonl")
    table_csv_path = os.path.join(args.out_dir, f"preds_table_{shard_tag}.csv")
    progress_path = os.path.join(args.out_dir, f"progress_{shard_tag}.json")

    # resume
    last_done = load_progress(progress_path)  # Original index
    logger.info(f"[shard {args.shard_id}] resume last_done_index={last_done}")

    # open files append
    f_all = open(pred_all_path, "a", encoding="utf-8")
    f_err = open(pred_err_path, "a", encoding="utf-8")

    # CSV: Do not extract key values, just keep the fields you want + raw data
    need_header = not os.path.exists(table_csv_path) or os.path.getsize(table_csv_path) == 0
    f_csv = open(table_csv_path, "a", encoding="utf-8-sig", newline="")
    fieldnames = [
        "index",
        "gold",
        "pred",
        "pN",
        "pP",
        "prob_P",
        "user_raw",
        "input_text",
    ]
    w = csv.DictWriter(f_csv, fieldnames=fieldnames)
    if need_header:
        w.writeheader()
        f_csv.flush()
        os.fsync(f_csv.fileno())

    # model
    logger.info(f"[shard {args.shard_id}] Building model on {device}...")
    model = build_model_with_lora(args.base_model_name, args.cache_dir, device)
    model.eval()
    logger.info(f"[shard {args.shard_id}] Loading ckpt...")
    load_trained_weights(model, args.ckpt_dir)
    logger.info(f"[shard {args.shard_id}] ready.")

    # running accuracy
    gold_seen = 0
    correct = 0

    # iterate indices belonging to this shard
    def belongs(i: int) -> bool:
        return (i % args.num_shards) == args.shard_id and i > last_done

    indices = [i for i in range(n) if belongs(i)]
    total_local = len(indices)
    steps = math.ceil(total_local / args.batch_size)
    logger.info(f"[shard {args.shard_id}] local_samples={total_local}, steps={steps}, bs={args.batch_size}")

    try:
        pbar = tqdm(range(steps), desc=f"Infer {shard_tag}")
        for step in pbar:
            batch_idxs = indices[step * args.batch_size: (step + 1) * args.batch_size]

            # build batch texts
            input_texts = []
            user_raws = []
            golds = []

            for orig_i in batch_idxs:
                msgs = messages_all[orig_i]
                user_raw, gold_raw = get_user_before_last_assistant(msgs)

                # Continue using your original standardize_messages to ensure user content is a string
                std_msgs, gold_std = standardize_messages(msgs)
                gold = gold_std if gold_std else gold_raw

                input_text = apply_chat_template(std_msgs)

                input_texts.append(input_text)
                user_raws.append(user_raw if user_raw is not None else "")
                golds.append(gold)

            pred_int, prob_p_list = predict_batch(
                model=model,
                tokenizer=tokenizer,
                device=device,
                input_texts=input_texts,
                max_length=args.max_length,
            )

            # write results immediately
            for j, orig_i in enumerate(batch_idxs):
                pred = "P" if int(pred_int[j]) == 1 else "N"
                prob_p = float(prob_p_list[j])
                user_raw = user_raws[j]
                gold = golds[j]

                # CSV: Do not extract key values, 
                # just keep the fields you want + raw data
                row = {
                    "index": orig_i,
                    "gold": gold,
                    "pred": pred,
                    "pN": f"{(1.0 - prob_p):.8f}",
                    "pP": f"{prob_p:.8f}",
                    "prob_P": f"{prob_p:.8f}",
                    "user_raw": user_raw,
                    "input_text": input_texts[j],
                }
                w.writerow(row)

                # jsonl only if gold exists
                if gold in ("P", "N"):
                    item = {
                        "index": orig_i,
                        "gold": gold,
                        "pred": pred,
                        "prob_P": prob_p,
                        "user_raw": user_raw,
                        "input_text": input_texts[j],
                    }
                    f_all.write(json.dumps(item, ensure_ascii=False) + "\n")
                    if pred != gold:
                        f_err.write(json.dumps(item, ensure_ascii=False) + "\n")

                    gold_seen += 1
                    if pred == gold:
                        correct += 1

            # progress: last done = max index processed in this shard so far
            save_progress(progress_path, max(batch_idxs))

            # flush policy
            if (step + 1) % args.flush_every_steps == 0:
                f_all.flush(); os.fsync(f_all.fileno())
                f_err.flush(); os.fsync(f_err.fileno())
                f_csv.flush(); os.fsync(f_csv.fileno())

            # eval policy (Keep your original parameters and output)
            if (step + 1) % args.eval_every_steps == 0:
                if gold_seen > 0:
                    acc = correct / gold_seen
                    logger.info(f"[{shard_tag} step {step+1}/{steps}] running_acc={acc:.6f} (gold_seen={gold_seen})")
                else:
                    logger.info(f"[{shard_tag} step {step+1}/{steps}] running_acc=N/A (no gold)")

    finally:
        try:
            f_all.close()
        except Exception:
            pass
        try:
            f_err.close()
        except Exception:
            pass
        try:
            f_csv.close()
        except Exception:
            pass

    logger.info(f"[shard {args.shard_id}] DONE. outputs: {pred_all_path} {pred_err_path} {table_csv_path}")

if __name__ == "__main__":
    main()