"""Project scaffolding logic.

This module implements:
- `init_project`: create the baseline project structure
- `add_model`: create model template code under `models/<ModelName>/`
- `list_builtin_models`: list built-in model templates

Design goals:
- Safe by default (do not overwrite existing files unless `force=True`)
- Minimal dependencies (stdlib only)
- Predictable, cross-platform directory names (ASCII)
"""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Dict, List, Optional


# -------------------------
# Small filesystem helpers
# -------------------------


def _mkdir(path: Path) -> None:
    path.mkdir(parents=True, exist_ok=True)


def _write_text(path: Path, content: str, force: bool) -> bool:
    """Write `content` to `path`.

    Returns True if a write happened; False if skipped due to existing file.
    """
    if path.exists() and not force:
        return False
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(content, encoding="utf-8")
    return True


# -------------------------
# Model templates
# -------------------------


@dataclass(frozen=True)
class BuiltinModel:
    name: str
    description: str
    files: Dict[str, str]  # relative paths under models/<name>/


def _normalize_model_name(raw: str) -> str:
    """Normalize user input into canonical built-in model directory name.

    We accept a few common aliases to improve UX.
    """
    s = raw.strip()
    low = s.lower().replace("-", "").replace("_", "")

    if low in {"embeddingmlp", "mlp", "embedding"}:
        return "embeddingMLP"
    if low in {"xgboost", "xgb", "xgbboost", "xgbbooster", "xgbtree"}:
        return "xgBoost"
    if low in {"randomforest", "rf"}:
        return "randomForest"

    # For unknown models, we keep the input as-is (directory name).
    return s


def list_builtin_models() -> List[dict]:
    """Return built-in model list for CLI display."""
    return [{"name": m.name, "description": m.description} for m in _BUILTIN_MODELS.values()]


def init_project(base_dir: Path, force: bool) -> None:
    """Initialize the project scaffold in `base_dir`."""
    _mkdir(base_dir)

    # 公共工具层 (Global Utilities)
    _mkdir(base_dir / "global_utils" / "featureUtils")
    _mkdir(base_dir / "global_utils" / "modelUtils")

    _write_text(
        base_dir / "global_utils" / "featureUtils" / "featureTransform.py",
        _T_FEATURE_TRANSFORM,
        force,
    )
    _write_text(
        base_dir / "global_utils" / "featureUtils" / "featureStats.py",
        _T_FEATURE_STATS,
        force,
    )
    _write_text(
        base_dir / "global_utils" / "modelUtils" / "modelConverter.py",
        _T_MODEL_CONVERTER,
        force,
    )
    _write_text(
        base_dir / "global_utils" / "modelUtils" / "modelTrainer.py",
        _T_MODEL_TRAINER,
        force,
    )

    # 数据层
    _mkdir(base_dir / "data" / "raw")
    _mkdir(base_dir / "data" / "features")
    _write_text(base_dir / "data" / "raw" / "README.md", _T_DATA_RAW_README, force)
    _write_text(base_dir / "data" / "features" / "README.md", _T_DATA_FEATURES_README, force)

    # 特征工程层
    _mkdir(base_dir / "feature_engineering" / "feature_generation")
    _mkdir(base_dir / "feature_engineering" / "feature_config")
    _mkdir(base_dir / "feature_engineering" / "transform_scripts")
    _write_text(
        base_dir / "feature_engineering" / "feature_generation" / "README.md",
        _T_FE_GEN_README,
        force,
    )
    _write_text(
        base_dir / "feature_engineering" / "feature_config" / "README.md",
        _T_FE_CFG_README,
        force,
    )
    _write_text(
        base_dir / "feature_engineering" / "transform_scripts" / "README.md",
        _T_FE_TRANSFORM_README,
        force,
    )

    # 模型层
    _mkdir(base_dir / "models")

    # Create all three built-in models by default to match the requested structure.
    add_model(base_dir=base_dir, name="embeddingMLP", force=force)
    add_model(base_dir=base_dir, name="xgBoost", force=force)
    add_model(base_dir=base_dir, name="randomForest", force=force)

    # pipeline.py
    _write_text(base_dir / "pipeline.py", _T_PIPELINE, force)


