"""
Orchard ML Command-Line Interface.

Provides the ``orchard`` entry point with two commands:

- ``orchard init`` — generate a starter recipe YAML with all defaults
- ``orchard run``  — execute a training pipeline from a YAML recipe

Usage::

    orchard init
    orchard run recipes/config_mini_cnn.yaml
    orchard run recipes/optuna_mini_cnn.yaml --set training.epochs=30
"""

from __future__ import annotations

from pathlib import Path
from typing import TYPE_CHECKING, Annotated, Any

import typer

from .exceptions import OrchardError

if TYPE_CHECKING:  # pragma: no cover
    from pydantic import BaseModel

# ── App instance ───────────────────────────────────────────────────────────

app = typer.Typer(
    name="orchard",
    add_completion=True,
    rich_markup_mode="rich",
    no_args_is_help=True,
)

# ── Constants ──────────────────────────────────────────────────────────────

_INIT_HEADER = """\
# yaml-language-server: $schema=
# ==============================================================================
# Orchard ML — Starter Recipe (generated by `orchard init`)
# ==============================================================================
# Usage:   orchard run {filename}
# Docs:    https://github.com/tomrussobuilds/orchard-ml
#
# Edit the values you need. Remove optional sections to disable them.
# Optional sections: tracking, export, optuna
# ==============================================================================

"""

_COMMENT_MAX_WIDTH = 80

_YAML_INDENT = "    "
_YAML_KEYWORDS = frozenset({"true", "false", "null", "yes", "no", "on", "off", ""})


# ── Version callback ──────────────────────────────────────────────────────


def _version_callback(value: bool) -> None:
    """Print version and exit."""
    if value:
        from importlib.metadata import version as pkg_version

        typer.echo(f"orchard-ml {pkg_version('orchard-ml')}")
        raise typer.Exit()


# ── Commands ───────────────────────────────────────────────────────────────


@app.callback()
def main(
    _: Annotated[
        bool | None,
        typer.Option(
            "--version",
            "-V",
            callback=_version_callback,
            is_eager=True,
            help="Show version and exit.",
        ),
    ] = None,
) -> None:
    """Orchard ML: type-Safe Deep Learning for Reproducible Research."""
    ...  # pragma: no cover


@app.command()
def init(
    output: Annotated[
        Path,
        typer.Argument(help="Output YAML file path."),
    ] = Path("recipe.yaml"),
    force: Annotated[
        bool,
        typer.Option("--force", "-f", help="Overwrite existing file."),
    ] = False,
) -> None:
    """Generate a starter recipe with all config fields and defaults."""
    if output.exists() and not force:
        typer.echo(f"Error: '{output}' already exists. Use --force to overwrite.", err=True)
        raise typer.Exit(code=1)

    data = _build_init_dict()
    yaml_body = _build_commented_yaml(data)
    content = _INIT_HEADER.format(filename=output.name) + yaml_body

    output.write_text(content, encoding="utf-8")
    typer.echo(f"Recipe created: {output}")
    typer.echo(f"Run it with:   orchard run {output}")


