from __future__ import annotations

import argparse
import json
import os
import sys
from pathlib import Path
from typing import Any, Dict, Optional

from iac.logging import configure_logging
from iac.runbook.discovery import split_module_function
from iac.runbook.module_contract import load_validated_callable
from iac.runbook.module_types import StepContext, StepResult


def _write_result(path: Optional[str], payload: Dict[str, Any]) -> None:
    if not path:
        return
    p = Path(path)
    p.parent.mkdir(parents=True, exist_ok=True)
    p.write_text(json.dumps(payload, indent=2), encoding="utf-8")


def main(argv: Optional[list[str]] = None) -> int:
    parser = argparse.ArgumentParser(prog="iac.module_exec")
    parser.add_argument("--ref", default="", help="module@function")
    parser.add_argument("--book-path", default="")
    parser.add_argument("--step-name", default="")
    parser.add_argument("--run-id", default="")
    parser.add_argument("--vars-json", default="")
    parser.add_argument("--result-file", default="")
    args = parser.parse_args(argv)

    ref = args.ref or os.environ.get("IAC_MODULE_REF", "")
    book_path = args.book_path or os.environ.get("IAC_MODULE_BOOK_PATH", "")
    step_name = args.step_name or os.environ.get("IAC_MODULE_STEP_NAME", "")
    run_id = args.run_id or os.environ.get("IAC_MODULE_RUN_ID", "")
    vars_json = args.vars_json or os.environ.get("IAC_MODULE_VARS_JSON", "{}")
    result_file = args.result_file or os.environ.get("IAC_MODULE_RESULT_FILE", "")

    if not ref.strip():
        sys.stderr.write("missing module ref (IAC_MODULE_REF or --ref)\n")
        return 2
    if not step_name.strip():
        sys.stderr.write("missing step name (IAC_MODULE_STEP_NAME or --step-name)\n")
        return 2
    configure_logging(os.environ.get("IAC_LOG_LEVEL"))
    module_name, fn_name = split_module_function(ref)
    try:
        vars_map = json.loads(vars_json or "{}")
        if not isinstance(vars_map, dict):
            raise ValueError("vars-json must decode to object")
    except Exception as e:
        sys.stderr.write(f"invalid --vars-json: {e}\n")
        return 2

    ctx = StepContext(
        vars=vars_map,
        env=dict(os.environ),
        book_path=book_path,
        step_name=step_name,
        run_id=run_id,
        cwd=os.getcwd(),
        refs={
            "module": module_name,
            "function": fn_name,
            "step_ref": f"{Path(book_path).name}@{step_name}",
        },
    )

    try:
        fn = load_validated_callable(module_name, fn_name)
        out = fn(ctx)
    except Exception as e:
        sys.stderr.write(f"module step failed: {e}\n")
        return 1

    if not isinstance(out, StepResult):
        sys.stderr.write(
            f"module step failed: {module_name}.{fn_name} returned {type(out).__name__}, expected StepResult\n"
        )
        return 1
    payload: Dict[str, Any] = {"message": out.message, "outputs": out.outputs}

    _write_result(result_file or None, payload)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
