import json
import logging
from datetime import timedelta
from pathlib import Path
from typing import Any

import yaml

# Type aliases for better readability
StatisticsDict = dict[str, dict[str, int | float]]
"""Type alias for statistics dictionary containing nested metrics.
Example: {"runtime": {"mean": 0.5, "std": 0.1}, "complexity": {"max": 10, "min": 1}}"""

StringProblemOrAnswer = str | list[str] | list[list[str]] | list[str | list[str]]
"""Type alias for string-formatted problems and answers.
Supports single strings, simple lists, nested lists, and mixed structures.
Maximum nesting depth is 2 levels.

Polynomial examples:
    - Single: "x^2 + 2*x + 1"
    - Simple list: ["x + y", "x - y"]
    - Nested list: [["x^2", "y^2"], ["z^2", "w^2"]]
    - Mixed: [["x", "y"], "z"] 
    
Arithmetic examples:
    - Single: "2"
    - Simple list: ["2", "3", "5", "7"]
    - Nested list: [["2", "3"], ["5", "7"]]
    - Mixed: ["2", ["3", "5"]]
"""

StringSample = tuple[StringProblemOrAnswer, StringProblemOrAnswer]
"""Type alias for a single problem-answer pair in string format.
Example: ("x^2 + 2*x + 1", "(x + 1)^2") or (["x + y", "x - y"], ["2*x", "2*y"])"""

StringSampleList = list[StringSample]
"""Type alias for a list of problem-answer pairs in string format.
Example: [("problem1", "answer1"), ("problem2", "answer2")]"""


class TimedeltaDumper(yaml.SafeDumper):
    """Custom YAML dumper that safely handles timedelta objects."""

    pass


def timedelta_representer(dumper: TimedeltaDumper, data: timedelta) -> yaml.ScalarNode:
    """Convert timedelta to float seconds for YAML serialization.

    Args:
        dumper: YAML SafeDumper instance.
        data: timedelta to serialize.

    Returns:
        YAML ScalarNode representing total seconds as float.
    """
    return dumper.represent_float(data.total_seconds())


