"""Notebook generation for exporting harnessml projects to Colab, Kaggle, or local."""
from __future__ import annotations

from pathlib import Path
from typing import Optional

import nbformat
import yaml

_VALID_DESTINATIONS = ("colab", "kaggle", "local")


def generate_notebook(
    project_dir: str | Path,
    *,
    destination: str,
    output_path: Optional[str | Path] = None,
) -> Path:
    """Generate an .ipynb notebook for the given project and destination platform.

    Parameters
    ----------
    project_dir : str | Path
        Root directory of the harnessml project (must contain config/config.yaml).
    destination : str
        Target platform: "colab", "kaggle", or "local".
    output_path : str | Path | None
        Where to write the notebook. If None, defaults to
        ``{project_dir}/{project_name}_{destination}.ipynb``.

    Returns
    -------
    Path
        Path to the generated notebook file.

    Raises
    ------
    ValueError
        If destination is not one of the valid options.
    """
    if destination not in _VALID_DESTINATIONS:
        raise ValueError(
            f"Invalid destination {destination!r}; must be one of {_VALID_DESTINATIONS}"
        )

    project_dir = Path(project_dir)
    config_path = project_dir / "config" / "config.yaml"
    config_text = config_path.read_text()
    config = yaml.safe_load(config_text)

    project_name = config.get("project", {}).get("name", "harnessml-project")

    nb = nbformat.v4.new_notebook()
    nb.cells = _build_cells(project_name, config_text, destination)

    if output_path is None:
        out = project_dir / f"{project_name}_{destination}.ipynb"
    else:
        out = Path(output_path)

    out.parent.mkdir(parents=True, exist_ok=True)
    nbformat.write(nb, str(out))
    return out


def _build_cells(project_name: str, config_text: str, destination: str) -> list:
    """Build the ordered list of notebook cells."""
    cells: list = []

    # 1 - Title
    cells.append(nbformat.v4.new_markdown_cell(
        f"# {project_name}\n\nAuto-generated by harnessml for {destination}"
    ))

    # 2 - Install
    if destination == "colab":
        install_src = "!pip install harness-core"
    elif destination == "kaggle":
        install_src = "!pip install -q harness-core"
    else:
        install_src = "# Local environment — install harness-core manually if needed\n# pip install harness-core"
    cells.append(nbformat.v4.new_code_cell(install_src))

    # 3 - Setup header
    cells.append(nbformat.v4.new_markdown_cell("## Setup"))

    # 4 - Platform setup
    cells.append(nbformat.v4.new_code_cell(_platform_setup(project_name, destination)))

    # 5 - Configuration header
    cells.append(nbformat.v4.new_markdown_cell("## Configuration"))

    # 6 - Inline config
    cells.append(nbformat.v4.new_code_cell(
        "import yaml\n\n"
        f'CONFIG_YAML = """\n{config_text}"""\n\n'
        "config = yaml.safe_load(CONFIG_YAML)"
    ))

    # 7 - Load Data header
    cells.append(nbformat.v4.new_markdown_cell("## Load Data"))

    # 8 - Data loading
    cells.append(nbformat.v4.new_code_cell(
        "import pandas as pd\n\n"
        "features_path = DATA_DIR / \"features.parquet\"\n"
        "df = pd.read_parquet(features_path)\n"
        "print(f\"Loaded {len(df)} rows, {len(df.columns)} columns\")\n"
        "df.head()"
    ))

    # 9 - Train header
    cells.append(nbformat.v4.new_markdown_cell("## Train & Evaluate"))

    # 10 - Training
    cells.append(nbformat.v4.new_code_cell(
        "import shutil\n\n"
        "# Write config to output dir\n"
        "config_out = OUTPUT_DIR / \"config\"\n"
        "config_out.mkdir(parents=True, exist_ok=True)\n"
        "(config_out / \"config.yaml\").write_text(CONFIG_YAML)\n\n"
        "# Copy data\n"
        "data_out = OUTPUT_DIR / \"data\"\n"
        "data_out.mkdir(parents=True, exist_ok=True)\n"
        "shutil.copy2(features_path, data_out / \"features.parquet\")\n\n"
        "# Run pipeline\n"
        "from harnessml.core.runner.pipeline import PipelineRunner\n\n"
        "runner = PipelineRunner(str(OUTPUT_DIR))\n"
        "result = runner.run()"
    ))

    # 11 - Results header
    cells.append(nbformat.v4.new_markdown_cell("## Results"))

    # 12 - Results display
    cells.append(nbformat.v4.new_code_cell(
        "# Display metrics\n"
        "if hasattr(result, 'metrics'):\n"
        "    for k, v in result.metrics.items():\n"
        "        print(f\"{k}: {v}\")\n"
        "else:\n"
        "    print(result)"
    ))

    return cells


def _platform_setup(project_name: str, destination: str) -> str:
    """Return the platform-specific setup code."""
    if destination == "colab":
        return (
            "from google.colab import drive\n"
            "drive.mount('/content/drive')\n\n"
            "from pathlib import Path\n\n"
            f"PROJECT_DIR = Path('/content/drive/MyDrive/{project_name}')\n"
            "DATA_DIR = PROJECT_DIR / 'data'\n"
            "OUTPUT_DIR = PROJECT_DIR / 'output'"
        )
    elif destination == "kaggle":
        return (
            "from pathlib import Path\n\n"
            f"PROJECT_DIR = Path('/kaggle/input/{project_name}')\n"
            "DATA_DIR = PROJECT_DIR / 'data'\n"
            "OUTPUT_DIR = Path('/kaggle/working/')"
        )
    else:
        return (
            "from pathlib import Path\n\n"
            "PROJECT_DIR = Path.cwd()\n"
            "DATA_DIR = PROJECT_DIR / 'data'\n"
            "OUTPUT_DIR = PROJECT_DIR / 'output'"
        )
