from __future__ import annotations

import re
from pathlib import Path

import yaml

from one_run.errors import OneRunError


def _ask(prompt: str, *, default: str | None = None) -> str:
    suffix = f" [{default}]" if default is not None else ""
    raw = input(f"{prompt}{suffix}: ").strip()
    if raw:
        return raw
    if default is not None:
        return default
    return ""


def _ask_choice(prompt: str, *, choices: tuple[str, ...], default: str) -> str:
    options = "/".join(choices)
    while True:
        value = _ask(f"{prompt} ({options})", default=default).lower()
        if value in choices:
            return value
        print(f"invalid choice: {value}. expected one of: {options}")


def _ask_int(prompt: str, *, default: int | None = None, required: bool = False) -> int | None:
    default_text = str(default) if default is not None else None
    while True:
        raw = _ask(prompt, default=default_text)
        if not raw:
            if required:
                print("value is required")
                continue
            return None
        try:
            return int(raw)
        except ValueError:
            print("expected integer value")


def _ask_optional_csv(prompt: str) -> list[str] | None:
    raw = _ask(prompt, default="")
    if not raw:
        return None
    values = [part.strip() for part in raw.split(",")]
    tags = [value for value in values if value]
    return tags or None


def _slugify_filename(value: str) -> str:
    slug = re.sub(r"[^a-zA-Z0-9._-]+", "-", value.strip().lower()).strip("-")
    return slug or "manifest"


def init_manifest(*, output: str | None, force: bool = False) -> int:

    print("one-run init manifest")
    print("press enter to accept default values")

    experiment_name = _ask("experiment name", default="local-test")
    seed = _ask_int("seed", default=42, required=True)
    tags = _ask_optional_csv("tags (comma-separated, optional)")
    color = _ask("experiment color (optional)", default="").lower() or None
    command = _ask("entrypoint command", default='python3 -c "print(\\"ok\\")"')

    environment_kind = _ask_choice("environment kind", choices=("local", "docker"), default="local")
    backend_kind = _ask_choice("backend kind", choices=("local", "slurm"), default="local")

    cpus = _ask_int("resources.cpus (optional)")
    mem_gb = _ask_int("resources.mem_gb (optional)")
    gpus = _ask_int("resources.gpus (optional)")
    time_min = _ask_int("resources.time_min (optional)")

    manifest: dict = {
        "experiment": {
            "name": experiment_name,
            "seed": seed,
        },
        "entrypoint": {
            "command": command,
        },
        "environment": {
            "kind": environment_kind,
        },
        "backend": {
            "kind": backend_kind,
        },
    }

    if tags:
        manifest["experiment"]["tags"] = tags
    if color:
        manifest["experiment"]["color"] = color

    if environment_kind == "docker":
        image = _ask("environment.image", default="python:3.11-slim")
        workdir = _ask("environment.workdir (optional)", default="") or None
        mounts = _ask_optional_csv("environment.mounts (comma-separated host:container, optional)")
        manifest["environment"]["image"] = image
        if workdir:
            manifest["environment"]["workdir"] = workdir
        if mounts:
            manifest["environment"]["mounts"] = mounts

    resources = {k: v for k, v in {
        "cpus": cpus,
        "mem_gb": mem_gb,
        "gpus": gpus,
        "time_min": time_min,
    }.items() if v is not None}
    if resources:
        manifest["resources"] = resources

    if backend_kind == "slurm":
        partition = _ask("backend.slurm.partition (optional)", default="") or None
        manifest["backend"]["slurm"] = {}
        if partition:
            manifest["backend"]["slurm"]["partition"] = partition

    if output:
        output_path = Path(output).resolve()
    else:
        filename = f"{_slugify_filename(experiment_name)}.manifest.yaml"
        output_path = (Path.cwd() / "manifests" / filename).resolve()

    if output_path.exists() and not force:
        raise OneRunError(f"file already exists: {output_path} (use --force to overwrite)")

    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_text(yaml.safe_dump(manifest, sort_keys=False), encoding="utf-8")

    print(f"written={output_path}")
    return 0


def init_listener(*, output: str | None, force: bool = False) -> int:
    print("one-run init listener")
    print("press enter to accept default values")

    kind = _ask_choice("listener kind", choices=("local", "webhook"), default="local")
    name = _ask("listener name", default="listener-profile")
    interval_sec = _ask_int("interval_sec (optional)", default=1, required=False)
    events = _ask_optional_csv("events (comma-separated, optional)")
    states = _ask_optional_csv("states (comma-separated, optional)")
    tags = _ask_optional_csv("tags (comma-separated, optional)")

    listener: dict = {
        "kind": kind,
        "name": name,
    }

    if events:
        listener["events"] = events
    if states:
        listener["states"] = states
    if tags:
        listener["tags"] = tags

    if kind == "local":
        handler = _ask("handler module:function (optional)", default="") or None
        if handler:
            listener["handler"] = handler
    else:
        webhook_url = _ask("webhook_url", default="http://127.0.0.1:8088/events")
        action = _ask("action (optional)", default="") or None
        timeout_sec = _ask_int("timeout_sec (optional)", default=3, required=False)
        listener["webhook_url"] = webhook_url
        if action:
            listener["action"] = action
        if timeout_sec is not None:
            listener["timeout_sec"] = timeout_sec

    payload: dict = {
        "listeners": [listener],
    }
    if interval_sec is not None:
        payload["interval_sec"] = interval_sec

    if output:
        output_path = Path(output).resolve()
    else:
        filename = f"{_slugify_filename(name)}.listener.yaml"
        output_path = (Path.cwd() / "listeners" / filename).resolve()

    if output_path.exists() and not force:
        raise OneRunError(f"file already exists: {output_path} (use --force to overwrite)")

    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_text(yaml.safe_dump(payload, sort_keys=False), encoding="utf-8")
    print(f"written={output_path}")
    return 0
