main.py

"""
Command-line entry point for WorldSmithAI.

This module is intended to live at the repository root beside app.py and
callbacks.py. It provides a local CLI for generating, validating, building,
simulating, visualizing, and narrating WorldSmithAI worlds.

It does not import Gradio and does not assume an app/ package.

Example usage:
    python main.py --example research --steps 60

    python main.py --dsl examples/farm.json --steps 80 --output-dir outputs/farm_run

    python main.py --prompt "A fantasy kingdom with dragons, mages, trade, and mana" --steps 70

    python main.py --dsl examples/research.json --validate-only

    python main.py --prompt "A startup economy with founders and investors" --generate-only

Future extensibility:
    - Add optional model-client backends for prompt generation.
    - Add multi-run experiments.
    - Add benchmark mode.
    - Add artifact bundle export.
    - Add CI smoke-test command.
    - Add run comparison summaries.
"""

from __future__ import annotations

import argparse
import json
import logging
import math
import os
import sys
import traceback
from collections.abc import Mapping, Sequence
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Any

from callbacks import (
    CallbackConfig,
    CallbackResult,
    apply_step_override,
    callback_failure_result,
    generate_world_from_prompt,
    parse_constraints,
    run_worldsmith_pipeline,
    simulate_dsl_pipeline,
)
from dsl.parser import parse_world_spec
from dsl.validator import validate_world_spec


logger = logging.getLogger(__name__)

APP_NAME = "WorldSmithAI"

DEFAULT_PROMPT = (
    "Create a compact research ecosystem where scientists, engineers, reviewers, "
    "curators, and funding agents exchange knowledge, collaborate, compete for "
    "attention, allocate funding, and adapt goals over time."
)

DEFAULT_OUTPUT_ROOT = "outputs"


@dataclass(frozen=True)
class MainRunResult:
    """Structured CLI result.

    Attributes:
        callback_result: Framework callback result.
        output_dir: Directory where CLI artifact files were written.
        written_files: Mapping from artifact label to file path.
        source_label: Human-readable source description.
    """

    callback_result: CallbackResult
    output_dir: Path
    written_files: Mapping[str, str] = field(default_factory=dict)
    source_label: str = "unknown"

    @property
    def success(self) -> bool:
        """Return whether the underlying callback result succeeded."""

        return self.callback_result.success

    def to_dict(self) -> dict[str, Any]:
        """Return a JSON-friendly run summary."""

        return {
            "success": self.success,
            "source_label": self.source_label,
            "output_dir": str(self.output_dir),
            "written_files": dict(self.written_files),
            "callback_result": self.callback_result.to_dict(),
        }


class MainError(RuntimeError):
    """Raised when CLI arguments or execution are invalid."""


def main(argv: Sequence[str] | None = None) -> int:
    """Run the WorldSmithAI CLI.

    Args:
        argv: Optional command-line arguments. When omitted, sys.argv is used.

    Returns:
        Process exit code.
    """

    parser = build_arg_parser()
    args = parser.parse_args(argv)

    configure_logging(args.log_level)

    try:
        if args.list_examples:
            print_examples(args.examples_dir)
            return 0

        result = run_from_args(args)

        if not args.quiet:
            print_console_summary(result)

        if args.print_summary:
            print(json.dumps(_json_safe(result.to_dict()), indent=2, sort_keys=True))

        return 0 if result.success else 2

    except KeyboardInterrupt:
        print("Interrupted.", file=sys.stderr)
        return 130

    except Exception as exc:
        logger.exception("WorldSmithAI CLI failed")
        failure = callback_failure_result(exc)
        if not args.quiet:
            print(failure.status_message, file=sys.stderr)
            if args.debug:
                print(traceback.format_exc(), file=sys.stderr)
        return 1


