"""
Dataset for custom fine-tuning pairs (JSON or JSONL).
Expected fields: prompt, code, optional language.
"""

from __future__ import annotations

import json
from pathlib import Path
from typing import Dict, List

from torch.utils.data import Dataset

from src.tokenizer.code_tokenizer import CodeTokenizer


class CustomPairDataset(Dataset):
    def __init__(self, path: str, tokenizer: CodeTokenizer, max_seq_len: int = 512) -> None:
        self.path = Path(path)
        if not self.path.exists():
            raise FileNotFoundError(f"Custom fine-tune data file not found: {self.path}")
        self.tokenizer = tokenizer
        self.max_seq_len = max_seq_len
        self.rows: List[List[int]] = []
        self._load()

    def _load(self) -> None:
        if self.path.suffix.lower() == ".jsonl":
            data = []
            for line in self.path.read_text(encoding="utf-8-sig").splitlines():
                line = line.strip().lstrip("\ufeff")
                if not line:
                    continue
                data.append(json.loads(line))
        elif self.path.suffix.lower() == ".json":
            raw = json.loads(self.path.read_text(encoding="utf-8-sig"))
            if isinstance(raw, dict) and "data" in raw:
                data = raw["data"]
            elif isinstance(raw, list):
                data = raw
            else:
                raise ValueError("JSON fine-tune file must be a list or {'data': [...]}.")
        else:
            raise ValueError("Custom fine-tune file must be .json or .jsonl")

        for row in data:
            prompt = str(row.get("prompt", "")).strip()
            code = str(row.get("code", "")).strip()
            language = str(row.get("language", "python")).strip().lower() or "python"
            if not prompt or not code:
                continue
            text = self.tokenizer.format_training_sample(prompt=prompt, code=code, language=language)
            ids = self.tokenizer.encode(text)[: self.max_seq_len]
            if len(ids) >= 8:
                self.rows.append(ids)

        if not self.rows:
            raise ValueError("No valid samples found in custom fine-tune data.")

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

    def __getitem__(self, idx: int) -> List[int]:
        return self.rows[idx]