def add_model(base_dir: Path, name: str, force: bool) -> Path:
    """Create model scaffold under `base_dir/models/<name>/`."""
    model_name = _normalize_model_name(name)
    model_dir = base_dir / "models" / model_name

    _mkdir(model_dir)
    _mkdir(model_dir / "save")
    _mkdir(model_dir / "save" / "save1")
    _write_text(model_dir / "save" / "README.md", _T_SAVE_README, force)

    builtin = _BUILTIN_MODELS.get(model_name)
    if builtin is None:
        # Unknown model: generate a generic minimal template.
        _write_text(model_dir / "model.py", _T_GENERIC_MODEL_PY.format(model_name=model_name), force)
        _write_text(model_dir / "config.py", _t_generic_config_py(model_name), force)
        _write_text(model_dir / "sampleDataUpload.py", _T_SAMPLE_UPLOAD_PY, force)
        return model_dir

    for rel_path, content in builtin.files.items():
        _write_text(model_dir / rel_path, content, force)

    return model_dir


def _t_generic_config_py(model_name: str) -> str:
    safe = "".join(ch if ch.isalnum() else "_" for ch in model_name)
    class_name = f"{safe[:1].upper()}{safe[1:]}Config"
    return (
        "from dataclasses import dataclass\n\n\n"
        f"@dataclass\nclass {class_name}:\n"
        "    # TODO: fill config fields\n"
        "    pass\n"
    )


# -------------------------
# File templates (utilities / readmes / pipeline)
# -------------------------


_T_FEATURE_TRANSFORM = """\
\"\"\"特征转换工具（占位）。

建议职责：
- 纯函数/可复用转换（不直接依赖业务 I/O）
- 输入输出协议稳定（DataFrame/ndarray/自定义样本结构）
\"\"\"

from __future__ import annotations


def example_transform(x: float) -> float:
    \"\"\"示例：对数变换（请按业务替换）。\"\"\"
    if x <= 0:
        return 0.0
    # 简化示例：避免额外依赖 numpy
    import math

    return math.log(x)
"""


_T_FEATURE_STATS = """\
\"\"\"特征统计工具（占位）。\"\"\"

from __future__ import annotations

from dataclasses import dataclass


@dataclass
class Stats:
    count: int
    mean: float


def mean(values: list[float]) -> Stats:
    if not values:
        return Stats(count=0, mean=0.0)
    return Stats(count=len(values), mean=sum(values) / len(values))
"""


_T_MODEL_CONVERTER = """\
\"\"\"模型转换工具（占位）。

建议职责：
- 统一保存/加载协议（如：pickle / joblib / ONNX）
- 版本标记、兼容性校验、元信息记录
\"\"\"

from __future__ import annotations

from pathlib import Path


def save_bytes(path: Path, payload: bytes) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_bytes(payload)


def load_bytes(path: Path) -> bytes:
    return path.read_bytes()
"""


_T_MODEL_TRAINER = """\
\"\"\"通用训练工具（占位）。

建议职责：
- 训练流程编排（数据读取 -> 特征 -> 训练 -> 评估 -> 保存）
- 统一日志/指标落盘协议
\"\"\"

from __future__ import annotations

from pathlib import Path
from typing import Any, Callable


def run_training(train_fn: Callable[[Path], Any], save_dir: Path) -> Any:
    \"\"\"统一的训练执行入口（示例）。\"\"\"
    save_dir.mkdir(parents=True, exist_ok=True)
    return train_fn(save_dir)
"""


_T_DATA_RAW_README = """\
这里放原始数据（raw）。
建议：不要把大文件提交到 git，使用 DVC/对象存储/数据平台。
"""


_T_DATA_FEATURES_README = """\
这里放特征数据（features）。
建议：明确特征生成版本与数据 schema。
"""


_T_FE_GEN_README = """\
特征生成层：建议放生成逻辑（例如按天批处理/增量生成）。
"""


