from __future__ import annotations

import os
import re
import shlex
import signal
import sys
import threading
import time
from dataclasses import dataclass
from typing import Annotated, Any, Dict, List, Optional, Sequence, Tuple, Mapping

import typer

from iac.utils.subproc import ControlledProcess

app = typer.Typer(
    add_completion=False,
    help="Run steps from a book.yaml (automation-first).",
    no_args_is_help=True,
)

# ----------------------------
# YAML loading
# ----------------------------


def _load_yaml(path: str) -> Dict[str, Any]:
    try:
        import yaml  # PyYAML
    except Exception as e:
        raise typer.BadParameter(
            "PyYAML is required to use runbooks. Install with: pip install pyyaml"
        ) from e

    try:
        with open(path, "r", encoding="utf-8") as f:
            data = yaml.safe_load(f)
    except FileNotFoundError as e:
        raise typer.BadParameter(f"Book file not found: {path}") from e
    except Exception as e:
        raise typer.BadParameter(f"Failed to parse YAML: {path}: {e}") from e

    if not isinstance(data, dict):
        raise typer.BadParameter("Book YAML must be a mapping/object at the top level.")
    return data


# ----------------------------
# Templating (minimal + deterministic)
#   - {{ var }} from merged vars
#   - ${ENV} from os.environ (pass-through)
# ----------------------------

_VAR_RE = re.compile(r"\{\{\s*([A-Za-z_][A-Za-z0-9_]*)\s*\}\}")
_ENV_RE = re.compile(r"\$\{([A-Za-z_][A-Za-z0-9_]*)\}")


def _render_str(s: str, *, vars_map: Dict[str, str], env_map: Mapping[str, str]) -> str:
    def repl_var(m: re.Match) -> str:
        key = m.group(1)
        return vars_map.get(key, "")

    def repl_env(m: re.Match) -> str:
        key = m.group(1)
        return env_map.get(key, "")

    s = _VAR_RE.sub(repl_var, s)
    s = _ENV_RE.sub(repl_env, s)
    return s


def _render_obj(
    obj: Any, *, vars_map: Dict[str, str], env_map: Mapping[str, str]
) -> Any:
    if obj is None:
        return None
    if isinstance(obj, str):
        return _render_str(obj, vars_map=vars_map, env_map=env_map)
    if isinstance(obj, list):
        return [_render_obj(x, vars_map=vars_map, env_map=env_map) for x in obj]
    if isinstance(obj, dict):
        return {
            str(k): _render_obj(v, vars_map=vars_map, env_map=env_map)
            for k, v in obj.items()
        }
    return obj


# ----------------------------
# Vars parsing (repeatable + env-friendly)
# ----------------------------


def _parse_kv(items: List[str]) -> Dict[str, str]:
    out: Dict[str, str] = {}
    for item in items:
        if "=" not in item:
            raise typer.BadParameter(
                f"Invalid --var value {item!r}. Expected KEY=VALUE."
            )
        k, v = item.split("=", 1)
        k = k.strip()
        if not k:
            raise typer.BadParameter(
                f"Invalid --var value {item!r}. KEY cannot be empty."
            )
        out[k] = v
    return out


def _parse_vars_env(s: Optional[str]) -> List[str]:
    """
    Parse IAC_VARS into ['k=v', 'k2=v2'].

    Accepts:
      - comma-separated: "k=v,k2=v2"
      - whitespace/newline-separated: "k=v k2=v2"
      - mixed
    """
    if not s:
        return []
    parts: List[str] = []
    for chunk in re.split(r"[,\n\r\t ]+", s.strip()):
        if not chunk:
            continue
        parts.append(chunk)
    return parts


# ----------------------------
# Book model validation + step lookup
# ----------------------------


@dataclass(frozen=True)
class StepDef:
    name: str
    desc: str
    cwd: Optional[str]
    env: Dict[str, str]
    argv: Optional[List[str]]
    shell: Optional[str]
    timeout: Optional[float]