def build_arg_parser() -> argparse.ArgumentParser:
    """Build the CLI argument parser."""

    parser = argparse.ArgumentParser(
        prog="worldsmithai",
        description="Run WorldSmithAI from the command line.",
    )

    input_group = parser.add_mutually_exclusive_group()
    input_group.add_argument(
        "--prompt",
        "-p",
        type=str,
        default=None,
        help="Natural-language world prompt. Uses deterministic fallback generation unless a client is wired externally.",
    )
    input_group.add_argument(
        "--dsl",
        type=str,
        default=None,
        help="Path to a WorldSpec JSON file, or '-' to read from stdin.",
    )
    input_group.add_argument(
        "--example",
        "-e",
        type=str,
        default=None,
        help="Example name or path. Examples include farm, civilization, and research.",
    )

    parser.add_argument(
        "--list-examples",
        action="store_true",
        help="List available examples and exit.",
    )
    parser.add_argument(
        "--examples-dir",
        type=str,
        default="examples",
        help="Directory containing example JSON files.",
    )
    parser.add_argument(
        "--constraints",
        "-c",
        type=str,
        default=None,
        help="Optional constraints as JSON text, plain text, or @path/to/file.",
    )
    parser.add_argument(
        "--steps",
        "-s",
        type=int,
        default=None,
        help="Override simulation steps.",
    )
    parser.add_argument(
        "--output-dir",
        "-o",
        type=str,
        default=None,
        help="Directory where output artifacts should be written.",
    )
    parser.add_argument(
        "--animation-format",
        choices=("gif", "mp4"),
        default="gif",
        help="Animation format. GIF is safest for local and Spaces demos.",
    )
    parser.add_argument(
        "--max-animation-frames",
        type=int,
        default=int(os.getenv("WORLDSMITHAI_MAX_FRAMES", "80")),
        help="Maximum number of animation frames to capture.",
    )

    parser.add_argument(
        "--generate-only",
        action="store_true",
        help="Generate DSL from --prompt and write JSON without running simulation.",
    )
    parser.add_argument(
        "--validate-only",
        action="store_true",
        help="Validate DSL input without building or simulating the world.",
    )
    parser.add_argument(
        "--strict",
        action="store_true",
        help="Treat semantic validation issues as errors where configured.",
    )

    parser.add_argument(
        "--no-animation",
        action="store_true",
        help="Do not generate animation artifacts.",
    )
    parser.add_argument(
        "--no-charts",
        action="store_true",
        help="Do not generate population or resource charts.",
    )
    parser.add_argument(
        "--no-final-image",
        action="store_true",
        help="Do not generate final world-state image.",
    )
    parser.add_argument(
        "--no-narrative",
        action="store_true",
        help="Do not generate narrative summary.",
    )
    parser.add_argument(
        "--no-metric-history",
        action="store_true",
        help="Do not collect per-step metric history.",
    )

    parser.add_argument(
        "--print-summary",
        action="store_true",
        help="Print full JSON run summary to stdout.",
    )
    parser.add_argument(
        "--quiet",
        "-q",
        action="store_true",
        help="Suppress human-readable console output.",
    )
    parser.add_argument(
        "--debug",
        action="store_true",
        help="Print tracebacks on failure.",
    )
    parser.add_argument(
        "--log-level",
        type=str,
        default=os.getenv("WORLDSMITHAI_LOG_LEVEL", "INFO"),
        help="Logging level.",
    )

    return parser


def run_from_args(args: argparse.Namespace) -> MainRunResult:
    """Run the requested CLI action from parsed arguments."""

    output_dir = resolve_output_dir(args.output_dir)
    config = build_callback_config(args, output_dir=output_dir)

    if args.generate_only:
        return run_generate_only(args, config=config, output_dir=output_dir)

    if args.validate_only:
        return run_validate_only(args, config=config, output_dir=output_dir)

    source_kind, source_value, source_label = resolve_world_source(args)

    if source_kind == "prompt":
        callback_result = run_worldsmith_pipeline(
            prompt=source_value,
            constraints=parse_constraints(read_constraints_argument(args.constraints)),
            client=None,
            config=config,
        )
    elif source_kind == "dsl":
        callback_result = simulate_dsl_pipeline(
            dsl_json=source_value,
            client=None,
            config=config,
        )
    else:
        raise MainError(f"Unsupported source kind: {source_kind}")

    written_files = write_callback_result_files(
        callback_result,
        output_dir=output_dir,
        source_label=source_label,
    )

    return MainRunResult(
        callback_result=callback_result,
        output_dir=output_dir,
        written_files=written_files,
        source_label=source_label,
    )


