import argparse
import os
import random
import re
from glob import glob

import pandas as pd
import yaml
from datasets import Dataset, load_dataset
from vllm import LLM, SamplingParams

from ctx_to_lora.data.definitions import (
    CLOSED_QA_INTX_TEMPLATES,
    RAW_DATA_DIR,
    SELF_GEN_DATA_DIR,
)
from ctx_to_lora.data.processing import (
    filter_none,
    get_preprocessing_fn,
    load_and_process_dataset,
)
from ctx_to_lora.utils import clear_gpu

STOP_STRINGS = {
    "google/gemma-2-2b-it": ["<eos>", "<end_of_turn>"],
}

# TODO: !!!!
# see https://huggingface.co/datasets/YuxinJiang/LTE_train_data/viewer/default/train?row=97&views%5B%5D=train&sql_row=2
# based on https://arxiv.org/pdf/2402.11905
SYSTEM_TEMPLATE = (
    "### SYSTEM INSTRUCTION ###\n"
    "You are an honest and helpful assistant.\n"
    "You must use the information provided in the context for responding to the question.\n"
    # "If the context does not contain enough information to answer the question, you must say so.\n"
    # "If that is the case, you can answer based on your knowledge, but you must clearly state that you are doing so.\n"
    "**DO NOT** hallucinate or make up information.\n"
    "### END OF SYSTEM INSTRUCTION ###"
)

PROMPT_TEMPLATE = "### Context ###\n{context}\n\n\n### Question ###\n{question}"

MODEL_CTX_LEN = {
    "google/gemma-2-27b-it": 8192,
    "google/gemma-2-2b-it": 8192,
    "google/gemma-2-9b-it": 8192,
}


def get_prompt(context: str, q: str) -> str:
    return PROMPT_TEMPLATE.format(context=context, question=q)


def add_closed_qa_prompt(q: str, closed_qa_prob: float = 0.1) -> str:
    if random.random() <= closed_qa_prob:
        q = random.choice(CLOSED_QA_INTX_TEMPLATES).format(input=q)
    return q


def load_config(config_path: str) -> dict:
    """Load dataset names from YAML config file."""
    with open(config_path) as f:
        config = yaml.safe_load(f)
    return config


def check_should_skip(txt: str, vllm_model: str) -> bool:
    """Check if the response should be skipped based on stop strings."""
    for stop in STOP_STRINGS[vllm_model]:
        if stop in txt[-len(stop) :]:
            return (txt.split(stop)[0], False)  # Found a valid stop string
    return (txt, True)  # No valid stop string found, skip this response


def get_dataset_configs(
    ds_names: list[str] | None,
    config: dict | None,
    split: str | None,
) -> list[tuple[str, str]]:
    assert not (ds_names and config), "Cannot provide both ds_names and config"
    if ds_names:
        assert split, "When using ds_names, --split must be provided"
        # Validate ds_names format
        for ds_name in ds_names:
            if not isinstance(ds_name, str):
                raise ValueError(f"Invalid dataset name: {ds_name}")
        return [(ds_name, split) for ds_name in ds_names]

    if config:
        dataset_configs = []

        # Process train datasets
        train_ds_names = config.get("train_ds_names", [])
        # self_gen_train_ds_names = [
        #     (ds_name.split("/")[-1], "train")
        #     for ds_name in train_ds_names
        #     if ds_name.startswith("self_gen/")
        # ]
        self_gen_train_ds_names = [
            (ds_name, "train")
            for ds_name in train_ds_names
            if ds_name.startswith("self_gen/")
        ]
        if not self_gen_train_ds_names:
            print("No self_gen datasets found in train_ds_names")
        dataset_configs.extend(self_gen_train_ds_names)

        # Process validation datasets
        val_ds_names = config.get("val_ds_names", [])
        self_gen_val_ds_names = [
            (ds_name, "validation")
            for ds_name in val_ds_names
            if ds_name.startswith("self_gen/")
        ]
        if not self_gen_val_ds_names:
            print("No self_gen datasets found in val_ds_names")
        dataset_configs.extend(self_gen_val_ds_names)

        return dataset_configs