@app.command()
def run(
    recipe: Annotated[
        Path,
        typer.Argument(help="Path to YAML recipe file."),
    ],
    set_: Annotated[
        list[str] | None,
        typer.Option(
            "--set",
            help="Override config value (repeatable): key.path=value",
        ),
    ] = None,
) -> None:
    """Run the ML pipeline from a YAML recipe."""
    from orchard import (
        MLRUNS_DB,
        Config,
        LogStyle,
        RootOrchestrator,
        create_tracker,
        log_pipeline_summary,
        run_export_phase,
        run_optimization_phase,
        run_training_phase,
    )

    if not recipe.exists():
        typer.echo(f"Error: recipe not found: {recipe}", err=True)
        raise typer.Exit(code=1)

    overrides = _parse_overrides(set_ or [])
    cfg = Config.from_recipe(recipe, overrides=overrides or None)

    with RootOrchestrator(cfg) as orchestrator:
        paths = orchestrator.paths
        run_logger = orchestrator.run_logger
        if paths is None or run_logger is None:
            raise RuntimeError("RootOrchestrator failed to initialize paths or logger")

        tracker = create_tracker(cfg)
        tracking_uri = f"sqlite:///{MLRUNS_DB}"
        tracker.start_run(cfg=cfg, run_name=paths.run_id, tracking_uri=tracking_uri)

        orchestrator.log_environment_report()

        training_cfg = cfg

        try:
            # Phase 1: Optimization (if optuna config present)
            if cfg.optuna is not None:
                _, best_config_path = run_optimization_phase(orchestrator, tracker=tracker)
                if best_config_path and best_config_path.exists():
                    training_cfg = Config.from_recipe(best_config_path)
                    run_logger.info(
                        "%s%s %-18s: %s",
                        LogStyle.INDENT,
                        LogStyle.ARROW,
                        "Optimized Config",
                        best_config_path.name,
                    )
            else:
                run_logger.info(
                    "%s%s Skipping optimization (no optuna config)",
                    LogStyle.INDENT,
                    LogStyle.ARROW,
                )

            # Phase 2: Training
            result = run_training_phase(orchestrator, cfg=training_cfg, tracker=tracker)

            # Phase 3: Export (if export config present)
            onnx_path = None
            if cfg.export is not None:
                onnx_path = run_export_phase(
                    orchestrator,
                    checkpoint_path=result.best_model_path,
                    cfg=training_cfg,
                )

            # Log final artifacts
            tracker.log_artifacts_dir(paths.figures)
            tracker.log_artifact(paths.final_report_path)
            tracker.log_artifact(paths.get_config_path())

            log_pipeline_summary(
                test_metrics=result.test_metrics,
                best_model_path=result.best_model_path,
                run_dir=paths.root,
                duration=orchestrator.time_tracker.elapsed_formatted,
                onnx_path=onnx_path,
                logger_instance=run_logger,
            )

        except KeyboardInterrupt:
            run_logger.warning("%s Interrupted by user.", LogStyle.WARNING)
            raise SystemExit(1)

        except OrchardError as e:
            run_logger.error("%s %s", LogStyle.FAILURE, e)
            raise SystemExit(1)

        except Exception as e:  # top-level catch-all for logging; re-raises
            run_logger.error("%s Pipeline failed: %s", LogStyle.FAILURE, e, exc_info=True)
            raise

        finally:
            tracker.end_run()


# ── CLI override helpers ───────────────────────────────────────────────────


def _auto_cast(value: str) -> Any:
    """
    Cast a CLI string to the appropriate Python scalar type.

    Args:
        value: Raw string from the command line.

    Returns:
        Converted bool, None, int, float, or the original string.
    """
    low = value.lower()
    if low in ("true", "false"):
        return low == "true"
    if low in ("null", "none"):
        return None
    try:
        return int(value)
    except ValueError:
        pass
    try:
        return float(value)
    except ValueError:
        pass
    return value


def _parse_overrides(raw: list[str]) -> dict[str, Any]:
    """
    Parse ``key.path=value`` strings into a flat override dict.

    Args:
        raw: list of "dotted.key=value" strings from ``--set`` flags.

    Returns:
        dict mapping dotted keys to auto-casted values.

    Raises:
        typer.BadParameter: If an item has no ``=``, an empty key,
            or the section/field does not exist in the config schema.
    """
    overrides: dict[str, Any] = {}
    for item in raw:
        if "=" not in item:
            raise typer.BadParameter(f"Override must use key=value format, got: '{item}'")
        key, _, val = item.partition("=")
        key = key.strip()
        if not key:
            raise typer.BadParameter(f"Empty key in override: '{item}'")
        _validate_override_key(key)
        overrides[key] = _auto_cast(val.strip())
    return overrides


def _validate_override_key(key: str) -> None:
    """
    Validate that a dotted override key matches a real config section and field.

    Raises:
        typer.BadParameter: If section or field is unknown.
    """
    parts = key.split(".", maxsplit=1)
    if len(parts) != 2:
        raise typer.BadParameter(f"Override key must be 'section.field', got: '{key}'")
    section, field = parts
    models = _section_models()
    if section not in models:
        raise typer.BadParameter(
            f"Unknown config section '{section}'. Valid sections: {sorted(models.keys())}"
        )
    valid_fields = set(models[section].model_fields)
    if field not in valid_fields:
        raise typer.BadParameter(
            f"Unknown field '{field}' in [{section}]. Valid fields: {sorted(valid_fields)}"
        )