def _normalize_book(
    book: Dict[str, Any],
) -> Tuple[Dict[str, str], Dict[str, str], Dict[str, StepDef], List[StepDef]]:
    version = book.get("version", 1)
    if version != 1:
        raise typer.BadParameter(
            f"Unsupported book version: {version}. Expected version: 1"
        )

    vars_raw = book.get("vars", {}) or {}
    env_raw = book.get("env", {}) or {}
    steps_raw = book.get("steps", None)

    if not isinstance(vars_raw, dict):
        raise typer.BadParameter("book.vars must be a mapping")
    if not isinstance(env_raw, dict):
        raise typer.BadParameter("book.env must be a mapping")
    if not isinstance(steps_raw, list) or not steps_raw:
        raise typer.BadParameter("book.steps must be a non-empty list")

    vars_map: Dict[str, str] = {str(k): str(v) for k, v in vars_raw.items()}
    env_map: Dict[str, str] = {str(k): str(v) for k, v in env_raw.items()}

    steps_list: List[StepDef] = []
    steps_by_name: Dict[str, StepDef] = {}

    for i, s in enumerate(steps_raw):
        if not isinstance(s, dict):
            raise typer.BadParameter(f"step[{i}] must be a mapping")

        name = s.get("name")
        if not isinstance(name, str) or not name.strip():
            raise typer.BadParameter(f"step[{i}].name is required")
        name = name.strip()
        if name in steps_by_name:
            raise typer.BadParameter(f"Duplicate step name: {name}")

        desc = s.get("desc") or ""
        if not isinstance(desc, str):
            raise typer.BadParameter(f"step[{i}].desc must be a string if provided")

        cwd = s.get("cwd")
        if cwd is not None and not isinstance(cwd, str):
            raise typer.BadParameter(f"step[{i}].cwd must be a string if provided")

        step_env = s.get("env", {}) or {}
        if not isinstance(step_env, dict):
            raise typer.BadParameter(f"step[{i}].env must be a mapping if provided")

        argv = s.get("argv")
        shell = s.get("shell")

        if argv is not None and shell is not None:
            raise typer.BadParameter(
                f"step[{i}] must specify only one of 'argv' or 'shell'"
            )
        if argv is None and shell is None:
            raise typer.BadParameter(f"step[{i}] must specify one of 'argv' or 'shell'")

        if argv is not None:
            if (
                not isinstance(argv, list)
                or not all(isinstance(x, str) for x in argv)
                or not argv
            ):
                raise typer.BadParameter(
                    f"step[{i}].argv must be a non-empty list of strings"
                )
            argv_norm = list(argv)
            shell_norm = None
        else:
            if not isinstance(shell, str) or not shell.strip():
                raise typer.BadParameter(f"step[{i}].shell must be a non-empty string")
            argv_norm = None
            shell_norm = shell

        timeout = s.get("timeout")
        timeout_norm: Optional[float]
        if timeout is None:
            timeout_norm = None
        else:
            try:
                timeout_norm = float(timeout)
            except Exception:
                raise typer.BadParameter(
                    f"step[{i}].timeout must be a number (seconds)"
                )

        step = StepDef(
            name=name,
            desc=desc,
            cwd=cwd,
            env={str(k): str(v) for k, v in step_env.items()},
            argv=argv_norm,
            shell=shell_norm,
            timeout=timeout_norm,
        )
        steps_by_name[name] = step
        steps_list.append(step)

    return vars_map, env_map, steps_by_name, steps_list


# ----------------------------
# Process runner (mirrors your exec behavior)
# ----------------------------


@dataclass
class _CaptureSinks:
    stdout_fh: Optional[object] = None
    stderr_fh: Optional[object] = None

    def close(self) -> None:
        for fh in (self.stdout_fh, self.stderr_fh):
            if fh is not None:
                try:
                    fh.close()
                except Exception:
                    pass


