from __future__ import annotations

import datetime as dt
import json
import os
import re
import shlex
import signal
import sys
import tempfile
import threading
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Dict, List, Literal, Optional, Sequence

from iac.utils.subproc import ControlledProcess

from .cel_eval import CelEngine
from .model import Book, Step
from .sensitive import SecretRedactor, sensitive_values

Status = Literal["success", "skipped", "failed", "failed_nonblocking"]


@dataclass(frozen=True)
class StepResult:
    book_path: str
    step: str
    status: Status
    exit_code: int
    duration_ms: int
    message: str = ""
    command: str = ""
    cwd: str = ""
    when: str = ""
    started_at: str = ""
    ended_at: str = ""
    stdin_file: str = ""
    stdout_file: str = ""
    stderr_file: str = ""
    attempts: int = 1
    max_attempts: int = 1
    outputs: Dict[str, Any] = field(default_factory=dict)


@dataclass(frozen=True)
class RunOptions:
    quiet: bool = False
    stdin_file: Optional[str] = None
    stdout_file: Optional[str] = None
    stderr_file: Optional[str] = None
    file_mode: str = "append"
    artifacts_dir: Optional[str] = None
    capture: bool = False
    print_captured: bool = False
    pty: bool = False
    killpg: bool = True
    forward_signals: bool = True
    timeout_override: Optional[float] = None
    retries_override: Optional[int] = None
    retry_delay_override: Optional[float] = None
    verbose: bool = False
    run_id: str = ""


@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], *, file_mode: str):
    if not path:
        return None
    mode = "ab" if file_mode == "append" else "wb"
    return open(path, mode)


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 _now_iso_utc() -> str:
    return dt.datetime.now(dt.timezone.utc).isoformat()


def _safe_name(name: str) -> str:
    return re.sub(r"[^A-Za-z0-9_.-]+", "_", name).strip("_") or "step"


def _resolve_file_setting(value: Any) -> Optional[str]:
    if value is None:
        return None
    s = str(value).strip()
    return s or None


def _ensure_parent(path: str) -> None:
    Path(path).expanduser().resolve().parent.mkdir(parents=True, exist_ok=True)


def _artifact_paths(
    opts: RunOptions, *, step_name: str, step_index: int
) -> tuple[Optional[str], Optional[str]]:
    if not opts.artifacts_dir:
        return None, None
    run_id = opts.run_id or "run"
    base = Path(opts.artifacts_dir).expanduser().resolve() / run_id
    base.mkdir(parents=True, exist_ok=True)
    stem = f"{step_index:03d}_{_safe_name(step_name)}"
    return (str(base / f"{stem}.stdout.log"), str(base / f"{stem}.stderr.log"))


def _execute_hook(
    hook_name: str,
    command: str,
    *,
    env: Dict[str, str],
    cwd: Optional[str],
    quiet: bool,
) -> int:
    proc = ControlledProcess(
        ["bash", "-lc", command],
        cwd=cwd,
        env=env,
        inherit_env=True,
        text=True,
        capture=not quiet,
        pty=False,
        merge_stderr=False,
        start_new_session=True,
    )
    proc.start()
    rc = proc.wait(timeout=None)
    return int(rc)


def _maybe_run_hook(
    book: Book,
    hook_name: str,
    *,
    book_path: Path,
    step_name: Optional[str],
    step_status: Optional[str],
    step_exit: Optional[int],
    run_id: str,
    report_path: Optional[str],
    quiet: bool,
) -> bool:
    hook = book.hooks.get(hook_name)
    if hook is None:
        return True

    env = {
        "IAC_HOOK_NAME": hook_name,
        "IAC_RUN_ID": run_id,
        "IAC_BOOK_PATH": str(book_path),
        "IAC_STEP_NAME": step_name or "",
        "IAC_STEP_STATUS": step_status or "",
        "IAC_STEP_EXIT_CODE": "" if step_exit is None else str(step_exit),
        "IAC_REPORT_PATH": report_path or "",
    }
    rc = _execute_hook(
        hook_name, hook.command, env=env, cwd=str(book_path.parent), quiet=quiet
    )
    if rc == 0:
        return True
    if hook.on_fail == "pass":
        return True
    if hook.on_fail == "warn":
        if not quiet:
            sys.stderr.write(
                f"WARNING: hook '{hook_name}' failed (warn), exit_code={rc}\n"
            )
            sys.stderr.flush()
        return True
    if not quiet:
        sys.stderr.write(f"hook '{hook_name}' failed (halt), exit_code={rc}\n")
        sys.stderr.flush()
    return False