# ── YAML generation helpers ────────────────────────────────────────────────


def _section_models() -> dict[str, type[BaseModel]]:
    """Lazy-load config classes and return the section → model mapping."""
    from orchard.core.config import (
        ArchitectureConfig,
        AugmentationConfig,
        DatasetConfig,
        EvaluationConfig,
        ExportConfig,
        HardwareConfig,
        OptunaConfig,
        TelemetryConfig,
        TrackingConfig,
        TrainingConfig,
    )

    return {
        "dataset": DatasetConfig,
        "architecture": ArchitectureConfig,
        "training": TrainingConfig,
        "augmentation": AugmentationConfig,
        "hardware": HardwareConfig,
        "telemetry": TelemetryConfig,
        "evaluation": EvaluationConfig,
        "tracking": TrackingConfig,
        "export": ExportConfig,
        "optuna": OptunaConfig,
    }


def _format_yaml_value(value: Any) -> str:
    """Format a Python value as an inline YAML scalar."""
    if value is None:
        return "null"
    if isinstance(value, bool):
        return "true" if value else "false"
    if isinstance(value, float):
        import math

        import yaml

        if not math.isclose(value, 0.0) and (abs(value) < 1e-3 or abs(value) >= 1e7):
            # yaml.dump appends '\n...\n'; take first line only
            return yaml.dump(value, default_flow_style=True).split("\n")[0]
        return str(value)
    if isinstance(value, int):
        return str(value)
    if isinstance(value, str):
        if value.lower() in _YAML_KEYWORDS:
            return f"'{value}'"
        if any(c in value for c in ":#{}[]|>&*!?,'\""):
            return f"'{value}'"
        return value
    import yaml

    return yaml.dump(value, default_flow_style=True).split("\n")[0]


def _inline_ref(schema: dict[str, Any], defs: dict[str, Any]) -> dict[str, Any]:
    """Resolve a single ``$ref`` property against ``$defs``."""
    ref_key = schema["$ref"].rsplit("/", 1)[-1]
    base = dict(defs.get(ref_key, {}))
    # Only inherit description if the property itself provides one;
    # prevents generic type docstrings (FloatRange, IntRange) leaking.
    if "description" not in schema:
        base.pop("description", None)
    for k, v in schema.items():
        if k != "$ref":
            base[k] = v
    return base


_CONSTRAINT_KEYS = ("minimum", "maximum", "exclusiveMinimum", "exclusiveMaximum", "enum")


def _merge_any_of(schema: dict[str, Any]) -> dict[str, Any]:
    """Merge ``anyOf`` variants, picking constraints from the first non-null type."""
    merged = {k: v for k, v in schema.items() if k != "anyOf"}
    for variant in schema["anyOf"]:
        if variant.get("type") != "null":
            for k in _CONSTRAINT_KEYS:
                if k in variant:
                    merged[k] = variant[k]
            break
    return merged


def _resolve_refs(properties: dict[str, Any], defs: dict[str, Any]) -> dict[str, Any]:
    """Inline ``$ref`` and ``anyOf`` references from JSON schema ``$defs``."""
    resolved: dict[str, Any] = {}
    for name, schema in properties.items():
        if "$ref" in schema:
            resolved[name] = _inline_ref(schema, defs)
        elif "anyOf" in schema:
            resolved[name] = _merge_any_of(schema)
        else:
            resolved[name] = schema
    return resolved


def _format_range_suffix(field_schema: dict[str, Any]) -> str | None:
    """Return a human-readable range suffix like ``(0\u2013100)`` or ``(\u2265 1)``."""
    lo_val = field_schema.get("minimum")
    lo_exc = field_schema.get("exclusiveMinimum")
    hi_val = field_schema.get("maximum")
    hi_exc = field_schema.get("exclusiveMaximum")

    lo = lo_val if lo_val is not None else lo_exc
    hi = hi_val if hi_val is not None else hi_exc

    if lo is not None and hi is not None:
        return f"({lo}\u2013{hi})"
    if lo is not None:
        return f"(\u2265 {lo_val})" if lo_val is not None else f"(> {lo_exc})"
    if hi is not None:
        return f"(\u2264 {hi_val})" if hi_val is not None else f"(< {hi_exc})"
    return None