def run_generate_only(
    args: argparse.Namespace,
    *,
    config: CallbackConfig,
    output_dir: Path,
) -> MainRunResult:
    """Generate a WorldSpec from prompt and write it without simulation."""

    prompt = args.prompt or DEFAULT_PROMPT
    constraints = parse_constraints(read_constraints_argument(args.constraints))

    generation = generate_world_from_prompt(
        prompt=prompt,
        constraints=constraints,
        client=None,
        config=config,
    )
    spec = apply_step_override(generation.spec, config.steps)
    report = validate_world_spec(spec, config=config.resolved_validation_config())

    status = (
        f"Generated WorldSpec {spec.id!r}: "
        f"{len(spec.agents)} agent(s), {len(spec.resources)} resource(s), "
        f"{len(spec.behavior_names)} behavior reference(s). "
        f"Validation: {len(report.errors)} error(s), {len(report.warnings)} warning(s)."
    )

    callback_result = CallbackResult(
        success=report.is_valid,
        status_message=status,
        world_spec_json=spec.to_json_string(indent=2, exclude_none=True),
        validation_json=safe_json_dumps(report.to_dict()),
        metrics_json=safe_json_dumps(
            {
                "success": report.is_valid,
                "generation": generation.to_dict(),
                "validation": report.to_dict(),
            }
        ),
        narrative="Generation-only mode completed. No simulation was run.",
        metadata={
            "mode": "generate_only",
            "prompt": prompt,
        },
    )

    written_files = write_callback_result_files(
        callback_result,
        output_dir=output_dir,
        source_label="prompt:generate_only",
    )

    return MainRunResult(
        callback_result=callback_result,
        output_dir=output_dir,
        written_files=written_files,
        source_label="prompt:generate_only",
    )


def run_validate_only(
    args: argparse.Namespace,
    *,
    config: CallbackConfig,
    output_dir: Path,
) -> MainRunResult:
    """Validate a DSL source and write validation report without simulation."""

    source_kind, source_value, source_label = resolve_world_source(args)

    if source_kind == "prompt":
        raise MainError("--validate-only requires --dsl or --example, not --prompt")

    spec = parse_world_spec(source_value)
    spec = apply_step_override(spec, config.steps)
    report = validate_world_spec(spec, config=config.resolved_validation_config())

    status = (
        f"Validated WorldSpec {spec.id!r}. "
        f"Validation: {len(report.errors)} error(s), {len(report.warnings)} warning(s)."
    )

    callback_result = CallbackResult(
        success=report.is_valid,
        status_message=status,
        world_spec_json=spec.to_json_string(indent=2, exclude_none=True),
        validation_json=safe_json_dumps(report.to_dict()),
        metrics_json=safe_json_dumps({"success": report.is_valid, "validation": report.to_dict()}),
        narrative="Validation-only mode completed. No simulation was run.",
        metadata={
            "mode": "validate_only",
            "source_label": source_label,
        },
    )

    written_files = write_callback_result_files(
        callback_result,
        output_dir=output_dir,
        source_label=source_label,
    )

    return MainRunResult(
        callback_result=callback_result,
        output_dir=output_dir,
        written_files=written_files,
        source_label=source_label,
    )


