import argparse
import hashlib
import json
import math
import os
import random
import re
from collections import Counter, defaultdict
from typing import Any, Dict, Iterable, List, Tuple

import numpy as np


def iter_jsonl(path):
    with open(path, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line:
                yield json.loads(line)


def write_jsonl(path, rows):
    os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
    with open(path, "w", encoding="utf-8") as f:
        for row in rows:
            f.write(json.dumps(row, ensure_ascii=False) + "\n")


def stable_id(text):
    return hashlib.md5(text.encode("utf-8")).hexdigest()[:12]


def as_list(value):
    if value is None:
        return []
    if isinstance(value, list):
        return value
    return [value]


def text_of(value):
    if value is None:
        return ""
    if isinstance(value, str):
        return value
    if isinstance(value, dict):
        for key in ("text", "content", "paragraph_text", "passage_text", "snippet", "body"):
            if key in value:
                return str(value[key])
    return str(value)


def tokenize(text):
    text = str(text or "").lower()
    zh = re.findall(r"[\u4e00-\u9fff]", text)
    words = re.findall(r"[a-z0-9]+", text)
    return zh + words


def normalize_answers(raw):
    answers = []
    for item in as_list(raw):
        if isinstance(item, dict):
            for key in ("text", "answer", "value"):
                if key in item and str(item[key]).strip():
                    answers.append(str(item[key]).strip())
        elif str(item).strip():
            answers.append(str(item).strip())
    seen = set()
    dedup = []
    for answer in answers:
        key = answer.lower()
        if key not in seen:
            seen.add(key)
            dedup.append(answer)
    return dedup


def normalize_context_item(item, fallback_id):
    if isinstance(item, str):
        return {"id": fallback_id, "title": "", "text": item}
    doc_id = str(
        item.get("id")
        or item.get("doc_id")
        or item.get("pid")
        or item.get("passage_id")
        or item.get("wikipedia_id")
        or item.get("source")
        or fallback_id
    )
    title = str(item.get("title") or item.get("document_title") or item.get("name") or "")
    text = text_of(item)
    return {"id": doc_id, "title": title, "text": text}


def extract_contexts(sample):
    candidates = (
        sample.get("contexts")
        or sample.get("documents")
        or sample.get("snippets")
        or sample.get("search_results")
        or sample.get("ctxs")
        or sample.get("passages")
        or sample.get("positive_ctxs")
        or sample.get("retrieved_passages")
        or []
    )
    contexts = []
    for i, item in enumerate(candidates):
        doc = normalize_context_item(item, f"D{i + 1}")
        if doc["text"].strip():
            contexts.append(doc)
    return contexts


def extract_gold_doc_ids(sample, contexts):
    explicit = sample.get("gold_doc_ids") or sample.get("citations") or sample.get("supporting_doc_ids")
    context_ids = {doc["id"] for doc in contexts}
    if explicit:
        matched = [str(x) for x in as_list(explicit) if str(x) in context_ids]
        if matched:
            return matched
    ids = []
    for doc in contexts:
        text = doc["text"]
        if any(answer and answer in text for answer in extract_answers(sample)):
            ids.append(doc["id"])
    if not ids and contexts:
        ids.append(contexts[0]["id"])
    return ids[:3]


def keep_gold_contexts_first(contexts, gold_doc_ids, max_contexts):
    if max_contexts <= 0 or len(contexts) <= max_contexts:
        return contexts
    gold = set(str(x) for x in gold_doc_ids)
    positives = [doc for doc in contexts if doc["id"] in gold]
    others = [doc for doc in contexts if doc["id"] not in gold]
    return (positives + others)[:max_contexts]


def extract_question(sample):
    return str(
        sample.get("question")
        or sample.get("query")
        or sample.get("prompt")
        or sample.get("input")
        or sample.get("id")
        or ""
    ).strip()


def extract_answers(sample):
    raw = sample.get("answers")
    if raw is None:
        raw = sample.get("answer")
    if raw is None:
        raw = sample.get("short_answers")
    if raw is None:
        raw = sample.get("output")
    return normalize_answers(raw)


def normalize_sample(sample, source_name, index):
    question = extract_question(sample)
    answers = extract_answers(sample)
    contexts = extract_contexts(sample)
    if not question or not answers or not contexts:
        return None
    gold_doc_ids = extract_gold_doc_ids(sample, contexts)
    return {
        "id": str(sample.get("id") or f"{source_name}_{index}_{stable_id(question)}"),
        "source": source_name,
        "question": question,
        "answer": answers[0],
        "answers": answers,
        "gold_doc_ids": gold_doc_ids,
        "contexts": contexts,
    }


def collect_corpus(samples):
    corpus = {}
    for sample in samples:
        for doc in sample["contexts"]:
            if doc["id"] not in corpus and doc["text"].strip():
                corpus[doc["id"]] = doc
    return list(corpus.values())


def build_idf(corpus):
    df = defaultdict(int)
    for doc in corpus:
        for tok in set(tokenize(doc["title"] + " " + doc["text"])):
            df[tok] += 1
    n_docs = max(len(corpus), 1)
    return {tok: math.log((n_docs - freq + 0.5) / (freq + 0.5) + 1.0) for tok, freq in df.items()}


def bm25_score(query_tokens, doc_tokens, idf, avgdl, k1=1.2, b=0.75):
    if not doc_tokens:
        return 0.0
    tf = Counter(doc_tokens)
    score = 0.0
    dl = len(doc_tokens)
    denom_base = k1 * (1 - b + b * dl / max(avgdl, 1.0))
    for tok in query_tokens:
        freq = tf.get(tok, 0)
        if freq:
            score += idf.get(tok, 0.0) * freq * (k1 + 1) / (freq + denom_base)
    return score


def add_bm25_hard_negatives(samples, num_negatives, max_contexts, seed):
    if num_negatives <= 0:
        return samples
    rng = random.Random(seed)
    corpus = collect_corpus(samples)
    if not corpus:
        return samples
    idf = build_idf(corpus)
    doc_tokens = {doc["id"]: tokenize(doc["title"] + " " + doc["text"]) for doc in corpus}
    avgdl = sum(len(toks) for toks in doc_tokens.values()) / max(len(doc_tokens), 1)
    doc_by_id = {doc["id"]: doc for doc in corpus}

    for sample in samples:
        existing = {doc["id"] for doc in sample["contexts"]}
        gold = set(sample["gold_doc_ids"])
        query_tokens = tokenize(sample["question"] + " " + sample["answer"])
        scored = []
        for doc in corpus:
            if doc["id"] in existing or doc["id"] in gold:
                continue
            score = bm25_score(query_tokens, doc_tokens[doc["id"]], idf, avgdl)
            if score > 0:
                scored.append((score, doc["id"]))
        scored.sort(reverse=True)
        top_ids = [doc_id for _, doc_id in scored[: num_negatives * 4]]
        rng.shuffle(top_ids)
        negatives = [doc_by_id[doc_id] for doc_id in top_ids[:num_negatives]]
        sample["contexts"] = keep_gold_contexts_first(sample["contexts"] + negatives, sample["gold_doc_ids"], max_contexts)
    return samples


def add_embedding_hard_negatives(samples, num_negatives, max_contexts, seed, model_name, batch_size):
    if num_negatives <= 0:
        return samples
    try:
        from sentence_transformers import SentenceTransformer
    except ImportError as exc:
        raise ImportError("Embedding negative mining requires sentence_transformers. Install project requirements first.") from exc

    rng = random.Random(seed)
    corpus = collect_corpus(samples)
    if not corpus:
        return samples

    model = SentenceTransformer(model_name)
    doc_texts = [(doc["title"] + "\n" + doc["text"]).strip() for doc in corpus]
    query_texts = [(sample["question"] + "\n" + sample["answer"]).strip() for sample in samples]
    doc_emb = model.encode(
        doc_texts,
        batch_size=batch_size,
        show_progress_bar=True,
        normalize_embeddings=True,
    )
    query_emb = model.encode(
        query_texts,
        batch_size=batch_size,
        show_progress_bar=True,
        normalize_embeddings=True,
    )
    doc_emb = np.asarray(doc_emb, dtype=np.float32)
    query_emb = np.asarray(query_emb, dtype=np.float32)
    scores = query_emb @ doc_emb.T

    for i, sample in enumerate(samples):
        existing = {doc["id"] for doc in sample["contexts"]}
        gold = set(sample["gold_doc_ids"])
        ranked = np.argsort(-scores[i])
        candidate_ids = []
        for doc_idx in ranked:
            doc_id = corpus[int(doc_idx)]["id"]
            if doc_id in existing or doc_id in gold:
                continue
            candidate_ids.append(doc_idx)
            if len(candidate_ids) >= num_negatives * 4:
                break
        rng.shuffle(candidate_ids)
        negatives = [corpus[int(doc_idx)] for doc_idx in candidate_ids[:num_negatives]]
        sample["contexts"] = keep_gold_contexts_first(sample["contexts"] + negatives, sample["gold_doc_ids"], max_contexts)
    return samples


def add_hard_negatives(samples, num_negatives, max_contexts, seed, method="bm25", embedding_model="", embedding_batch_size=64):
    if method == "bm25":
        return add_bm25_hard_negatives(samples, num_negatives, max_contexts, seed)
    if method == "embedding":
        return add_embedding_hard_negatives(samples, num_negatives, max_contexts, seed, embedding_model, embedding_batch_size)
    raise ValueError(f"Unsupported negative_mining method: {method}")


def split_rows(rows, train_size, dev_size, test_size, seed):
    rng = random.Random(seed)
    rows = rows[:]
    rng.shuffle(rows)
    if train_size:
        train = rows[:train_size]
        rest = rows[train_size:]
    else:
        train, rest = rows, []
    dev = rest[:dev_size] if dev_size else []
    rest = rest[dev_size:] if dev_size else rest
    test = rest[:test_size] if test_size else []
    return train, dev, test


def load_sources(paths):
    rows = []
    for path in paths:
        source_name = os.path.splitext(os.path.basename(path))[0]
        for idx, sample in enumerate(iter_jsonl(path)):
            row = normalize_sample(sample, source_name, idx)
            if row is not None:
                rows.append(row)
    return rows


def main():
    parser = argparse.ArgumentParser(description="Build SearchShortQA JSONL from public QA/RAG style datasets")
    parser.add_argument("--inputs", nargs="+", required=True, help="One or more local JSONL files")
    parser.add_argument("--output_dir", default="dataset/search_shortqa")
    parser.add_argument("--train_size", type=int, default=30000)
    parser.add_argument("--dev_size", type=int, default=1000)
    parser.add_argument("--test_size", type=int, default=1000)
    parser.add_argument("--max_contexts", type=int, default=6)
    parser.add_argument("--num_negatives", type=int, default=3)
    parser.add_argument("--negative_mining", choices=["bm25", "embedding"], default="bm25")
    parser.add_argument("--embedding_model", default="BAAI/bge-small-zh-v1.5")
    parser.add_argument("--embedding_batch_size", type=int, default=64)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--limit", type=int, default=0)
    args = parser.parse_args()

    rows = load_sources(args.inputs)
    if args.limit:
        rows = rows[: args.limit]
    rows = add_hard_negatives(
        rows,
        args.num_negatives,
        args.max_contexts,
        args.seed,
        method=args.negative_mining,
        embedding_model=args.embedding_model,
        embedding_batch_size=args.embedding_batch_size,
    )
    rows = [row for row in rows if row["question"] and row["answers"] and row["contexts"]]
    train, dev, test = split_rows(rows, args.train_size, args.dev_size, args.test_size, args.seed)

    os.makedirs(args.output_dir, exist_ok=True)
    write_jsonl(os.path.join(args.output_dir, "search_shortqa_train.jsonl"), train)
    write_jsonl(os.path.join(args.output_dir, "search_shortqa_dev.jsonl"), dev)
    write_jsonl(os.path.join(args.output_dir, "search_shortqa_test.jsonl"), test)

    manifest = {
        "inputs": args.inputs,
        "total_normalized": len(rows),
        "train": len(train),
        "dev": len(dev),
        "test": len(test),
        "max_contexts": args.max_contexts,
        "num_negatives": args.num_negatives,
        "negative_mining": args.negative_mining,
        "embedding_model": args.embedding_model if args.negative_mining == "embedding" else "",
        "seed": args.seed,
        "format": "SearchShortQA v1",
    }
    with open(os.path.join(args.output_dir, "manifest.json"), "w", encoding="utf-8") as f:
        json.dump(manifest, f, ensure_ascii=False, indent=2)
    print(json.dumps(manifest, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()