def _open_out_file(path: Optional[str]):
    if not path:
        return None
    return open(path, "ab")


class SignalForwarder:
    def __init__(
        self,
        proc: ControlledProcess,
        *,
        use_process_group: bool,
        forward_signals: Sequence[int],
    ) -> None:
        self.proc = proc
        self.use_process_group = use_process_group
        self.forward_signals = list(forward_signals)

        self._old_handlers: Dict[int, object] = {}
        self._received = threading.Event()
        self._last_signal: Optional[int] = None
        self._lock = threading.Lock()

    @property
    def received(self) -> bool:
        return self._received.is_set()

    @property
    def last_signal(self) -> Optional[int]:
        with self._lock:
            return self._last_signal

    def install(self) -> None:
        for sig in self.forward_signals:
            try:
                self._old_handlers[sig] = signal.getsignal(sig)
                signal.signal(sig, self._handler)  # type: ignore[arg-type]
            except Exception:
                continue

    def uninstall(self) -> None:
        for sig, old in self._old_handlers.items():
            try:
                signal.signal(sig, old)  # type: ignore[arg-type]
            except Exception:
                continue
        self._old_handlers.clear()

    def _handler(self, signum: int, _frame) -> None:
        with self._lock:
            self._last_signal = signum
        self._received.set()
        try:
            if self.use_process_group:
                self.proc.killpg(signum)
            else:
                self.proc.send_signal(signum)
        except Exception:
            pass


def _exit_code_for_signal(signum: int) -> int:
    return 128 + int(signum)


def _run_controlled(
    argv: List[str],
    *,
    cwd: Optional[str],
    env: Dict[str, str],
    timeout: Optional[float],
    pty: bool,
    killpg: bool,
    forward_signals: bool,
    capture: bool,
    stdout_file: Optional[str],
    stderr_file: Optional[str],
    print_captured: bool,
    quiet: bool,
    verbose: bool,
    display: str,
) -> None:
    sinks = _CaptureSinks(
        stdout_fh=_open_out_file(stdout_file), stderr_fh=_open_out_file(stderr_file)
    )

    def on_out_line(line: str) -> None:
        if quiet:
            return
        sys.stdout.write(line)
        sys.stdout.flush()

    def on_err_line(line: str) -> None:
        if quiet:
            return
        sys.stderr.write(line)
        sys.stderr.flush()

    def on_out_bytes(b: bytes) -> None:
        if sinks.stdout_fh is not None:
            try:
                sinks.stdout_fh.write(b)
                sinks.stdout_fh.flush()
            except Exception:
                pass

    def on_err_bytes(b: bytes) -> None:
        if sinks.stderr_fh is not None:
            try:
                sinks.stderr_fh.write(b)
                sinks.stderr_fh.flush()
            except Exception:
                pass

    p = ControlledProcess(
        argv,
        cwd=cwd,
        env=env or None,
        inherit_env=True,  # we already merged pass-through into env as desired
        text=True,
        pty=pty,
        capture=capture,
        merge_stderr=False,
        start_new_session=True,
    )

    sigs: List[int] = []
    if forward_signals:
        sigs.extend([signal.SIGINT, signal.SIGTERM])
        if hasattr(signal, "SIGHUP"):
            sigs.append(signal.SIGHUP)

    forwarder = SignalForwarder(p, use_process_group=killpg, forward_signals=sigs)

    try:
        p.start(
            on_stdout=on_out_line,
            on_stderr=on_err_line,
            on_stdout_bytes=on_out_bytes,
            on_stderr_bytes=on_err_bytes,
        )

        if verbose:
            typer.echo(f"[pid={p.pid} pgid={p.pgid}] {display}", err=True)

        forwarder.install()

        start_t = time.time()
        while True:
            rc = p.poll()
            if rc is not None:
                break

            if timeout is not None and (time.time() - start_t) > timeout:
                typer.echo(
                    f"Step timed out after {timeout} seconds; stopping...", err=True
                )
                p.stop(grace=2.0, use_process_group=killpg)
                raise typer.Exit(code=124)

            if forwarder.received:
                time.sleep(0.2)
                if p.is_running():
                    p.stop(grace=2.0, use_process_group=killpg)
                sig = forwarder.last_signal or signal.SIGINT
                raise typer.Exit(code=_exit_code_for_signal(sig))

            time.sleep(0.1)

    except FileNotFoundError as e:
        missing = getattr(e, "filename", None) or str(e)
        typer.echo(f"Executable not found: {missing}", err=True)
        raise typer.Exit(code=127)

    except KeyboardInterrupt:
        typer.echo("Interrupted; stopping child process...", err=True)
        try:
            p.stop(grace=2.0, use_process_group=killpg)
        finally:
            raise typer.Exit(code=130)

    finally:
        forwarder.uninstall()
        sinks.close()

    rc_final = p.returncode if p.returncode is not None else 1

    if capture and print_captured and not quiet:
        out = p.stdout_text()
        err = p.stderr_text()
        if out:
            sys.stdout.write(out)
            sys.stdout.flush()
        if err:
            sys.stderr.write(err)
            sys.stderr.flush()

    raise typer.Exit(code=int(rc_final))