def create_messages(
    ctxs: list[str], questions: list[list[str]], vllm_model: str, system_template: str
) -> list[list[dict]]:
    """Create chat messages for the model."""
    if "gemma" in vllm_model:
        # gemma models do not support system messages
        return [
            [
                {
                    "role": "user",
                    "content": system_template + "\n\n\n" + get_prompt(ctx, q),
                }
            ]
            for ctx, q_list in zip(ctxs, questions)
            for q in q_list
        ]
    else:
        return [
            [
                {"role": "system", "content": system_template},
                {"role": "user", "content": get_prompt(ctx, q)},
            ]
            for ctx, q_list in zip(ctxs, questions)
            for q in q_list
        ]


def self_generate(
    ds_name: str,
    split: str,
    args: argparse.Namespace,
    llm: LLM,
    system_template: str,
    parquet_file: str | None = None,
) -> None:
    """Process a single dataset and generate QA pairs."""

    shard_name = ""
    temp = 0.0
    closed_qa_prob = 0.0
    if ds_name is not None and "_temp_" in ds_name:
        temp_match = re.search(r"_temp_([\d.]+)", ds_name)
        if temp_match:
            temp = float(temp_match.group(1))
    if ds_name is not None and "_closed_qa_prob_" in ds_name:
        prob_match = re.search(r"_closed_qa_prob_([\d.]+)", ds_name)
        if prob_match:
            closed_qa_prob = float(prob_match.group(1))
    print(f"Processing dataset: {ds_name}, split: {split}")
    print(f"Using temperature: {temp}")
    print(f"Using closed QA prompt probability: {closed_qa_prob}")

    if parquet_file:
        print(f"Loading dataset from parquet file: {parquet_file}")

        split = "train"
        ds_name = "/".join(parquet_file.split(RAW_DATA_DIR)[-1].split("/")[:-1])

        shard_name = "_" + os.path.basename(parquet_file).replace(".parquet", "")
        ds = load_dataset(path="parquet", data_files=[parquet_file], split="train")
        processing_fn = get_preprocessing_fn(ds_name, is_eval=False, is_pretrain=False)
        ds = ds.map(processing_fn, num_proc=8)

    else:
        ds_name = ds_name.split("/")[-1]  # Extract just the dataset name

        print(f"Loading dataset: {ds_name} with split: {split}")
        kwargs = dict(
            ds_name=ds_name,
            split=split,
            add_negative_prompt=False,
            add_repeat_prompt=False,
            repeat_prob=0,
            is_pretrain=False,
        )

        ds = load_and_process_dataset(**kwargs, streaming=False, num_proc=8)
    print(f"Loaded dataset: {ds_name} with split: {split}")

    if args.debug:
        ds = ds.take(10)

    ds = ds.filter(filter_none, batched=False, num_proc=8)

    ctxs = [sample["context"] for sample in ds]
    questions = [
        [add_closed_qa_prompt(q, closed_qa_prob) for q in sample["prompts"] if q]
        for sample in ds
    ]
    print(f"Loaded {len(ctxs)} contexts and {len(questions)} questions")

    messages = create_messages(ctxs, questions, args.vllm_model, SYSTEM_TEMPLATE)

    print(f"Generating from {len(messages)} contexts")
    # TODO (distillation): make vllm outputs logits here too
    completions = llm.chat(
        messages,
        sampling_params=SamplingParams(
            max_tokens=1024,
            temperature=temp,
            # needed for checking if stop tokens are present
            skip_special_tokens=False,
            include_stop_str_in_output=True,
        ),
    )

    self_gen_data = {ctx: {"prompts": [], "responses": []} for ctx in ctxs}
    c = 0
    n_skips = 0
    for ctx, q_list in zip(ctxs, questions):
        for i, q in enumerate(q_list):
            response = completions[c + i].outputs[0].text
            response, skip = check_should_skip(response, args.vllm_model)
            # skip = True
            # for stop in STOP_STRINGS[args.vllm_model]:
            #     if stop == response[-len(stop) :]:
            #         # Check if response ends with stop string
            #         response = response.split(stop)[0]
            #         skip = False
            #         break
            if skip:
                print(f"Skipping due to missing stop string")
                n_skips += 1
                continue

            self_gen_data[ctx]["prompts"].append(q)
            self_gen_data[ctx]["responses"].append(response)
        c += i + 1

    print(f"Skipped {n_skips} responses due to missing stop strings")
    samples = [
        {
            "context": ctx,
            "prompts": q_list,
            "responses": self_gen_data[ctx]["responses"],
        }
        for ctx, q_list in zip(ctxs, questions)
    ]

    if args.debug:
        for sample in samples:
            print(f"context={sample['context']}")
            print(f"prompt={sample['prompts']}")
            print(f"response={sample['responses']}")
            print("=" * 80)

    print(f"Generated {len(samples)} samples")
    # random.shuffle(samples)

    # Save results
    if not args.debug:
        df = pd.DataFrame(samples)
        ds_out = Dataset.from_pandas(df)
        fpath = f"{SELF_GEN_DATA_DIR}/{args.vllm_model}_temp_{temp}_closed_qa_prob_{closed_qa_prob}/{ds_name}/{split}/ds{shard_name}"
        os.makedirs(os.path.dirname(fpath), exist_ok=True)
        fpath = f"{fpath}.parquet"
        ds_out.to_parquet(fpath)
        print(f"Saved to {fpath}")

        # Cleanup
        del samples, df, ds_out, completions, messages, ctxs, questions
        clear_gpu()


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Generate QA pairs using VLLM")
    parser.add_argument(
        "--vllm_model",
        type=str,
        required=True,
        help="VLLM model name (e.g., google/gemma-2-2b-it)",
    )
    parser.add_argument(
        "--debug",
        action="store_true",
        help="Enable debug mode (process only 10 samples)",
    )

    # Either config file OR ds_names + split
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument(
        "--config",
        type=str,
        help="Path to YAML config file with train_ds_names/val_ds_names",
    )
    group.add_argument(
        "--ds_names",
        type=str,
        nargs="+",
        help="List of dataset names/shard patterns",
    )
    group.add_argument(
        "--glob_pattern",
        type=str,
        help="Glob pattern to match dataset names (e.g., 'data/raw_datasets/fw_qa_3/*')",
    )

    parser.add_argument(
        "--split",
        type=str,
        help="Dataset split to use when using --ds_names (required with --ds_names)",
    )
    return parser.parse_args()


