# dods/datagen/main.py
import sys
import inspect
from pathlib import Path
import yaml
from dods.datagen.helpers.engine import DataScenario


def run(config=None):
    """
    Universal entry point for datagen.

    Supports two main usage modes:
      1. CLI usage  ->  $ datagen path/to/config.yaml
      2. Python API ->  run("path/to/config.yaml") or run(CONFIG_DICT)

    The function automatically detects:
      - whether `config` is a dict or a YAML path
      - how to resolve relative paths correctly
      - whether to save the generated dataset automatically
    """

    # ---------------------------
    # STEP 1: Load the configuration
    # ---------------------------

    # Case 1: no config passed directly -> check CLI arguments
    if config is None and len(sys.argv) > 1:
        # e.g. `datagen contest.yaml`
        config = sys.argv[1]

    # Case 2: still no config → explicit error
    if config is None:
        raise ValueError("You must provide a config dict or YAML file path.")

    # Case 3: config is a path → load YAML file
    if isinstance(config, (str, Path)):
        config_path = Path(config)

        # --- Handle relative vs. absolute paths ---
        # If path is relative, resolve it differently depending on how datagen is executed.
        if not config_path.is_absolute():
            if len(sys.argv) > 0 and "datagen" in sys.argv[0].lower():
                # (CLI case) -> relative to current working directory
                config_path = Path.cwd() / config_path
            else:
                # (Python/Jupyter case) -> relative to the calling script
                caller_file = Path(inspect.stack()[1].filename)
                config_path = (caller_file.parent / config_path).resolve()

        if not config_path.exists():
            raise FileNotFoundError(f"Config file not found: {config_path}")

        # Load the YAML into a Python dictionary
        with open(config_path, "r", encoding="utf-8") as f:
            config = yaml.safe_load(f)

    # Case 4: invalid input type
    elif not isinstance(config, dict):
        raise TypeError("Config must be a dict or YAML file path.")

    # ---------------------------
    # STEP 2: Run the data generation
    # ---------------------------

    scenario = DataScenario(config)
    df = scenario.generate()

    # ---------------------------
    # STEP 3: Save output (if configured)
    # ---------------------------

    # The user can specify either:
    # - output_path: directory (e.g. "data/output")
    # - filename: file name (e.g. "test.csv")
    # or both together → will be combined
    output_dir = config.get("output_path") or config.get("output")
    filename = config.get("filename")

    if output_dir or filename:
        # Ensure target directory exists
        output_dir = Path(output_dir) if output_dir else Path.cwd()
        output_dir.mkdir(parents=True, exist_ok=True)

        # Build full output path
        if filename:
            output_path = output_dir / filename
        else:
            # Default filename if not provided
            scenario_name = config.get("scenario", "synthetic")
            output_path = output_dir / f"{scenario_name}_output.csv"

        # Save dataset
        df.to_csv(output_path, index=False)
        print(f"✅ Dataset successfully generated and saved to: {output_path}")
    else:
        print("ℹ️ Dataset generated (not saved because no output path or filename was provided).")

    # Always return the DataFrame (so you can use it in Python or Jupyter)
    return df


if __name__ == "__main__":
    """
    When executed directly (e.g. via CLI entry point),
    this ensures `run()` is called automatically.
    """
    run()