# ----------------------------
# CLI
# ----------------------------


@app.command("steps")
def steps_cmd(
    book: Annotated[
        str,
        typer.Option(
            "--book", envvar="IAC_BOOK", help="Path to book.yaml (or set IAC_BOOK)."
        ),
    ] = "book.yaml",
) -> None:
    """List available steps in the run book."""
    data = _load_yaml(book)
    _vars, _env, _by_name, steps_list = _normalize_book(data)

    for s in steps_list:
        if s.desc:
            typer.echo(f"{s.name}\t{s.desc}")
        else:
            typer.echo(f"{s.name}")


@app.command("step")
def step_cmd(
    step: Annotated[
        str, typer.Argument(help="Step name to run (from book.steps[].name).")
    ],
    book: Annotated[
        str,
        typer.Option(
            "--book", envvar="IAC_BOOK", help="Path to book.yaml (or set IAC_BOOK)."
        ),
    ] = "book.yaml",
    var: Annotated[
        List[str],
        typer.Option(
            "--var",
            help="Override book vars (repeatable): --var key=value",
            show_default=False,
        ),
    ] = [],
    vars_env: Annotated[
        Optional[str],
        typer.Option(
            "--vars-env",
            envvar="IAC_VARS",
            help="Additional vars from env (IAC_VARS). Format: 'k=v,k2=v2' or whitespace-separated.",
        ),
    ] = None,
    dry_run: Annotated[
        bool,
        typer.Option(
            "--dry-run",
            help="Print the resolved step (cwd/env/command) without executing.",
        ),
    ] = False,
    # passthrough runner options (automation friendly)
    quiet: Annotated[
        bool, typer.Option("-q", "--quiet", help="Do not stream output to console.")
    ] = False,
    stdout_file: Annotated[
        Optional[str],
        typer.Option("--stdout-file", help="Append stdout bytes to this file"),
    ] = None,
    stderr_file: Annotated[
        Optional[str],
        typer.Option("--stderr-file", help="Append stderr bytes to this file"),
    ] = None,
    capture: Annotated[
        bool,
        typer.Option(
            "--capture/--no-capture",
            help="Capture output in memory (still streams live).",
        ),
    ] = False,
    print_captured: Annotated[
        bool,
        typer.Option(
            "--print-captured/--no-print-captured",
            help="If --capture, print captured output at end.",
        ),
    ] = False,
    pty: Annotated[
        bool, typer.Option("--pty", help="Run under a PTY (TTY-like).")
    ] = False,
    timeout: Annotated[
        Optional[float],
        typer.Option("--timeout", help="Override step timeout (seconds)."),
    ] = None,
    killpg: Annotated[
        bool,
        typer.Option(
            "--killpg/--no-killpg", help="Use process-group signals/termination."
        ),
    ] = True,
    forward_signals: Annotated[
        bool,
        typer.Option(
            "--forward-signals/--no-forward-signals",
            help="Forward SIGINT/SIGTERM(/SIGHUP) to child.",
        ),
    ] = True,
    verbose: Annotated[
        bool, typer.Option("-v", "--verbose", help="Print extra process metadata.")
    ] = False,
) -> None:
    """Run a single step from a book.yaml"""

    data = _load_yaml(book)
    book_vars, book_env, steps_by_name, _steps_list = _normalize_book(data)

    if step not in steps_by_name:
        typer.echo(f"Unknown step: {step}", err=True)
        typer.echo(
            "Tip: run `iac run steps --book <path>` to list available steps.", err=True
        )
        raise typer.Exit(code=2)

    s = steps_by_name[step]

    # Merge vars in precedence order:
    # book.vars < IAC_VARS/--vars-env < repeated --var
    vars_items: List[str] = []
    vars_items.extend(_parse_vars_env(vars_env))
    vars_items.extend(var)
    cli_vars = _parse_kv(vars_items)

    merged_vars: Dict[str, str] = dict(book_vars)
    merged_vars.update(cli_vars)

    # Render global env first (may reference vars and host env)
    rendered_book_env = _render_obj(book_env, vars_map=merged_vars, env_map=os.environ)
    rendered_step_env = _render_obj(s.env, vars_map=merged_vars, env_map=os.environ)

    if not isinstance(rendered_book_env, dict) or not isinstance(
        rendered_step_env, dict
    ):
        raise typer.BadParameter("Rendered env must be a mapping")

    merged_env: Dict[str, str] = {}
    merged_env.update({str(k): str(v) for k, v in rendered_book_env.items()})
    merged_env.update({str(k): str(v) for k, v in rendered_step_env.items()})

    # Render cwd + command
    cwd = _render_obj(s.cwd, vars_map=merged_vars, env_map=os.environ)
    cwd = str(cwd) if cwd is not None else None

    # Resolve command into argv + display string
    if s.argv is not None:
        argv = _render_obj(s.argv, vars_map=merged_vars, env_map=os.environ)
        if (
            not isinstance(argv, list)
            or not all(isinstance(x, str) for x in argv)
            or not argv
        ):
            raise typer.BadParameter(
                "Rendered argv must be a non-empty list of strings"
            )
        display = shlex.join(argv)
    else:
        shell_str = _render_obj(s.shell, vars_map=merged_vars, env_map=os.environ)
        if not isinstance(shell_str, str) or not shell_str.strip():
            raise typer.BadParameter("Rendered shell must be a non-empty string")
        argv = ["bash", "-lc", shell_str]
        display = shell_str

    # Timeout: CLI override > step timeout > none
    eff_timeout = timeout if timeout is not None else s.timeout

    if dry_run:
        typer.echo(f"step: {s.name}")
        if s.desc:
            typer.echo(f"desc: {s.desc}")
        if cwd:
            typer.echo(f"cwd:  {cwd}")
        typer.echo(f"cmd:  {display}")
        if eff_timeout is not None:
            typer.echo(f"timeout: {eff_timeout}s")
        # keep env output concise for automation: show keys + a couple common ones
        if merged_env:
            keys = sorted(merged_env.keys())
            typer.echo(f"env:  {', '.join(keys)}")
        raise typer.Exit(code=0)

    _run_controlled(
        argv=argv,
        cwd=cwd,
        env=merged_env,
        timeout=eff_timeout,
        pty=pty,
        killpg=killpg,
        forward_signals=forward_signals,
        capture=capture,
        stdout_file=stdout_file,
        stderr_file=stderr_file,
        print_captured=print_captured,
        quiet=quiet,
        verbose=verbose,
        display=display,
    )
