from __future__ import annotations

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

import typer

from iac.cli.formatting import ensure_output_format, emit_structured, build_table
from iac.utils.subproc import ControlledProcess

app = typer.Typer(
    add_completion=False, no_args_is_help=True, help="Run one-off commands"
)

_SHELL_META = ("|", "&&", "||", ";", ">", "<", "$(", "`", "\n")


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


def _default_shell_argv(cmd_str: str) -> List[str]:
    return ["bash", "-lc", cmd_str]


def _looks_like_shell_string(s: str) -> bool:
    return any(tok in s for tok in _SHELL_META)


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


def _open_out_file(path: Optional[str], *, mode: str):
    if not path:
        return None
    _ensure_parent(path)
    file_mode = "ab" if mode == "append" else "wb"
    return open(path, file_mode)


def _load_stdin_bytes(
    *,
    stdin_file: Optional[str],
    stdin_text: Optional[str],
    stdin_json: Optional[str],
    stdin_env: Optional[str],
) -> Optional[bytes]:
    methods = [
        stdin_file is not None,
        stdin_text is not None,
        stdin_json is not None,
        stdin_env is not None,
    ]
    if sum(1 for m in methods if m) > 1:
        raise typer.BadParameter(
            "Use at most one of --stdin-file, --stdin-text, --stdin-json, or --stdin-env."
        )

    if stdin_file is not None:
        p = Path(stdin_file).expanduser()
        if not p.exists():
            raise typer.BadParameter(f"--stdin-file not found: {stdin_file}")
        return p.read_bytes()
    if stdin_text is not None:
        return stdin_text.encode("utf-8")
    if stdin_json is not None:
        try:
            obj = json.loads(stdin_json)
        except Exception as e:
            raise typer.BadParameter(f"Invalid --stdin-json payload: {e}") from e
        return json.dumps(obj).encode("utf-8")
    if stdin_env is not None:
        val = os.environ.get(stdin_env)
        if val is None:
            raise typer.BadParameter(f"--stdin-env variable is not set: {stdin_env}")
        return val.encode("utf-8")
    return None


@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


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)