def build_callback_config(args: argparse.Namespace, *, output_dir: Path) -> CallbackConfig:
    """Build callback configuration from CLI args."""

    return CallbackConfig(
        output_dir=output_dir,
        steps=args.steps,
        max_animation_frames=max(1, int(args.max_animation_frames)),
        animation_format=args.animation_format,
        generate_animation=not bool(args.no_animation),
        generate_charts=not bool(args.no_charts),
        generate_final_image=not bool(args.no_final_image),
        generate_narrative=not bool(args.no_narrative),
        collect_metric_history=not bool(args.no_metric_history),
        strict_validation_for_app=bool(args.strict),
    )


def resolve_world_source(args: argparse.Namespace) -> tuple[str, str, str]:
    """Resolve CLI input into a source kind, source text, and source label."""

    if args.prompt:
        return "prompt", str(args.prompt).strip(), "prompt"

    if args.dsl:
        dsl_text = read_text_source(args.dsl)
        return "dsl", dsl_text, f"dsl:{args.dsl}"

    if args.example:
        example_path = resolve_example_path(args.example, examples_dir=args.examples_dir)
        return "dsl", example_path.read_text(encoding="utf-8"), f"example:{example_path.name}"

    default_example = resolve_default_example(examples_dir=args.examples_dir)
    if default_example is not None:
        return "dsl", default_example.read_text(encoding="utf-8"), f"example:{default_example.name}"

    return "prompt", DEFAULT_PROMPT, "default_prompt"


def discover_examples(examples_dir: str | Path = "examples") -> dict[str, Path]:
    """Discover available example JSON files."""

    directory = Path(examples_dir)
    if not directory.exists():
        return {}

    examples: dict[str, Path] = {}

    for path in sorted(directory.glob("*.json")):
        examples[path.stem.lower()] = path
        examples[path.stem.replace("_", " ").lower()] = path
        examples[path.name.lower()] = path

    return examples


def resolve_example_path(example: str, *, examples_dir: str | Path = "examples") -> Path:
    """Resolve an example name or path to a JSON file path."""

    candidate = Path(example)
    if candidate.exists():
        return candidate

    if candidate.suffix != ".json":
        json_candidate = candidate.with_suffix(".json")
        if json_candidate.exists():
            return json_candidate

    examples = discover_examples(examples_dir)
    key = example.strip().lower()
    if key in examples:
        return examples[key]

    simplified_key = key.replace("_", " ")
    if simplified_key in examples:
        return examples[simplified_key]

    known = sorted(set(str(path) for path in examples.values()))
    raise MainError(f"Unknown example {example!r}. Known example files: {known}")


def resolve_default_example(*, examples_dir: str | Path = "examples") -> Path | None:
    """Return the preferred default example path, if available."""

    examples = discover_examples(examples_dir)
    for key in ("research", "farm", "civilization"):
        if key in examples:
            return examples[key]

    unique_paths = sorted(set(examples.values()))
    return unique_paths[0] if unique_paths else None


def print_examples(examples_dir: str | Path = "examples") -> None:
    """Print available example worlds."""

    examples = discover_examples(examples_dir)
    unique_paths = sorted(set(examples.values()))

    if not unique_paths:
        print(f"No examples found in {examples_dir!s}.")
        return

    print("Available examples:")
    for path in unique_paths:
        print(f"  - {path.stem}: {path}")


def read_text_source(source: str) -> str:
    """Read DSL text from a path or stdin marker."""

    if source == "-":
        return sys.stdin.read()

    path = Path(source)
    if not path.exists():
        raise MainError(f"DSL file not found: {source}")

    return path.read_text(encoding="utf-8")


def read_constraints_argument(value: str | None) -> str | Mapping[str, Any] | None:
    """Read constraints from CLI text or @file syntax."""

    if value is None:
        return None

    text = str(value).strip()
    if not text:
        return None

    if text.startswith("@"):
        path = Path(text[1:])
        if not path.exists():
            raise MainError(f"Constraints file not found: {path}")
        return path.read_text(encoding="utf-8")

    return text