class DatasetWriter:
    """
    Dataset writer for saving problem-answer pairs in multiple formats.

    This class handles saving datasets with nested structure support up to 2 levels.
    It can save data in raw text (.txt) and JSON Lines (.jsonl) formats.

    Attributes:
        INNER_SEP (str): Separator for single-level lists (" | ")
        OUTER_SEP (str): Separator for nested lists (" || ")
        save_dir (Path): Base directory for saving datasets
        save_text (bool): Whether to save raw text files
        save_json (bool): Whether to save JSON Lines files
        _file_handles (dict): Dictionary to store open file handles
    """

    # Separator constants for string formatting
    INNER_SEP = " | "  # Separator for single-level lists
    OUTER_SEP = " || "  # Separator for nested lists

    def __init__(
        self,
        save_dir: str | None = None,
        save_text: bool = True,
        save_json: bool = True,
    ) -> None:
        """
        Initialize dataset writer.

        Args:
            save_dir: Base directory for saving datasets. If None, uses current working directory.
            save_text: Whether to save raw text files. Text files use "#" as separator
                      between problem and answer, with nested structures joined by separators.
            save_json: Whether to save JSON Lines files. JSON Lines files preserve the original
                      nested structure format, with one sample per line.

        Usage:
            ```python
            # Efficient batch processing with file handle management
            writer = DatasetWriter(save_dir="./datasets")
            writer.open("train")  # Open file handles once
            try:
                for batch_idx, samples in enumerate(batches):
                    writer.save_batch(samples, tag="train", batch_idx=batch_idx)
            finally:
                writer.close("train")  # Close file handles

            # Or use context manager
            with DatasetWriter(save_dir="./datasets") as writer:
                writer.open("train")
                for batch_idx, samples in enumerate(batches):
                    writer.save_batch(samples, tag="train", batch_idx=batch_idx)
                writer.close("train")

            # Support for various dataset splits
            writer.open("validation")  # Validation set
            writer.open("dev")         # Development set
            writer.open("eval")        # Evaluation set
            ```
        """
        self.save_dir = Path(save_dir) if save_dir else Path.cwd()
        self.save_text = save_text
        self.save_json = save_json
        self.logger = logging.getLogger(__name__)
        self._file_handles: dict[
            str, dict[str, Any]
        ] = {}  # {tag: {file_type: file_handle}}
        TimedeltaDumper.add_representer(timedelta, timedelta_representer)

    def _validate_tag(self, tag: str) -> None:
        """
        Validate tag parameter.

        Args:
            tag: Tag to validate

        Raises:
            ValueError: If tag is empty or None
        """
        if not tag or not isinstance(tag, str):
            raise ValueError(f"tag must be a non-empty string, got '{tag}'")

    def _create_dataset_dir(self) -> Path:
        """
        Create and return dataset directory path.

        Returns:
            Path object pointing to the created dataset directory
        """
        self.save_dir.mkdir(parents=True, exist_ok=True)
        return self.save_dir

    def _validate_nested_structure(self, obj: StringProblemOrAnswer) -> None:
        """
        Validate that the object has at most 2 levels of nesting and contains only strings.

        Args:
            obj: Object to validate

        Raises:
            ValueError: If obj contains lists deeper than 2 levels or non-string elements
        """
        if isinstance(obj, str):
            # Single string is valid
            return None
        elif isinstance(obj, list):
            # Check if this list contains any nested lists
            has_nested_lists = any(isinstance(item, list) for item in obj)

            if has_nested_lists:
                # Check that inner lists don't contain further nested lists
                for item in obj:
                    if isinstance(item, list):
                        # First check for deeper nesting (3+ levels)
                        if any(isinstance(subitem, list) for subitem in item):
                            raise ValueError(
                                f"Lists deeper than 2 levels are not supported. "
                                f"Found 3+ level structure in: {obj}"
                            )
                        # Then validate inner list contains only strings
                        if not all(isinstance(subitem, str) for subitem in item):
                            raise ValueError(
                                f"All elements in nested lists must be strings. "
                                f"Found non-string element in: {obj}"
                            )
                    else:
                        # Validate single item is string
                        if not isinstance(item, str):
                            raise ValueError(
                                f"All elements must be strings. "
                                f"Found non-string element: {item} in {obj}"
                            )
            else:
                # Simple list: validate all elements are strings
                if not all(isinstance(item, str) for item in obj):
                    raise ValueError(
                        f"All elements in simple lists must be strings. "
                        f"Found non-string element in: {obj}"
                    )
        else:
            # Neither string nor list
            raise ValueError(
                f"Object must be a string or list. "
                f"Found type {type(obj).__name__}: {obj}"
            )

    def _join_with_separators(self, obj: StringProblemOrAnswer) -> str:
        """
        Join object elements with appropriate separators.

        The following rules are applied to stringify obj:
        - Single string -> return as is
        - Single-level list -> join with INNER_SEP (" | ")
        - List containing nested lists -> join each nested list with INNER_SEP (" | "),
          then join all parts with OUTER_SEP (" || ")

        Args:
            obj: Problem or answer in string format (str or list with at most 2 levels).

        Returns:
            String with elements joined by INNER_SEP and/or OUTER_SEP.

        Raises:
            ValueError: If obj contains lists deeper than 2 levels
        """
        # Validate structure first
        self._validate_nested_structure(obj)

        if isinstance(obj, list):
            # Check if this list contains any nested lists
            has_nested_lists = any(isinstance(item, list) for item in obj)

            if has_nested_lists:
                # Mixed structure: join with OUTER_SEP
                parts = []
                for item in obj:
                    if isinstance(item, list):
                        # Inner list: join with INNER_SEP
                        parts.append(self.INNER_SEP.join(item))
                    else:
                        # Single item: already a string
                        parts.append(item)
                return self.OUTER_SEP.join(parts)
            else:
                # Simple list: join with INNER_SEP
                return self.INNER_SEP.join(obj)
        else:
            # Single value: already a string
            return obj

    def _format_sample_strings(
        self,
        problem_str: StringProblemOrAnswer,
        answer_str: StringProblemOrAnswer,
    ) -> str:
        """
        Format problem and answer to string representation with separators.

        Args:
            problem_str: Problem in string format (can be nested)
            answer_str: Answer in string format (can be nested)

        Returns:
            Formatted string with problem and answer separated by " # " (problem # answer)

        Examples:
            >>> writer._format_sample_strings("x", "y")
            "x # y"
            >>> writer._format_sample_strings(["a", "b"], ["c", "d"])
            "a | b # c | d"
            >>> writer._format_sample_strings([["a", "b"], "c"], [["d", "e"], "f"])
            "a | b || c # d | e || f"
        """
        problem_formatted = self._join_with_separators(problem_str)
        answer_formatted = self._join_with_separators(answer_str)
        return f"{problem_formatted} # {answer_formatted}"

    def _get_json_data(
        self,
        problem_str: StringProblemOrAnswer,
        answer_str: StringProblemOrAnswer,
    ) -> dict[str, StringProblemOrAnswer]:
        """
        Format problem and answer to JSON format.

        This method creates a dictionary structure suitable for JSON serialization.
        The original nested structure is preserved exactly as provided.

        Args:
            problem_str: Problem in string format (can be nested)
            answer_str: Answer in string format (can be nested)

        Returns:
            Dictionary with "problem" and "answer" keys preserving original structure

        Examples:
            >>> writer._get_json_data("x^2 + 1", "x^2")
            {"problem": "x^2 + 1", "answer": "x^2"}
            >>> writer._get_json_data(["x + y", "x - y"], ["2*x", "2*y"])
            {"problem": ["x + y", "x - y"], "answer": ["2*x", "2*y"]}
            >>> writer._get_json_data([["x", "y"], ["z"]], [["a", "b"], ["c"]])
            {"problem": [["x", "y"], ["z"]], "answer": [["a", "b"], ["c"]]}
        """
        return {"problem": problem_str, "answer": answer_str}

    def open(self, tag: str) -> None:
        """
        Open file handles for the specified tag.

        This method should be called before starting batch processing to avoid
        repeated file open/close operations.

        Args:
            tag: Dataset tag (e.g., "train", "test", "validation", "dev", "eval")

        Raises:
            ValueError: If tag is invalid
        """
        self._validate_tag(tag)

        if tag in self._file_handles:
            self.logger.warning(f"File handles for tag '{tag}' are already open")
            return

        dataset_dir = self._create_dataset_dir()
        self._file_handles[tag] = {}

        # Open text file if enabled
        if self.save_text:
            raw_path = dataset_dir / f"{tag}_raw.txt"
            self._file_handles[tag]["text"] = open(raw_path, "w")

        # Open JSON Lines file if enabled
        if self.save_json:
            json_path = dataset_dir / f"{tag}_data.jsonl"
            self._file_handles[tag]["json"] = open(json_path, "w")

    def close(self, tag: str) -> None:
        """
        Close file handles for the specified tag.

        This method should be called after finishing batch processing.

        Args:
            tag: Dataset tag (e.g., "train", "test", "validation", "dev", "eval")

        Raises:
            ValueError: If tag is invalid
        """
        self._validate_tag(tag)

        if tag not in self._file_handles:
            self.logger.warning(f"No open file handles found for tag '{tag}'")
            return

        # Close all open file handles
        for file_type, file_handle in self._file_handles[tag].items():
            if hasattr(file_handle, "close"):  # Only close actual file handles
                try:
                    file_handle.close()
                except Exception as e:
                    self.logger.error(
                        f"Error closing {file_type} file for tag '{tag}': {e}"
                    )

        del self._file_handles[tag]

    def close_all(self) -> None:
        """
        Close all open file handles.

        Closes every tag that was opened with open(tag). Should be called when
        the writer is no longer needed (e.g., before exiting or when using as context manager).
        """
        for tag in list(self._file_handles.keys()):
            self.close(tag)

    def __enter__(self):
        """Context manager entry."""
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit - close all files."""
        self.close_all()

    def save_batch(
        self,
        samples: StringSampleList,
        tag: str = "train",
        batch_idx: int = 0,
    ) -> None:
        """
        Save a batch of samples to files.

        This method saves samples in:
        - Raw text (.txt) - Human-readable format with separators (if save_text=True)
        - JSON Lines (.jsonl) - Structured format preserving nested structure (if save_json=True)

        Args:
            samples: List of (problem, answer) pairs in string format
            tag: Dataset tag (e.g., "train", "test", "validation", "dev", "eval")
            batch_idx: Batch index for incremental saving. Use 0 for first batch,
                      subsequent batches will append to existing files.

        Raises:
            ValueError: If tag is invalid or samples contain invalid nested structures

        Examples:
            >>> # Simple string samples (single problem-answer pairs)
            >>> writer = DatasetWriter(save_dir="./datasets", save_text=True, save_json=True)
            >>> samples = [
            ...     ("x^2 + 2*x + 1", "(x + 1)^2"),
            ...     ("2*x^3 - 3*x^2", "x^2*(2*x - 3)"),
            ... ]
            >>> # Creates: train_raw.txt, train_data.jsonl
            >>> writer.save_batch(samples, tag="train", batch_idx=0)
            >>>
            >>> # 1 level nested structure samples (multiple problems/answers)
            >>> samples = [
            ...     (["x + y", "x - y"], ["2*x", "2*y"]),
            ...     (["x^2 + y^2", "x^2 - y^2"], ["2*x^2", "2*y^2"]),
            ... ]
            >>> # Text output: "x + y | x - y # 2*x | 2*y"
            >>> writer.save_batch(samples, tag="test", batch_idx=0)
            >>>
            >>> # 2 level nested structure samples (complex nested problems)
            >>> samples = [
            ...     ([["x", "y"], ["z", "w"]], [["x", "z"], ["y", "w"]]),
            ...     ([["x + y", "x - y"], ["z + w", "z - w"]], [["x + y", "z + w"], ["x - y", "z - w"]]),
            ... ]
            >>> # Text output: "x | y || z | w # x | z || y | w"
            >>> writer.save_batch(samples, tag="test", batch_idx=0)
            >>>
            >>> # Append more samples to existing dataset
            >>> more_samples = [
            ...     ([["a", "b"], ["c", "d"]], [["a", "c"], ["b", "d"]]),
            ...     ([["e", "f"], ["g", "h"]], [["e", "g"], ["f", "h"]]),
            ... ]
            >>> # Appends to existing files instead of overwriting
            >>> writer.save_batch(more_samples, tag="train", batch_idx=1)
        """
        self._validate_tag(tag)

        # Validate samples
        if not samples:
            self.logger.warning(
                "Empty samples list provided - no files will be created"
            )
            return

        # Check if file handles are open for this tag
        if tag not in self._file_handles:
            # Fallback to old method if file handles are not open
            self._save_batch_legacy(samples, tag, batch_idx)
            return

        # Save raw text data (optional)
        if self.save_text:
            for problem_str, answer_str in samples:
                formatted_line = self._format_sample_strings(problem_str, answer_str)
                self._file_handles[tag]["text"].write(f"{formatted_line}\n")
            self._file_handles[tag]["text"].flush()

        # Save JSON Lines data (optional)
        if self.save_json:
            for problem_str, answer_str in samples:
                json_data = self._get_json_data(problem_str, answer_str)
                json_line = json.dumps(json_data, ensure_ascii=False)
                self._file_handles[tag]["json"].write(f"{json_line}\n")
            self._file_handles[tag]["json"].flush()

    def _save_batch_legacy(
        self,
        samples: StringSampleList,
        tag: str = "train",
        batch_idx: int = 0,
    ) -> None:
        """
        Legacy save_batch method for backward compatibility.

        Used when file handles are not open (open(tag) was not called), maintaining
        the original behavior of opening/closing files for each batch.

        Args:
            samples: List of (problem, answer) pairs in string format.
            tag: Dataset tag (e.g., "train", "test", "validation").
            batch_idx: Batch index; 0 overwrites/writes new files, >0 appends.
        """
        dataset_dir = self._create_dataset_dir()

        # Save raw text data (optional)
        if self.save_text:
            raw_path = dataset_dir / f"{tag}_raw.txt"
            mode = "w" if batch_idx == 0 else "a"
            with open(raw_path, mode) as f:
                for problem_str, answer_str in samples:
                    formatted_line = self._format_sample_strings(
                        problem_str, answer_str
                    )
                    f.write(f"{formatted_line}\n")

        # Save JSON Lines data (optional)
        if self.save_json:
            json_path = dataset_dir / f"{tag}_data.jsonl"
            mode = "w" if batch_idx == 0 else "a"
            with open(json_path, mode) as f:
                for problem_str, answer_str in samples:
                    json_data = self._get_json_data(problem_str, answer_str)
                    json_line = json.dumps(json_data, ensure_ascii=False)
                    f.write(f"{json_line}\n")

    def save_final_statistics(
        self,
        statistics: StatisticsDict,
        tag: str = "train",
    ) -> None:
        """
        Save final overall statistics to YAML file.

        Args:
            statistics: Dictionary containing dataset statistics
            tag: Dataset tag (e.g., "train", "test", "validation", "dev", "eval")

        Raises:
            ValueError: If tag is invalid

        Note:
            Statistics are saved in YAML format for human readability.
            The file is named "{tag}_stats.yaml" in the dataset directory.
        """
        self._validate_tag(tag)
        dataset_dir = self._create_dataset_dir()

        stats_path = dataset_dir / f"{tag}_stats.yaml"
        with open(stats_path, "w") as f:
            yaml.dump(
                statistics,
                f,
                Dumper=TimedeltaDumper,
                default_flow_style=False,
                sort_keys=False,
                indent=4,
            )

    def load_dataset_jsonl(self, tag: str) -> StringSampleList:
        """
        Load dataset from JSON Lines file.

        Args:
            tag: Dataset tag (e.g., "train", "test", "validation", "dev", "eval")

        Returns:
            List of (problem, answer) pairs in string format. Accepts JSON lines
            with either "answer" or "solution" key (for backward compatibility).

        Raises:
            ValueError: If tag is invalid
            FileNotFoundError: If the JSON Lines file doesn't exist

        Examples:
            >>> samples = writer.load_dataset_jsonl("train")
            >>> print(f"Loaded {len(samples)} samples")
            >>> for problem, answer in samples[:3]:
            ...     print(f"Problem: {problem}, Answer: {answer}")
        """
        self._validate_tag(tag)
        jsonl_path = self.save_dir / f"{tag}_data.jsonl"

        if not jsonl_path.exists():
            raise FileNotFoundError(f"JSON Lines file not found: {jsonl_path}")

        samples = []
        with open(jsonl_path, "r") as f:
            for line_num, line in enumerate(f, 1):
                line = line.strip()
                if not line:  # Skip empty lines
                    continue
                try:
                    data = json.loads(line)
                    problem = data.get("problem")
                    answer = data.get("answer") or data.get("solution")
                    if problem is None or answer is None:
                        self.logger.warning(
                            f"Line {line_num}: missing 'problem' or 'answer'/'solution' key"
                        )
                        continue
                    samples.append((problem, answer))
                except (json.JSONDecodeError, KeyError) as e:
                    self.logger.warning(f"Error parsing line {line_num}: {e}")
                    continue

        return samples