@app.command("cmd")
def cmd_(
    cmd: Annotated[
        str,
        typer.Argument(
            help="Command string. In --no-shell mode this is the executable."
        ),
    ],
    arg: Annotated[
        List[str],
        typer.Option(
            "--arg", "-a", help="Argument token for --no-shell mode (repeatable)"
        ),
    ] = [],
    shell: Annotated[
        bool,
        typer.Option(
            "--shell/--no-shell",
            help="Shell string (bash -lc) vs direct argv execution.",
        ),
    ] = True,
    cwd: Annotated[
        Optional[str], typer.Option("--cwd", help="Working directory")
    ] = None,
    env: Annotated[
        List[str],
        typer.Option("--env", help="Environment override(s), KEY=VALUE. Repeatable."),
    ] = [],
    stdin_file: Annotated[
        Optional[str], typer.Option("--stdin-file", help="Send stdin bytes from file")
    ] = None,
    stdin_text: Annotated[
        Optional[str], typer.Option("--stdin-text", help="Send UTF-8 text to stdin")
    ] = None,
    stdin_json: Annotated[
        Optional[str],
        typer.Option(
            "--stdin-json", help="Validate JSON then send compact JSON to stdin"
        ),
    ] = None,
    stdin_env: Annotated[
        Optional[str],
        typer.Option(
            "--stdin-env", help="Read stdin content from environment variable value"
        ),
    ] = None,
    clean_env: Annotated[
        bool, typer.Option("--clean-env", help="Do not inherit parent env")
    ] = False,
    pty: Annotated[bool, typer.Option("--pty", help="Run under a PTY")] = False,
    timeout: Annotated[
        Optional[float], typer.Option("--timeout", help="Timeout seconds")
    ] = None,
    killpg: Annotated[
        bool, typer.Option("--killpg/--no-killpg", help="Terminate process group")
    ] = True,
    forward_signals: Annotated[
        bool,
        typer.Option(
            "--forward-signals/--no-forward-signals",
            help="Forward SIGINT/SIGTERM(/SIGHUP)",
        ),
    ] = True,
    capture: Annotated[
        bool, typer.Option("--capture/--no-capture", help="Capture output in memory")
    ] = False,
    stdout_file: Annotated[
        Optional[str], typer.Option("--stdout-file", help="Append stdout bytes to file")
    ] = None,
    stderr_file: Annotated[
        Optional[str], typer.Option("--stderr-file", help="Append stderr bytes to file")
    ] = None,
    file_mode: Annotated[
        Literal["append", "truncate"],
        typer.Option(
            "--file-mode", help="File mode for stdout/stderr files: append|truncate"
        ),
    ] = "append",
    output: Annotated[
        str,
        typer.Option(
            "--output", help="Output format: table|json|yaml", show_default=True
        ),
    ] = "table",
    print_captured: Annotated[
        bool,
        typer.Option(
            "--print-captured/--no-print-captured",
            help="If --capture, print captured output at end",
        ),
    ] = False,
    quiet: Annotated[
        bool, typer.Option("-q", "--quiet", help="Do not stream output to console")
    ] = False,
    verbose: Annotated[
        bool, typer.Option("-v", "--verbose", help="Print extra metadata")
    ] = False,
) -> None:
    """Run a one-off command with shell/argv modes, stdin helpers, and output controls."""
    output_fmt = ensure_output_format(output)
    started_at = dt.datetime.now(dt.timezone.utc).isoformat()
    env_map = _parse_env_kv(env)
    stdin_bytes = _load_stdin_bytes(
        stdin_file=stdin_file,
        stdin_text=stdin_text,
        stdin_json=stdin_json,
        stdin_env=stdin_env,
    )

    if shell:
        if arg:
            raise typer.BadParameter(
                "In --shell mode, do not pass --arg values; include args in CMD."
            )
        proc_argv = _default_shell_argv(cmd)
        display = cmd
    else:
        proc_argv = [cmd] + arg
        display = shlex.join(proc_argv)
        if len(proc_argv) == 1 and _looks_like_shell_string(proc_argv[0]):
            raise typer.BadParameter(
                "It looks like you passed a shell expression while using --no-shell.\n"
                "Use: iac run cmd 'echo hi | wc -c'\n"
                "Or:  iac run cmd --no-shell echo --arg hi\n"
            )

    sinks = _CaptureSinks(
        stdout_fh=_open_out_file(stdout_file, mode=file_mode),
        stderr_fh=_open_out_file(stderr_file, mode=file_mode),
    )

    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(
        proc_argv,
        cwd=cwd,
        env=env_map if env_map else None,
        inherit_env=not clean_env,
        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)

    start_t = time.time()
    exit_code = 1
    message = ""

    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 stdin_bytes is not None:
            p.write(stdin_bytes)
            p.close_stdin()
        if verbose:
            typer.echo(f"[pid={p.pid} pgid={p.pgid}] {display}", err=True)
        forwarder.install()

        while True:
            rc = p.poll()
            if rc is not None:
                exit_code = int(rc)
                break
            if timeout is not None and (time.time() - start_t) > timeout:
                typer.echo(
                    f"Command timed out after {timeout} seconds; stopping...", err=True
                )
                p.stop(grace=2.0, use_process_group=killpg)
                exit_code = 124
                message = "timeout"
                break
            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
                exit_code = _exit_code_for_signal(sig)
                message = f"terminated by signal {sig}"
                break
            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)
        exit_code = 127
        message = f"executable not found: {missing}"

    except KeyboardInterrupt:
        typer.echo("Interrupted; stopping child process...", err=True)
        try:
            p.stop(grace=2.0, use_process_group=killpg)
        finally:
            exit_code = 130
            message = "keyboard interrupt"

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

    ended_at = dt.datetime.now(dt.timezone.utc).isoformat()
    rc_final = (
        exit_code
        if exit_code is not None
        else (p.returncode if p.returncode is not None else 1)
    )
    duration_ms = int((time.time() - start_t) * 1000)

    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()

    payload = {
        "kind": "cmd",
        "started_at": started_at,
        "ended_at": ended_at,
        "duration_ms": duration_ms,
        "exit_code": int(rc_final),
        "status": "success" if int(rc_final) == 0 else "failed",
        "message": message,
        "shell": shell,
        "command": display,
        "argv": proc_argv,
        "cwd": cwd or "",
        "stdin_source": (
            "file"
            if stdin_file is not None
            else (
                "text"
                if stdin_text is not None
                else (
                    "json"
                    if stdin_json is not None
                    else "env" if stdin_env is not None else ""
                )
            )
        ),
        "stdin_bytes": len(stdin_bytes) if stdin_bytes is not None else 0,
        "stdout_file": stdout_file or "",
        "stderr_file": stderr_file or "",
        "file_mode": file_mode,
    }

    if output_fmt in ("json", "yaml"):
        emit_structured(payload, output_fmt)
        raise typer.Exit(code=int(rc_final))

    table = build_table(
        [
            "STATUS",
            "EXIT",
            "DURATION_MS",
            "COMMAND",
            "CWD",
            "STDIN",
            "STDOUT_FILE",
            "STDERR_FILE",
            "MESSAGE",
        ],
        [
            [
                payload["status"],
                payload["exit_code"],
                payload["duration_ms"],
                payload["command"],
                payload["cwd"],
                payload["stdin_source"],
                payload["stdout_file"],
                payload["stderr_file"],
                payload["message"],
            ]
        ],
    )
    typer.echo("Run Cmd")
    typer.echo(table.get_string())
    raise typer.Exit(code=int(rc_final))