def resolve_output_dir(output_dir_arg: str | None) -> Path:
    """Resolve and create the CLI output directory."""

    if output_dir_arg:
        output_dir = Path(output_dir_arg)
    elif os.getenv("WORLDSMITHAI_OUTPUT_DIR"):
        output_dir = Path(os.environ["WORLDSMITHAI_OUTPUT_DIR"])
    else:
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        output_dir = Path(DEFAULT_OUTPUT_ROOT) / f"run_{timestamp}"

    output_dir.mkdir(parents=True, exist_ok=True)
    return output_dir


def write_callback_result_files(
    result: CallbackResult,
    *,
    output_dir: Path,
    source_label: str,
) -> dict[str, str]:
    """Write standard JSON, Markdown, and summary files for a callback result."""

    output_dir.mkdir(parents=True, exist_ok=True)

    written: dict[str, str] = {}

    written["world_spec_json"] = write_text_file(
        output_dir / "world_spec.json",
        result.world_spec_json or "{}",
    )
    written["validation_report_json"] = write_text_file(
        output_dir / "validation_report.json",
        result.validation_json or "{}",
    )
    written["metrics_json"] = write_text_file(
        output_dir / "metrics.json",
        result.metrics_json or "{}",
    )
    written["narrative_md"] = write_text_file(
        output_dir / "narrative.md",
        result.narrative or "No narrative generated.",
    )

    summary_payload = {
        "app": APP_NAME,
        "source_label": source_label,
        "success": result.success,
        "status_message": result.status_message,
        "standard_files": written,
        "artifacts": None if result.artifacts is None else result.artifacts.to_dict(),
        "callback_result": result.to_dict(),
    }

    written["run_summary_json"] = write_text_file(
        output_dir / "run_summary.json",
        safe_json_dumps(summary_payload),
    )

    return written


def write_text_file(path: Path, text: str) -> str:
    """Write UTF-8 text to a file and return the path string."""

    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(text, encoding="utf-8")
    return str(path)


def print_console_summary(result: MainRunResult) -> None:
    """Print a concise human-readable run summary."""

    callback_result = result.callback_result

    print()
    print(f"{APP_NAME} run summary")
    print("=" * 72)
    print(f"Success: {callback_result.success}")
    print(f"Source: {result.source_label}")
    print(f"Output directory: {result.output_dir}")
    print(f"Status: {callback_result.status_message}")

    print()
    print("Written files:")
    for label, path in sorted(result.written_files.items()):
        print(f"  - {label}: {path}")

    if callback_result.artifacts is not None:
        artifacts = callback_result.artifacts.to_dict()
        artifact_items = [(key, value) for key, value in artifacts.items() if value and key != "output_dir"]
        if artifact_items:
            print()
            print("Generated artifacts:")
            for label, path in artifact_items:
                print(f"  - {label}: {path}")

    print()


def configure_logging(level: str) -> None:
    """Configure logging for CLI execution."""

    logging.basicConfig(
        level=str(level).upper(),
        format="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
    )


def safe_json_dumps(value: Any) -> str:
    """Serialize a value as pretty JSON."""

    return json.dumps(
        _json_safe(value),
        indent=2,
        sort_keys=True,
        ensure_ascii=False,
    )


def _json_safe(value: Any) -> Any:
    """Return a JSON-friendly representation of arbitrary values."""

    if value is None or isinstance(value, (str, bool)):
        return value

    if isinstance(value, int) and not isinstance(value, bool):
        return value

    if isinstance(value, float):
        if not math.isfinite(value):
            return None
        return value

    if isinstance(value, Mapping):
        return {str(key): _json_safe(nested) for key, nested in value.items()}

    if isinstance(value, Sequence) and not isinstance(value, (str, bytes)):
        return [_json_safe(item) for item in value]

    if hasattr(value, "to_dict") and callable(value.to_dict):
        return _json_safe(value.to_dict())

    if hasattr(value, "model_dump") and callable(value.model_dump):
        return _json_safe(value.model_dump(mode="json"))

    return str(value)


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