def _build_comment(field_schema: dict[str, Any]) -> str | None:
    """Build a compact comment string from a resolved JSON schema property."""
    desc = field_schema.get("description")
    if not desc:
        return None

    if "enum" in field_schema:
        opts = ", ".join(str(v) for v in field_schema["enum"])
        return f"{desc} [{opts}]"

    suffix = _format_range_suffix(field_schema)
    if suffix:
        return f"{desc} {suffix}"
    return str(desc)


def _append_wrapped_comment(lines: list[str], comment: str, prefix: str) -> None:
    """Append a comment, wrapping long lines at word boundaries."""
    max_text = _COMMENT_MAX_WIDTH - len(prefix) - 2  # subtract "# "
    if len(comment) <= max_text:
        lines.append(f"{prefix}# {comment}")
        return
    import textwrap

    for wrapped in textwrap.wrap(comment, width=max_text):
        lines.append(f"{prefix}# {wrapped}")


def _render_fields(
    lines: list[str],
    data: dict[str, Any],
    properties: dict[str, Any],
    indent: int,
    defs: dict[str, Any],
) -> None:
    """Recursively render dict fields as commented YAML lines."""
    prefix = _YAML_INDENT * indent
    for key, value in data.items():
        field_schema = properties.get(key, {})

        comment = _build_comment(field_schema)
        if comment:
            _append_wrapped_comment(lines, comment, prefix)

        if isinstance(value, dict):
            lines.append(f"{prefix}{key}:")
            sub_props = _resolve_refs(field_schema.get("properties", {}), defs)
            _render_fields(lines, value, sub_props, indent + 1, defs)
        elif isinstance(value, list):
            lines.append(f"{prefix}{key}:")
            for item in value:
                if isinstance(item, dict):
                    lines.append(f"{prefix}{_YAML_INDENT}-")
                    _render_fields(lines, item, {}, indent + 2, defs)
                else:
                    lines.append(f"{prefix}{_YAML_INDENT}- {_format_yaml_value(item)}")
        else:
            lines.append(f"{prefix}{key}: {_format_yaml_value(value)}")


def _build_commented_yaml(data: dict[str, Any]) -> str:
    """Build a YAML string with descriptive comments above each field."""
    models: dict[str, type[BaseModel]] = _section_models()

    lines: list[str] = []
    for section_name, section_data in data.items():
        lines.append(f"{section_name}:")

        config_cls = models.get(section_name)
        if config_cls is not None:
            raw_schema = config_cls.model_json_schema()
            defs = raw_schema.get("$defs", {})
            properties = _resolve_refs(raw_schema.get("properties", {}), defs)
        else:
            properties = {}
            defs = {}

        _render_fields(lines, section_data, properties, indent=1, defs=defs)
        lines.append("")

    return "\n".join(lines) + "\n"


def _build_init_dict() -> dict[str, Any]:
    """
    Build a complete config dict with all defaults for recipe generation.

    Returns:
        Ordered dict with every config section dumped via ``model_dump(mode="json")``,
        paths sanitized to portable relative strings, and device reset to ``"auto"``.
    """
    models = _section_models()

    def dump(m: BaseModel) -> dict[str, Any]:
        """Serialize a config sub-model to a JSON-compatible dict."""
        return m.model_dump(mode="json")

    ds = dump(models["dataset"]())
    ds.pop("metadata", None)
    ds.pop("img_size", None)
    ds["data_root"] = "./dataset"

    hw = dump(models["hardware"]())
    hw["device"] = "auto"

    tel = dump(models["telemetry"]())
    tel["output_dir"] = "./outputs"

    return {
        "dataset": ds,
        "architecture": dump(models["architecture"]()),
        "training": dump(models["training"]()),
        "augmentation": dump(models["augmentation"]()),
        "hardware": hw,
        "telemetry": tel,
        "evaluation": dump(models["evaluation"]()),
        "tracking": dump(models["tracking"]()),
        "export": dump(models["export"]()),
        "optuna": dump(models["optuna"]()),
    }