def resolve_step(
    step: Step,
    *,
    book_path: Path,
    vars_map: Dict[str, Any],
    module_dirs: Optional[List[Path]] = None,
) -> tuple[
    List[str], str, Optional[str], Dict[str, str], Optional[float], Optional[str]
]:
    """Returns: argv, display, cwd, env, timeout, when_expr."""
    env: Dict[str, str] = {}
    cwd_s = step.cwd.strip() if isinstance(step.cwd, str) and step.cwd.strip() else None
    when_expr = step.when
    timeout = step.timeout

    if step.mode == "argv":
        if not step.argv:
            raise ValueError(f"step {step.name} is missing sh argv")
        argv = list(step.argv)
        return argv, shlex.join(argv), cwd_s, env, timeout, when_expr

    module_ref = (step.module or "").strip()
    if not module_ref:
        raise ValueError(f"step {step.name} is missing fn reference")

    env["IAC_MODULE_REF"] = module_ref
    env["IAC_MODULE_BOOK_PATH"] = str(book_path)
    env["IAC_MODULE_STEP_NAME"] = step.name
    env["IAC_MODULE_VARS_JSON"] = json.dumps(vars_map)
    if module_dirs:
        env["IAC_MODULE_DIRS"] = ":".join(str(d) for d in module_dirs)

    argv = [sys.executable, "-m", "iac.runbook.module_exec"]
    display = f"{sys.executable} -m iac.runbook.module_exec {module_ref}"
    return argv, display, cwd_s, env, timeout, when_expr


