# DocumentsReasoningDataset.py

from __future__ import annotations

import json
import random
import re
from pathlib import Path
from typing import Any, Optional

from datasets import load_dataset

try:
    import utils.configs_loader as ProjectConfigs
except Exception:
    ProjectConfigs = None


# ------------------------------------------------------------
# Split config
# ------------------------------------------------------------

SAMPLE_SPLITS = {
    "easy": 0.25,          # only supporting documents
    "normal": 0.25,        # original MuSiQue context
    "messy": 0.20,         # support + distractors, shuffled
    "messy_essay": 0.10,   # uglier real-world-ish document format
    "missing_one": 0.15,   # remove one required support document
    "missing_all": 0.05,   # remove all support documents
}


DEFAULT_SYSTEM_PROMPT = (
    "You are a concise document reasoning assistant. "
    "Use only the provided <document> tags as evidence. "
    "Reason inside <think> and </think>. "
    "If required evidence is missing, answer exactly: Not enough information. "
    "Keep reasoning short, grounded, and cite document ids when useful."
)


class DocumentsReasoningDataset:
    def __init__(
        self,
        path: str = "dataset/documents_reasoning",
        hf_dataset: str = "dgslibisey/MuSiQue",
        train_split: str = "train",
        validation_split: str = "validation",
        samples_split: dict[str, float] = SAMPLE_SPLITS,
        target_train_samples: int = 35_000,
        target_validation_samples: int = 1_000,
        seed: int = 42,
        max_distractors: int = 8,
        hide_title_prob: float = 0.20,
    ):
        self.path = Path(path)
        self.hf_dataset = hf_dataset
        self.train_split = train_split
        self.validation_split = validation_split

        self.samples_split = self._normalize_splits(samples_split)
        self.target_train_samples = target_train_samples
        self.target_validation_samples = target_validation_samples

        self.seed = seed
        self.rng = random.Random(seed)

        self.max_distractors = max_distractors
        self.hide_title_prob = hide_title_prob

        self.prompt = self._load_system_prompt()

    # ------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------

    def build(self) -> dict[str, Any]:
        """
        Builds:

            dataset/documents_reasoning/train.jsonl
            dataset/documents_reasoning/validation.jsonl
            dataset/documents_reasoning/validation_stress.jsonl

        Output format per line:

            {
                "messages": [
                    {"role": "system", "content": "..."},
                    {"role": "user", "content": "..."},
                    {"role": "assistant", "content": "..."}
                ],
                "meta": {...}
            }
        """

        self.path.mkdir(parents=True, exist_ok=True)

        print(f"[Dataset] Loading {self.hf_dataset} ...")
        ds = load_dataset(self.hf_dataset)

        train_ds = ds[self.train_split]
        val_ds = ds[self.validation_split]

        print("[Dataset] Building train records...")
        train_records = self._build_records(
            dataset=train_ds,
            target_samples=self.target_train_samples,
            split_name="train",
            sample_splits=self.samples_split,
        )

        print("[Dataset] Building validation records...")
        validation_records = self._build_records(
            dataset=val_ds,
            target_samples=self.target_validation_samples,
            split_name="validation",
            sample_splits={
                "normal": 0.70,
                "easy": 0.20,
                "messy": 0.10,
            },
        )

        print("[Dataset] Building stress validation records...")
        validation_stress_records = self._build_records(
            dataset=val_ds,
            target_samples=min(500, self.target_validation_samples),
            split_name="validation_stress",
            sample_splits={
                "missing_one": 0.60,
                "missing_all": 0.25,
                "messy_essay": 0.15,
            },
        )

        train_path = self.path / "train.jsonl"
        validation_path = self.path / "validation.jsonl"
        stress_path = self.path / "validation_stress.jsonl"

        self._write_jsonl(train_path, train_records)
        self._write_jsonl(validation_path, validation_records)
        self._write_jsonl(stress_path, validation_stress_records)

        summary = {
            "train_path": str(train_path),
            "validation_path": str(validation_path),
            "validation_stress_path": str(stress_path),
            "train_samples": len(train_records),
            "validation_samples": len(validation_records),
            "validation_stress_samples": len(validation_stress_records),
        }

        print("[Dataset] Done.")
        print(json.dumps(summary, indent=2, ensure_ascii=False))

        return summary

    # ------------------------------------------------------------
    # Record building
    # ------------------------------------------------------------

    def _build_records(
        self,
        dataset,
        target_samples: int,
        split_name: str,
        sample_splits: dict[str, float],
    ) -> list[dict[str, Any]]:
        records: list[dict[str, Any]] = []
        counts = self._target_counts(target_samples, sample_splits)

        all_indices = list(range(len(dataset)))

        for kind, target_count in counts.items():
            made = 0
            attempts = 0

            indices = all_indices[:]
            self.rng.shuffle(indices)

            print(f"[Build:{split_name}] kind={kind}, target={target_count}")

            while made < target_count and attempts < target_count * 30:
                idx = indices[attempts % len(indices)]
                sample = dataset[idx]

                attempts += 1

                record = self._make_variant(
                    sample=sample,
                    kind=kind,
                    split_name=split_name,
                )

                if record is None:
                    continue

                records.append(record)
                made += 1

            if made < target_count:
                print(
                    f"[WARN] kind={kind}: requested {target_count}, made {made}. "
                    f"This usually means many missing-info variants leaked answers."
                )

        self.rng.shuffle(records)
        return records

    def _make_variant(
        self,
        sample: dict[str, Any],
        kind: str,
        split_name: str,
    ) -> Optional[dict[str, Any]]:
        sample_id = self._scalar(sample.get("id", ""))

        paragraphs = self._records_from_nested(sample.get("paragraphs", []))
        subquestions = self._records_from_nested(sample.get("question_decomposition", []))

        question = self._scalar(sample.get("question", ""))
        answer = self._scalar(sample.get("answer", ""))
        answer_aliases = [self._scalar(x) for x in self._as_list(sample.get("answer_aliases", []))]
        answerable = bool(sample.get("answerable", True))

        if not paragraphs or not subquestions or not question:
            return None

        paragraph_by_idx = self._paragraph_by_idx(paragraphs)
        support_idxs = self._support_indices(subquestions)

        support_paragraphs = [
            paragraph_by_idx[idx]
            for idx in support_idxs
            if idx in paragraph_by_idx
        ]

        distractor_paragraphs = [
            p for p in paragraphs
            if int(p["idx"]) not in support_idxs
        ]

        if not support_paragraphs:
            return None

        hide_titles = self.rng.random() < self.hide_title_prob
        removed_support_idxs: set[int] = set()

        if kind == "easy":
            kept = support_paragraphs[:]
            self.rng.shuffle(kept)
            context_style = "clean"

        elif kind == "normal":
            kept = paragraphs[:]
            context_style = "clean"

        elif kind == "messy":
            kept = self._sample_support_plus_distractors(
                support=support_paragraphs,
                distractors=distractor_paragraphs,
                max_distractors=self.max_distractors,
            )
            self.rng.shuffle(kept)
            context_style = "clean"

        elif kind == "messy_essay":
            kept = self._sample_support_plus_distractors(
                support=support_paragraphs,
                distractors=distractor_paragraphs,
                max_distractors=self.max_distractors,
            )
            self.rng.shuffle(kept)
            context_style = "essay"
            hide_titles = self.rng.random() < 0.50

        elif kind == "missing_one":
            if len(support_paragraphs) < 2:
                return None

            removed = self.rng.choice(support_paragraphs)
            removed_idx = int(removed["idx"])
            removed_support_idxs.add(removed_idx)

            kept_support = [
                p for p in support_paragraphs
                if int(p["idx"]) != removed_idx
            ]

            kept = self._sample_support_plus_distractors(
                support=kept_support,
                distractors=distractor_paragraphs,
                max_distractors=min(self.max_distractors, 6),
            )
            self.rng.shuffle(kept)
            context_style = "clean"

            if self._missing_answer_leaks(
                subquestions=subquestions,
                kept_paragraphs=kept,
                kept_support_idxs={int(p["idx"]) for p in kept},
                final_answer=answer,
                answer_aliases=answer_aliases,
            ):
                return None

        elif kind == "missing_all":
            if len(distractor_paragraphs) < 2:
                return None

            kept = self.rng.sample(
                distractor_paragraphs,
                k=min(len(distractor_paragraphs), 6),
            )
            self.rng.shuffle(kept)

            removed_support_idxs = set(support_idxs)
            context_style = "clean"

            if self._missing_answer_leaks(
                subquestions=subquestions,
                kept_paragraphs=kept,
                kept_support_idxs={int(p["idx"]) for p in kept},
                final_answer=answer,
                answer_aliases=answer_aliases,
            ):
                return None

        else:
            raise ValueError(f"Unknown sample kind: {kind}")

        kept_idxs = {int(p["idx"]) for p in kept}

        user_prompt = self._make_user_prompt(
            question=question,
            paragraphs=kept,
            context_style=context_style,
            hide_titles=hide_titles,
        )

        assistant_content = self._make_assistant_response(
            answer=answer,
            answerable=answerable,
            subquestions=subquestions,
            paragraph_by_idx=paragraph_by_idx,
            kept_idxs=kept_idxs,
        )

        return {
            "messages": [
                {
                    "role": "system",
                    "content": self.prompt,
                },
                {
                    "role": "user",
                    "content": user_prompt,
                },
                {
                    "role": "assistant",
                    "content": assistant_content,
                },
            ],
            "meta": {
                "source": self.hf_dataset,
                "split": split_name,
                "sample_id": sample_id,
                "kind": kind,
                "question": question,
                "answer": answer,
                "support_idxs": sorted(list(support_idxs)),
                "kept_idxs": sorted(list(kept_idxs)),
                "removed_support_idxs": sorted(list(removed_support_idxs)),
            },
        }

    # ------------------------------------------------------------
    # Prompt / response formatting
    # ------------------------------------------------------------

    def _make_user_prompt(
        self,
        question: str,
        paragraphs: list[dict[str, Any]],
        context_style: str = "clean",
        hide_titles: bool = False,
    ) -> str:
        documents_text = self._format_documents(
            paragraphs=paragraphs,
            style=context_style,
            hide_titles=hide_titles,
        )

        return self._clean_text(f"""
Documents:
{documents_text}

Question:
{question}

Instructions:
Use only the information inside the <document> tags.
If a required fact is missing from the documents, answer exactly: Not enough information.
Keep the reasoning concise and grounded in the documents.
""")

    def _make_assistant_response(
        self,
        answer: str,
        answerable: bool,
        subquestions: list[dict[str, Any]],
        paragraph_by_idx: dict[int, dict[str, Any]],
        kept_idxs: set[int],
    ) -> str:
        lines: list[str] = []
        missing_any = not answerable

        lines.append("I will answer using only the provided documents.")
        lines.append("")
        lines.append("Decomposition:")

        for i, step in enumerate(subquestions, start=1):
            step_question = self._scalar(step.get("question", ""))
            step_answer = self._scalar(step.get("answer", ""))
            support_idx = self._safe_int(step.get("paragraph_support_idx", None))

            lines.append(f"{i}. {step_question}")

            if support_idx is not None and support_idx in kept_idxs:
                paragraph = paragraph_by_idx.get(support_idx)
                title = (
                    paragraph.get("title", f"Document {support_idx}")
                    if paragraph is not None
                    else f"Document {support_idx}"
                )

                lines.append(
                    f'   From document id="{support_idx}" titled "{title}", '
                    f"the answer is {step_answer}."
                )
            else:
                missing_any = True
                lines.append(
                    "   The required evidence is not found in the provided documents."
                )

        if missing_any:
            final = "Not enough information."
        else:
            lines.append("")
            lines.append(
                f"From these evidence steps, the final answer is {answer}."
            )
            final = answer

        reasoning = "\n".join(lines)

        return self._clean_text(f"""
<think>
{reasoning}
</think>

Final Answer:
{final}
""")

    def _format_documents(
        self,
        paragraphs: list[dict[str, Any]],
        style: str = "clean",
        hide_titles: bool = False,
    ) -> str:
        parts: list[str] = []

        for p in paragraphs:
            idx = int(p["idx"])

            raw_title = self._scalar(p.get("title", f"Document {idx}"))
            title = f"Document {idx}" if hide_titles else raw_title

            text = self._scalar(p.get("paragraph_text", ""))

            if style == "clean":
                body = text

            elif style == "essay":
                body = f"Source excerpt: {text}"

            else:
                raise ValueError(f"Unknown context style: {style}")

            parts.append(
                f'<document id="{idx}" title="{self._escape_attr(title)}">\n'
                f'{body}\n'
                f'</document>'
            )

        return "\n\n".join(parts)

    # ------------------------------------------------------------
    # MuSiQue helpers
    # ------------------------------------------------------------

    def _support_indices(self, subquestions: list[dict[str, Any]]) -> set[int]:
        out: set[int] = set()

        for step in subquestions:
            idx = self._safe_int(step.get("paragraph_support_idx", None))
            if idx is not None:
                out.add(idx)

        return out

    def _paragraph_by_idx(
        self,
        paragraphs: list[dict[str, Any]],
    ) -> dict[int, dict[str, Any]]:
        out: dict[int, dict[str, Any]] = {}

        for p in paragraphs:
            idx = self._safe_int(p.get("idx", None))
            if idx is not None:
                out[idx] = p

        return out

    def _sample_support_plus_distractors(
        self,
        support: list[dict[str, Any]],
        distractors: list[dict[str, Any]],
        max_distractors: int,
    ) -> list[dict[str, Any]]:
        kept = support[:]

        if distractors:
            k = min(len(distractors), max_distractors)
            kept += self.rng.sample(distractors, k=k)

        return kept

    # ------------------------------------------------------------
    # Leak checks for missing-information samples
    # ------------------------------------------------------------

    def _missing_answer_leaks(
        self,
        subquestions: list[dict[str, Any]],
        kept_paragraphs: list[dict[str, Any]],
        kept_support_idxs: set[int],
        final_answer: str,
        answer_aliases: list[str],
    ) -> bool:
        """
        Prevent poisoned missing-info samples.

        If we remove a support document but the final answer or missing
        intermediate answer still appears in the remaining context, then
        "Not enough information" may be wrong.
        """

        context = self._context_for_leak_check(kept_paragraphs)

        final_candidates = [final_answer] + answer_aliases
        if self._any_candidate_leaks(final_candidates, context):
            return True

        for step in subquestions:
            support_idx = self._safe_int(step.get("paragraph_support_idx", None))

            if support_idx is not None and support_idx not in kept_support_idxs:
                step_answer = self._scalar(step.get("answer", ""))

                if self._any_candidate_leaks([step_answer], context):
                    return True

        return False

    def _context_for_leak_check(self, paragraphs: list[dict[str, Any]]) -> str:
        chunks: list[str] = []

        for p in paragraphs:
            chunks.append(self._scalar(p.get("title", "")))
            chunks.append(self._scalar(p.get("paragraph_text", "")))

        return self._normalize_for_match(" ".join(chunks))

    def _any_candidate_leaks(
        self,
        candidates: list[str],
        normalized_context: str,
    ) -> bool:
        for candidate in candidates:
            candidate = self._scalar(candidate)

            if not candidate:
                continue

            norm = self._normalize_for_match(candidate)

            # Skip tiny noisy strings.
            # Example: "a", "an", "UK", "3"
            # But keep 4-digit years like 1960.
            if len(norm) < 3:
                continue

            if norm.isdigit() and len(norm) < 4:
                continue

            if norm in normalized_context:
                return True

        return False

    def _normalize_for_match(self, text: str) -> str:
        text = str(text).lower()
        text = re.sub(r"\s+", " ", text)
        return text.strip()

    # ------------------------------------------------------------
    # Generic normalization helpers
    # ------------------------------------------------------------

    def _records_from_nested(self, value: Any) -> list[dict[str, Any]]:
        """
        Handles both common HF nested formats:

        1. list-of-dicts:
            [
                {"idx": 0, "title": "..."},
                {"idx": 1, "title": "..."}
            ]

        2. dict-of-lists:
            {
                "idx": [0, 1],
                "title": ["...", "..."]
            }
        """

        if value is None:
            return []

        # Sometimes scalar nested value is wrapped once.
        if isinstance(value, list) and len(value) == 1 and isinstance(value[0], list):
            value = value[0]

        # list-of-dicts
        if isinstance(value, list):
            return [x for x in value if isinstance(x, dict)]

        # dict-of-lists
        if isinstance(value, dict):
            keys = list(value.keys())
            if not keys:
                return []

            max_len = 0
            for k in keys:
                v = value[k]
                if isinstance(v, (list, tuple)):
                    max_len = max(max_len, len(v))
                else:
                    max_len = max(max_len, 1)

            records: list[dict[str, Any]] = []

            for i in range(max_len):
                rec = {}

                for k in keys:
                    v = value[k]

                    if isinstance(v, (list, tuple)):
                        rec[k] = v[i] if i < len(v) else None
                    else:
                        rec[k] = v

                records.append(rec)

            return records

        return []

    def _as_list(self, value: Any) -> list[Any]:
        if value is None:
            return []

        if isinstance(value, list):
            return value

        if isinstance(value, tuple):
            return list(value)

        return [value]

    def _scalar(self, value: Any) -> str:
        if value is None:
            return ""

        if isinstance(value, list) and len(value) == 1:
            return self._scalar(value[0])

        return str(value)

    def _safe_int(self, value: Any) -> Optional[int]:
        try:
            if value is None:
                return None

            if isinstance(value, list) and len(value) == 1:
                value = value[0]

            return int(value)
        except Exception:
            return None

    def _escape_attr(self, text: str) -> str:
        text = str(text)

        return (
            text.replace("&", "&amp;")
            .replace('"', "&quot;")
            .replace("<", "&lt;")
            .replace(">", "&gt;")
        )

    def _clean_text(self, text: str) -> str:
        text = text.strip()
        text = re.sub(r"\n{3,}", "\n\n", text)
        return text

    # ------------------------------------------------------------
    # Config / IO
    # ------------------------------------------------------------

    def _load_system_prompt(self) -> str:
        if ProjectConfigs is not None:
            try:
                prompt = ProjectConfigs.get(("PROMPTS", "LLM"))
                if isinstance(prompt, str) and prompt.strip():
                    return prompt.strip()
            except Exception:
                pass

        return DEFAULT_SYSTEM_PROMPT

    def _normalize_splits(self, splits: dict[str, float]) -> dict[str, float]:
        total = sum(splits.values())

        if total <= 0:
            raise ValueError("samples_split must sum to > 0")

        return {k: v / total for k, v in splits.items()}

    def _target_counts(
        self,
        target_samples: int,
        splits: dict[str, float],
    ) -> dict[str, int]:
        splits = self._normalize_splits(splits)

        counts = {
            kind: int(target_samples * ratio)
            for kind, ratio in splits.items()
        }

        diff = target_samples - sum(counts.values())

        if diff != 0:
            largest_kind = max(splits, key=splits.get)
            counts[largest_kind] += diff

        return counts

    def _write_jsonl(
        self,
        path: Path,
        records: list[dict[str, Any]],
    ) -> None:
        with path.open("w", encoding="utf-8") as f:
            for record in records:
                f.write(json.dumps(record, ensure_ascii=False) + "\n")