_T_FE_CFG_README = """\
特征配置层：建议放特征开关、schema、特征列表、分桶配置等。
"""


_T_FE_TRANSFORM_README = """\
转换脚本：建议放一次性迁移脚本/数据修复脚本/特征回填脚本等。
"""


_T_SAVE_README = """\
这里存放多次训练的结果与图表。
建议结构：
- save1/ save2/ ... saveN/
每次训练固定落盘：metrics.json、模型文件、可视化图表等。
"""


_T_SAMPLE_UPLOAD_PY = """\
\"\"\"样本收集/上传占位脚本（通用）。

按你的业务实际改造：
- 从 data/raw 或线上拉取数据
- 做脱敏/采样
- 写入 data/features 或上传到对象存储
\"\"\"

from __future__ import annotations

from pathlib import Path


def main() -> None:
    project_root = Path(__file__).resolve().parents[2]
    raw_dir = project_root / "data" / "raw"
    print(f"[sampleDataUpload] raw data dir: {raw_dir}")
    print("TODO: implement your sample upload logic")


if __name__ == "__main__":
    main()
"""


_T_PIPELINE = """\
\"\"\"项目 pipeline 入口（占位）。

建议在这里做：
- 数据/特征/训练/评估的编排
- 统一参数解析与运行模式（本地/离线/线上）
\"\"\"

from __future__ import annotations

from pathlib import Path


def main() -> None:
    root = Path(__file__).resolve().parent
    print(f"Project root: {root}")
    print("TODO: implement pipeline orchestration")


if __name__ == "__main__":
    main()
"""


_T_GENERIC_MODEL_PY = """\
\"\"\"{model_name} 模型占位模板（通用）。

此文件为脚手架生成的最小可运行骨架：
- 不引入任何第三方依赖
- 通过写 metrics.json 的方式示范训练输出

后续你可以在这里接入 sklearn / xgboost / torch 等。
\"\"\"

from __future__ import annotations

from dataclasses import dataclass
import json
from pathlib import Path


@dataclass
class Config:
    # TODO: fill model hyperparameters
    pass


def train(cfg: Config, save_dir: Path) -> dict:
    save_dir.mkdir(parents=True, exist_ok=True)
    metrics = {{"status": "not_implemented", "model": "{model_name}"}}
    (save_dir / "metrics.json").write_text(json.dumps(metrics, indent=2), encoding="utf-8")
    return metrics


def main() -> None:
    here = Path(__file__).resolve().parent
    out = here / "save" / "save1"
    metrics = train(Config(), out)
    print(f"[{model_name}] placeholder done, metrics saved to: {{out}}")
    print(metrics)


if __name__ == "__main__":
    main()
"""


# -------------------------
# Built-in model templates
# -------------------------