def run_one_step(
    book: Book,
    step: Step,
    *,
    book_path: Path,
    vars_map: Dict[str, Any],
    opts: RunOptions,
    step_index: int = 0,
    report_path: Optional[str] = None,
    module_dirs: Optional[List[Path]] = None,
    cel: Optional[CelEngine] = None,
    inputs_map: Optional[Dict[str, Any]] = None,
    steps_ctx: Optional[Dict[str, Any]] = None,
    run_ctx: Optional[Dict[str, Any]] = None,
) -> StepResult:
    cel_engine = cel if cel is not None else CelEngine.default()
    with_values: Dict[str, Any] = {}

    argv, display, cwd, env, timeout, when_expr = resolve_step(
        step,
        book_path=book_path,
        vars_map=vars_map,
        module_dirs=module_dirs,
    )

    if step.vars_map:
        eval_ctx = {
            "inputs": inputs_map or vars_map,
            "steps": steps_ctx or {},
            "run": run_ctx or {},
        }
        for key, expr in step.vars_map.items():
            with_values[key] = cel_engine.eval(expr, eval_ctx)
        env = dict(env)
        env.update({k: str(v) for k, v in with_values.items()})
        if step.mode == "module":
            module_vars = dict(vars_map)
            module_vars.update(with_values)
            env["IAC_MODULE_VARS_JSON"] = json.dumps(module_vars)

    file_mode = opts.file_mode if opts.file_mode in ("append", "truncate") else "append"

    art_stdout, art_stderr = _artifact_paths(
        opts, step_name=step.name, step_index=step_index
    )
    eff_stdin = _resolve_file_setting(opts.stdin_file)
    eff_stdout = _resolve_file_setting(opts.stdout_file) or art_stdout
    eff_stderr = _resolve_file_setting(opts.stderr_file) or art_stderr

    redactor = SecretRedactor.from_values(
        sensitive_values(book, vars_map={k: str(v) for k, v in vars_map.items()})
    )
    display_redacted = redactor.redact_text(display)

    if eff_stdout:
        _ensure_parent(eff_stdout)
    if eff_stderr:
        _ensure_parent(eff_stderr)
    if eff_stdin and not Path(eff_stdin).expanduser().exists():
        return StepResult(
            book_path=str(book_path),
            step=step.name,
            status="failed",
            exit_code=1,
            duration_ms=0,
            message=f"stdin file not found: {eff_stdin}",
            command=display_redacted,
            cwd=cwd or "",
            when=when_expr or "",
            started_at="",
            ended_at="",
            stdin_file=eff_stdin,
            stdout_file=eff_stdout or "",
            stderr_file=eff_stderr or "",
        )

    eff_timeout = (
        opts.timeout_override if opts.timeout_override is not None else timeout
    )
    eff_retries = (
        opts.retries_override if opts.retries_override is not None else step.retries
    )
    eff_retry_delay = (
        opts.retry_delay_override
        if opts.retry_delay_override is not None
        else step.retry_delay
    )
    max_attempts = 1 + max(0, int(eff_retries))
    eff_retry_delay = max(0.0, float(eff_retry_delay))

    if when_expr:
        ok = bool(
            cel_engine.eval(
                when_expr,
                {
                    "inputs": inputs_map or vars_map,
                    "steps": steps_ctx or {},
                    "run": run_ctx or {},
                    "vars": with_values,
                },
            )
        )
        if not ok:
            return StepResult(
                book_path=str(book_path),
                step=step.name,
                status="skipped",
                exit_code=0,
                duration_ms=0,
                message="when evaluated false",
                command=display_redacted,
                cwd=cwd or "",
                when=when_expr or "",
                started_at="",
                ended_at="",
                stdin_file=eff_stdin or "",
                stdout_file=eff_stdout or "",
                stderr_file=eff_stderr or "",
                attempts=1,
                max_attempts=max_attempts,
            )

    t0 = time.time()
    started_at = _now_iso_utc()

    if not _maybe_run_hook(
        book,
        "on_step_start",
        book_path=book_path,
        step_name=step.name,
        step_status=None,
        step_exit=None,
        run_id=opts.run_id,
        report_path=report_path,
        quiet=opts.quiet,
    ):
        ended_at = _now_iso_utc()
        return StepResult(
            book_path=str(book_path),
            step=step.name,
            status="failed",
            exit_code=1,
            duration_ms=0,
            message="on_step_start hook failed",
            command=display_redacted,
            cwd=cwd or "",
            when=when_expr or "",
            started_at=started_at,
            ended_at=ended_at,
            stdin_file=eff_stdin or "",
            stdout_file=eff_stdout or "",
            stderr_file=eff_stderr or "",
            attempts=1,
            max_attempts=max_attempts,
        )

    rc: int = 1
    msg = ""
    status: Status = "success"
    attempts_made = 0
    captured_proc: Optional[ControlledProcess] = None
    interrupted = False
    timed_out = False
    module_result_file: Optional[str] = None
    module_step_outputs: Dict[str, Any] = {}

    if step.mode == "module":
        fd, path = tempfile.mkstemp(
            prefix=f"iac_mod_{_safe_name(step.name)}_", suffix=".json"
        )
        os.close(fd)
        module_result_file = path
        env = dict(env)
        env["IAC_MODULE_RESULT_FILE"] = module_result_file
        env["IAC_MODULE_RUN_ID"] = opts.run_id

    for attempt in range(1, max_attempts + 1):
        attempts_made = attempt
        attempt_file_mode = file_mode if attempt == 1 else "append"
        sinks = _CaptureSinks(
            stdout_fh=_open_out_file(eff_stdout, file_mode=attempt_file_mode),
            stderr_fh=_open_out_file(eff_stderr, file_mode=attempt_file_mode),
        )

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

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

        def on_out_bytes(b: bytes) -> None:
            if sinks.stdout_fh is not None:
                try:
                    b = redactor.redact_bytes(b)
                    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:
                    b = redactor.redact_bytes(b)
                    sinks.stderr_fh.write(b)
                    sinks.stderr_fh.flush()
                except Exception:
                    pass

        p = ControlledProcess(
            argv,
            cwd=cwd,
            env=env or None,
            inherit_env=True,
            text=True,
            pty=opts.pty,
            capture=True,
            merge_stderr=False,
            start_new_session=True,
        )
        captured_proc = p

        sigs: List[int] = []
        if opts.forward_signals:
            sigs.extend([signal.SIGINT, signal.SIGTERM])
            if hasattr(signal, "SIGHUP"):
                sigs.append(signal.SIGHUP)
        forwarder = SignalForwarder(
            p, use_process_group=opts.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 eff_stdin:
                with open(eff_stdin, "rb") as fh:
                    p.write(fh.read())
                p.close_stdin()

            if opts.verbose:
                sys.stderr.write(f"[pid={p.pid} pgid={p.pgid}] {display_redacted}\n")
                sys.stderr.flush()

            forwarder.install()

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

                if eff_timeout is not None and (time.time() - start_t) > eff_timeout:
                    sys.stderr.write(
                        f"Step timed out after {eff_timeout} seconds; stopping...\n"
                    )
                    sys.stderr.flush()
                    p.stop(grace=2.0, use_process_group=opts.killpg)
                    rc = 124
                    msg = "timeout"
                    timed_out = True
                    status = "failed"
                    break

                if forwarder.received:
                    time.sleep(0.2)
                    if p.is_running():
                        p.stop(grace=2.0, use_process_group=opts.killpg)
                    sig = forwarder.last_signal or signal.SIGINT
                    rc = _exit_code_for_signal(sig)
                    msg = f"terminated by signal {sig}"
                    interrupted = True
                    status = "failed"
                    break

                time.sleep(0.1)

        except FileNotFoundError as e:
            missing = getattr(e, "filename", None) or str(e)
            rc = 127
            msg = f"executable not found: {missing}"
            status = "failed"

        except KeyboardInterrupt:
            sys.stderr.write("Interrupted; stopping child process...\n")
            sys.stderr.flush()
            try:
                p.stop(grace=2.0, use_process_group=opts.killpg)
            finally:
                rc = 130
                msg = "keyboard interrupt"
                interrupted = True
                status = "failed"

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

        if rc == 0:
            status = "success"
            break
        if interrupted:
            break
        if attempt < max_attempts:
            if not opts.quiet:
                sys.stderr.write(
                    f"retrying step '{step.name}' ({attempt + 1}/{max_attempts}) after exit_code={rc}\n"
                )
                sys.stderr.flush()
            if eff_retry_delay > 0:
                time.sleep(eff_retry_delay)

    dt_ms = int((time.time() - t0) * 1000)
    ended_at = _now_iso_utc()

    if rc != 0:
        if attempts_made > 1 and not msg:
            msg = f"failed after {attempts_made}/{max_attempts} attempts"
        elif timed_out and attempts_made > 1:
            msg = f"timeout after {attempts_made}/{max_attempts} attempts"
        if step.on_fail in ("pass", "warn"):
            status = "failed_nonblocking"
            if step.on_fail == "warn":
                sys.stderr.write(
                    f"WARNING: step '{step.name}' failed (warn), exit_code={rc}\n"
                )
                sys.stderr.flush()
        else:
            status = "failed"
    elif module_result_file:
        try:
            p = Path(module_result_file)
            if p.exists():
                data = json.loads(p.read_text(encoding="utf-8"))
                if isinstance(data, dict):
                    msg = str(data.get("message", "") or "")
                    raw_out = data.get("outputs")
                    if isinstance(raw_out, dict):
                        module_step_outputs = dict(raw_out)
        except Exception:
            pass
        finally:
            try:
                Path(module_result_file).unlink(missing_ok=True)
            except Exception:
                pass

    if (
        opts.capture
        and opts.print_captured
        and not opts.quiet
        and captured_proc is not None
    ):
        out = captured_proc.stdout_text()
        err = captured_proc.stderr_text()
        if out:
            sys.stdout.write(redactor.redact_text(out))
            sys.stdout.flush()
        if err:
            sys.stderr.write(redactor.redact_text(err))
            sys.stderr.flush()

    if not _maybe_run_hook(
        book,
        "on_step_end",
        book_path=book_path,
        step_name=step.name,
        step_status=status,
        step_exit=rc,
        run_id=opts.run_id,
        report_path=report_path,
        quiet=opts.quiet,
    ):
        status = "failed"
        rc = 1
        msg = "on_step_end hook failed"

    captured_stdout = ""
    captured_stderr = ""
    if captured_proc is not None:
        try:
            captured_stdout = captured_proc.stdout_text() or ""
            captured_stderr = captured_proc.stderr_text() or ""
        except Exception:
            captured_stdout = ""
            captured_stderr = ""

    step_outputs: Dict[str, Any] = {}
    if step.mode == "module" and module_step_outputs:
        step_outputs.update(module_step_outputs)
    if step.outputs_map:
        eval_ctx = {
            "inputs": inputs_map or vars_map,
            "steps": steps_ctx or {},
            "run": run_ctx or {},
            "vars": with_values,
            "result": {
                "stdout": captured_stdout,
                "stderr": captured_stderr,
                "exit_code": rc,
                "message": msg,
                "outputs": dict(step_outputs),
            },
        }
        for key, expr in step.outputs_map.items():
            step_outputs[key] = cel_engine.eval(expr, eval_ctx)

    return StepResult(
        book_path=str(book_path),
        step=step.name,
        status=status,
        exit_code=rc,
        duration_ms=dt_ms,
        message=msg,
        command=display_redacted,
        cwd=cwd or "",
        when=when_expr or "",
        started_at=started_at,
        ended_at=ended_at,
        stdin_file=eff_stdin or "",
        stdout_file=eff_stdout or "",
        stderr_file=eff_stderr or "",
        attempts=attempts_made,
        max_attempts=max_attempts,
        outputs=step_outputs,
    )


def run_book_sequential(
    book: Book,
    *,
    book_path: Path,
    vars_map: Dict[str, Any],
    opts: RunOptions,
    report_path: Optional[str] = None,
    module_dirs: Optional[List[Path]] = None,
) -> tuple[List[StepResult], Dict[str, Any]]:
    results: List[StepResult] = []
    run_outputs: Dict[str, Any] = {}
    cel = CelEngine.default()
    steps_ctx: Dict[str, Any] = {}
    run_ctx = {
        "id": opts.run_id,
        "book": str(book_path),
        "started_at": _now_iso_utc(),
    }
    if not _maybe_run_hook(
        book,
        "on_run_start",
        book_path=book_path,
        step_name=None,
        step_status=None,
        step_exit=None,
        run_id=opts.run_id,
        report_path=report_path,
        quiet=opts.quiet,
    ):
        return (
            [
                StepResult(
                    book_path=str(book_path),
                    step="__run_start__",
                    status="failed",
                    exit_code=1,
                    duration_ms=0,
                    message="on_run_start hook failed",
                )
            ],
            {},
        )

    for i, step in enumerate(book.steps):
        res = run_one_step(
            book,
            step,
            book_path=book_path,
            vars_map=vars_map,
            opts=opts,
            step_index=i,
            report_path=report_path,
            module_dirs=module_dirs,
            cel=cel,
            inputs_map=vars_map,
            steps_ctx=steps_ctx,
            run_ctx=run_ctx,
        )
        results.append(res)
        steps_ctx[step.name] = {
            "status": res.status,
            "succeeded": res.status in ("success", "failed_nonblocking"),
            "exit_code": res.exit_code,
            "message": res.message,
            "outputs": dict(res.outputs),
            "started_at": res.started_at,
            "ended_at": res.ended_at,
        }
        if res.status == "failed":
            break

    _maybe_run_hook(
        book,
        "on_run_end",
        book_path=book_path,
        step_name=(results[-1].step if results else None),
        step_status=(results[-1].status if results else None),
        step_exit=(results[-1].exit_code if results else None),
        run_id=opts.run_id,
        report_path=report_path,
        quiet=opts.quiet,
    )

    out_defs = book.outputs or {}
    eval_ctx = {
        "inputs": vars_map,
        "steps": steps_ctx,
        "run": run_ctx,
    }
    for name, spec in out_defs.items():
        when_expr = spec.get("when")
        if isinstance(when_expr, str) and when_expr.strip():
            try:
                ok = bool(cel.eval(when_expr, eval_ctx))
            except Exception:
                ok = False
            if not ok:
                continue
        value_expr = spec.get("value")
        if isinstance(value_expr, str) and value_expr.strip():
            try:
                run_outputs[name] = cel.eval(value_expr, eval_ctx)
            except Exception:
                run_outputs[name] = None

    return results, run_outputs
