import json
from dataclasses import dataclass
from typing import Iterable, List, Dict, Optional, Tuple, Union

import torch
from torch.utils.data import Dataset, DataLoader

from ..loader import load_model_and_tokenizer
from .. import (
    inject_lora,
    inject_adapter,
    inject_bitfit,
    inject_prompt_tuning,
    PEFTTrainer,
)


class QADataset(Dataset):
    def __init__(self, pairs: List[Dict[str, str]]):
        self.pairs = pairs

    def __len__(self) -> int:
        return len(self.pairs)

    def __getitem__(self, idx: int) -> Dict[str, str]:
        return self.pairs[idx]


def _format_qa_example(question: str, answer: str) -> str:
    return f"Question: {question}\nAnswer: {answer}"


def load_qa_from_jsonl(path: str, question_field: str = "question", answer_field: str = "answer") -> List[Dict[str, str]]:
    pairs: List[Dict[str, str]] = []
    with open(path, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            obj = json.loads(line)
            if question_field not in obj or answer_field not in obj:
                continue
            pairs.append({"question": obj[question_field], "answer": obj[answer_field]})
    return pairs


def load_qa_from_csv(path: str, question_field: str = "question", answer_field: str = "answer", delimiter: str = ",") -> List[Dict[str, str]]:
    import csv

    pairs: List[Dict[str, str]] = []
    with open(path, "r", encoding="utf-8") as f:
        reader = csv.DictReader(f, delimiter=delimiter)
        for row in reader:
            if question_field not in row or answer_field not in row:
                continue
            pairs.append({"question": row[question_field], "answer": row[answer_field]})
    return pairs


def load_qa_from_hf_dataset(
    dataset: Optional["datasets.Dataset"] = None,
    dataset_name: Optional[str] = None,
    split: str = "train",
    question_field: str = "question",
    answer_field: str = "answer",
) -> List[Dict[str, str]]:
    if dataset is None:
        if dataset_name is None:
            raise ValueError("Either `dataset` or `dataset_name` must be provided.")
        try:
            from datasets import load_dataset
        except ImportError as exc:
            raise ImportError("`datasets` library is required to load HuggingFace datasets") from exc
        dataset = load_dataset(dataset_name, split=split)

    pairs: List[Dict[str, str]] = []
    for row in dataset:
        if question_field not in row or answer_field not in row:
            continue
        pairs.append({"question": row[question_field], "answer": row[answer_field]})
    return pairs


@dataclass
class QAConfig:
    model_name: str
    peft_type: str = "lora"  # "lora", "adapters", "bitfit", "prompt_tuning"
    lr: float = 1e-4
    batch_size: int = 4
    num_epochs: int = 1
    max_length: int = 256
    device: Optional[str] = None
    lora_r: int = 8
    lora_alpha: int = 16
    lora_dropout: float = 0.0
    target_modules: Optional[List[str]] = None
    match_mode: str = "contains"


def _build_dataloader(pairs: List[Dict[str, str]], tokenizer, max_length: int, batch_size: int) -> DataLoader:
    dataset = QADataset(pairs)

    def collate(batch: List[Dict[str, str]]):
        texts = [_format_qa_example(ex["question"], ex["answer"]) for ex in batch]
        enc = tokenizer(
            texts,
            padding=True,
            truncation=True,
            max_length=max_length,
            return_tensors="pt",
        )
        input_ids = enc["input_ids"]
        attention_mask = enc.get("attention_mask")
        labels = input_ids.clone()
        return {"input_ids": input_ids, "attention_mask": attention_mask, "labels": labels}

    return DataLoader(dataset, batch_size=batch_size, shuffle=True, collate_fn=collate)


def _inject_peft(model, config: QAConfig):
    peft_type = config.peft_type.lower()
    target_modules = config.target_modules

    if peft_type == "lora":
        return inject_lora(
            model,
            r=config.lora_r,
            alpha=config.lora_alpha,
            target_modules=target_modules or ["c_attn", "c_fc", "c_proj", "linear", "fc"],
            match_mode=config.match_mode,
            dropout=config.lora_dropout,
        )
    if peft_type == "adapters":
        return inject_adapter(
            model,
            bottleneck=64,
            target_modules=target_modules or ["mlp", "fc1", "linear", "fc"],
            match_mode=config.match_mode,
        )
    if peft_type == "bitfit":
        return inject_bitfit(model)
    if peft_type == "prompt_tuning":
        wrapped, peft = inject_prompt_tuning(model, length=20)
        return wrapped, peft
    raise ValueError(f"Unknown peft_type: {config.peft_type}")


def fine_tune_qa(
    pairs: List[Dict[str, str]],
    config: QAConfig,
) -> Tuple[torch.nn.Module, object, object]:
    model: Optional[torch.nn.Module] = None
    tokenizer = None

    # If the caller wants a completely custom stack, they can pass their own
    # model and tokenizer and call inject_* + PEFTTrainer manually. This
    # helper focuses on the common case where we load a HF model by name.
    model, tokenizer = load_model_and_tokenizer(config.model_name, device=config.device)

    if getattr(tokenizer, "pad_token", None) is None:
        eos = getattr(tokenizer, "eos_token", None)
        if eos is not None:
            tokenizer.pad_token = eos
        else:
            tokenizer.add_special_tokens({"pad_token": "[PAD]"})
            model.resize_token_embeddings(len(tokenizer))

    model, peft = _inject_peft(model, config)

    pefts = peft if isinstance(peft, list) else [peft]
    trainer = PEFTTrainer(model, pefts, lr=config.lr)

    dataloader = _build_dataloader(pairs, tokenizer, config.max_length, config.batch_size)

    for _ in range(config.num_epochs):
        for batch in dataloader:
            batch = {k: v.to(trainer.device) if isinstance(v, torch.Tensor) else v for k, v in batch.items()}
            trainer.train_step(batch)

    return model, peft, tokenizer