_BUILTIN_MODELS: Dict[str, BuiltinModel] = {
    "embeddingMLP": BuiltinModel(
        name="embeddingMLP",
        description="最小可运行的 MLP 示例（占位版本，方便后续替换成真实 embedding/深度模型）",
        files={
            "model.py": """\
\"\"\"embeddingMLP 模型训练入口（最小可运行示例）。

目标：
- 给脚手架一个“能跑通”的示例，避免生成后用户无从下手
- 真实项目中你可以将此处替换为 PyTorch/TF/自研训练流程
\"\"\"

from __future__ import annotations

from dataclasses import dataclass
import json
from pathlib import Path
import random


@dataclass
class Config:
    # 这里是占位配置；你可以扩展为从 yaml/cli 读取
    seed: int = 42
    epochs: int = 3
    lr: float = 1e-3


def train(cfg: Config, save_dir: Path) -> dict:
    \"\"\"占位训练函数：模拟 loss 下降并写出 metrics.json。\"\"\"
    random.seed(cfg.seed)
    save_dir.mkdir(parents=True, exist_ok=True)

    loss = 1.0
    history = []
    for _ in range(cfg.epochs):
        loss *= (0.7 + random.random() * 0.1)
        history.append(loss)

    metrics = {"final_loss": loss, "history": history, "epochs": cfg.epochs, "lr": cfg.lr}
    (save_dir / "metrics.json").write_text(json.dumps(metrics, indent=2), encoding="utf-8")
    return metrics


def main() -> None:
    cfg = Config()
    here = Path(__file__).resolve().parent
    out = here / "save" / "save1"
    metrics = train(cfg, out)
    print(f"[embeddingMLP] done, metrics saved to: {out}")
    print(metrics)


if __name__ == "__main__":
    main()
""",
            "config.py": """\
\"\"\"embeddingMLP 配置（可按需替换为 pydantic/yaml/hydra 等）\"\"\"

from dataclasses import dataclass


@dataclass
class EmbeddingMLPConfig:
    seed: int = 42
    epochs: int = 3
    lr: float = 1e-3
""",
            "sampleDataUpload.py": _T_SAMPLE_UPLOAD_PY,
        },
    ),
    "xgBoost": BuiltinModel(
        name="xgBoost",
        description="XGBoost 占位模板（只提供接口骨架，不强依赖 xgboost 包）",
        files={
            "model.py": """\
\"\"\"xgBoost 模型占位模板。

说明：
- 为了让脚手架保持轻量，这里不强制依赖 xgboost
- 你可以在此处接入 xgboost.XGBClassifier / XGBRegressor
\"\"\"

from __future__ import annotations

from dataclasses import dataclass
import json
from pathlib import Path


@dataclass
class Config:
    # 占位参数：按需扩展
    n_estimators: int = 200
    max_depth: int = 6
    learning_rate: float = 0.1


def train(cfg: Config, save_dir: Path) -> dict:
    save_dir.mkdir(parents=True, exist_ok=True)
    metrics = {"status": "not_implemented", "cfg": cfg.__dict__}
    (save_dir / "metrics.json").write_text(json.dumps(metrics, indent=2), encoding="utf-8")
    return metrics


def main() -> None:
    here = Path(__file__).resolve().parent
    out = here / "save" / "save1"
    metrics = train(Config(), out)
    print(f"[xgBoost] placeholder done, metrics saved to: {out}")
    print(metrics)


if __name__ == "__main__":
    main()
""",
            "config.py": """\
from dataclasses import dataclass


@dataclass
class XGBoostConfig:
    n_estimators: int = 200
    max_depth: int = 6
    learning_rate: float = 0.1
""",
            "sampleDataUpload.py": _T_SAMPLE_UPLOAD_PY,
        },
    ),
    "randomForest": BuiltinModel(
        name="randomForest",
        description="RandomForest 占位模板（可对接 sklearn）",
        files={
            "model.py": """\
\"\"\"randomForest 模型占位模板。

说明：
- 你可以在此处接入 sklearn.ensemble.RandomForestClassifier/Regressor
- 脚手架默认不引入 sklearn 依赖，避免安装过重；按需在项目侧添加
\"\"\"

from __future__ import annotations

from dataclasses import dataclass
import json
from pathlib import Path
from typing import Optional


@dataclass
class Config:
    n_estimators: int = 300
    max_depth: Optional[int] = None
    random_state: int = 42


def train(cfg: Config, save_dir: Path) -> dict:
    save_dir.mkdir(parents=True, exist_ok=True)
    metrics = {"status": "not_implemented", "cfg": cfg.__dict__}
    (save_dir / "metrics.json").write_text(json.dumps(metrics, indent=2), encoding="utf-8")
    return metrics


def main() -> None:
    here = Path(__file__).resolve().parent
    out = here / "save" / "save1"
    metrics = train(Config(), out)
    print(f"[randomForest] placeholder done, metrics saved to: {out}")
    print(metrics)


if __name__ == "__main__":
    main()
""",
            "config.py": """\
from dataclasses import dataclass
from typing import Optional


@dataclass
class RandomForestConfig:
    n_estimators: int = 300
    max_depth: Optional[int] = None
    random_state: int = 42
""",
            "sampleDataUpload.py": _T_SAMPLE_UPLOAD_PY,
        },
    ),
}