if __name__ == "__main__":
    args = parse_args()

    # Validate arguments
    if args.ds_names and not args.split:
        raise ValueError("--split is required when using --ds_names")

    vllm_model = args.vllm_model
    print(f"Using model: {vllm_model}")

    # Setup model-specific configurations
    llm_kwargs = dict(
        model=vllm_model,
        dtype="bfloat16",
        enable_prefix_caching=True,
        enable_chunked_prefill=True,
        max_model_len=MODEL_CTX_LEN.get(vllm_model),
    )

    print(f"{llm_kwargs=}")
    llm = LLM(**llm_kwargs)

    # Get dataset configs from config or CLI args
    config = load_config(args.config) if args.config else None
    if args.ds_names or args.config:
        dataset_configs = get_dataset_configs(
            ds_names=args.ds_names,
            config=config,
            split=args.split,
        )

        # Process each dataset
        for ds_name, split in dataset_configs:
            print(f"Processing dataset: {ds_name}, split: {split}")
            self_generate(ds_name, split, args, llm, SYSTEM_TEMPLATE)
    else:
        assert args.glob_pattern, (
            "glob_pattern must be provided if no ds_names or config"
        )
        files = glob(args.glob_pattern)
        for file in files:
            print(f"Processing file: {file}")
            self_generate(
                ds_name=None,
                parquet_file=file,
                split=args.split,
                args=args,
                llm=llm,
                system_template=SYSTEM_TEMPLATE,
            )