"""CLI for Motia framework."""

import argparse
import glob
import importlib.util
import json
import logging
import os
import sys
from collections.abc import Mapping
from pathlib import Path
from typing import Any

from dotenv import load_dotenv

log = logging.getLogger("motia.cli")


def load_module_from_path(path: str) -> Any | None:
    """Load a Python module from a file path and return it."""
    spec = importlib.util.spec_from_file_location("step_module", path)
    if spec and spec.loader:
        module = importlib.util.module_from_spec(spec)
        sys.modules[f"step_module:{path}"] = module
        spec.loader.exec_module(module)
        return module
    return None


def load_and_register_step(path: str) -> None:
    """Load a step module and register using Motia.add_step()."""
    from .runtime import Motia

    motia = Motia()
    module = load_module_from_path(path)
    if module is None:
        return

    config = getattr(module, "config", None)
    handler = getattr(module, "handler", None)

    if callable(handler) and (isinstance(config, Mapping) or hasattr(config, "model_dump")):
        motia.add_step(config, path, handler)


def _discover_files(directories: list[str], pattern: str) -> list[str]:
    files: list[str] = []
    for directory in directories:
        if not os.path.exists(directory):
            continue
        files.extend(glob.glob(os.path.join(directory, "**", pattern), recursive=True))
    return files


def discover_steps(directory: str, *, include_src: bool = False) -> list[str]:
    """Discover step files in a directory."""
    directories = [directory]
    if include_src and directory != "src":
        directories.append("src")
    return _discover_files(directories, "*_step.py")


def discover_streams(directory: str, *, include_src: bool = False) -> list[str]:
    """Discover stream files in a directory."""
    directories = [directory]
    if include_src and directory != "src":
        directories.append("src")
    return _discover_files(directories, "*_stream.py")


def configure_logging(verbose: bool = False) -> None:
    """Configure logging for the CLI."""
    level = logging.DEBUG if verbose else logging.INFO
    format = "%(asctime)s [%(levelname)s] %(name)s: %(message)s"
    logging.basicConfig(level=level, format=format)


def generate_index(step_files: list[str], stream_files: list[str]) -> str:
    """Generate a Python entrypoint that loads all steps and streams."""
    lines = [
        "# Auto-generated by motia build",
        "import asyncio",
        "import importlib.util",
        "import sys",
        "from collections.abc import Mapping",
        "from dotenv import load_dotenv",
        "from motia.runtime import Motia",
        "from motia.iii import get_instance",
        "",
        "load_dotenv()",
        "",
        "def load_module_from_path(path: str) -> None:",
        "    spec = importlib.util.spec_from_file_location('step_module', path)",
        "    if spec and spec.loader:",
        "        module = importlib.util.module_from_spec(spec)",
        "        sys.modules[f'step_module:{path}'] = module",
        "        spec.loader.exec_module(module)",
        "",
        "motia = Motia()",
        "",
        "def load_and_register_step(path: str) -> None:",
        "    spec = importlib.util.spec_from_file_location('step_module', path)",
        "    if spec and spec.loader:",
        "        module = importlib.util.module_from_spec(spec)",
        "        sys.modules[f'step_module:{path}'] = module",
        "        spec.loader.exec_module(module)",
        "        config = getattr(module, 'config', None)",
        "        handler = getattr(module, 'handler', None)",
        "        if callable(handler) and (isinstance(config, Mapping) or hasattr(config, 'model_dump')):",
        "            motia.add_step(config, path, handler)",
        "",
        "def main() -> None:",
    ]
    for stream_file in stream_files:
        lines.append(f"    load_module_from_path(r'{stream_file}')")
    for step_file in step_files:
        lines.append(f"    load_and_register_step(r'{step_file}')")
    lines.extend(
        [
            "    iii = get_instance()",
            "",
            "    async def run() -> None:",
            "        await iii.connect()",
            "        while True:",
            "            await asyncio.sleep(1)",
            "",
            "    asyncio.run(run())",
            "",
            "if __name__ == '__main__':",
            "    main()",
        ]
    )
    return "\n".join(lines)


def generate_schema_manifest(step_files: list[str], stream_files: list[str]) -> dict[str, Any]:
    """Generate a schema manifest from discovered steps and streams."""
    manifest: dict[str, Any] = {"steps": [], "streams": []}

    for file_path in stream_files + step_files:
        load_module_from_path(file_path)

    for file_path in step_files:
        module = sys.modules.get(f"step_module:{file_path}")
        if not module or not hasattr(module, "config"):
            continue
        config = getattr(module, "config")
        try:
            config_data = config.model_dump(by_alias=True) if hasattr(config, "model_dump") else config
        except Exception:
            config_data = config
        manifest["steps"].append({"filePath": file_path, "config": config_data})

    for file_path in stream_files:
        module = sys.modules.get(f"step_module:{file_path}")
        if not module or not hasattr(module, "config"):
            continue
        config = getattr(module, "config")
        try:
            config_data = config.model_dump(by_alias=True) if hasattr(config, "model_dump") else config
        except Exception:
            config_data = config
        manifest["streams"].append({"filePath": file_path, "config": config_data})

    return manifest


def main() -> None:
    """Main entry point for the CLI."""
    load_dotenv()
    parser = argparse.ArgumentParser(description="Motia CLI")
    parser.add_argument(
        "command",
        choices=["run", "dev", "build", "typegen"],
        help="Command to run",
    )
    parser.add_argument(
        "--dir",
        "-d",
        default="steps",
        help="Directory containing step files",
    )
    parser.add_argument(
        "--verbose",
        "-v",
        action="store_true",
        help="Enable verbose logging",
    )
    parser.add_argument(
        "--watch",
        "-w",
        action="store_true",
        help="Watch for file changes",
    )
    parser.add_argument(
        "--output",
        "-o",
        default="types.json",
        help="Output file for typegen",
    )

    args = parser.parse_args()
    configure_logging(args.verbose)

    if args.command == "run":
        step_files = discover_steps(args.dir, include_src=True)
        stream_files = discover_streams(args.dir, include_src=True)

        log.info("Found %d step(s) and %d stream(s)", len(step_files), len(stream_files))

        for stream_file in stream_files:
            log.debug("Loading stream: %s", stream_file)
            load_module_from_path(stream_file)

        for step_file in step_files:
            log.debug("Loading step: %s", step_file)
            load_and_register_step(step_file)

        log.info("All steps loaded. Connecting...")

        try:
            import asyncio

            from .iii import get_instance

            iii = get_instance()

            async def run() -> None:
                await iii.connect()
                log.info("Connected. Waiting for events...")
                while True:
                    await asyncio.sleep(1)

            asyncio.run(run())
        except KeyboardInterrupt:
            log.info("Shutting down...")
    elif args.command == "dev":
        step_files = discover_steps(args.dir, include_src=True)
        stream_files = discover_streams(args.dir, include_src=True)

        if args.watch:
            try:
                from watchfiles import watch

                log.info("Watching for changes...")
                for _ in watch(*step_files, *stream_files):
                    os.execv(sys.executable, [sys.executable] + sys.argv)
            except Exception:
                log.warning("watchfiles not available; running without watch mode")

        for stream_file in stream_files:
            load_module_from_path(stream_file)
        for step_file in step_files:
            load_and_register_step(step_file)

        try:
            import asyncio

            from .iii import get_instance

            iii = get_instance()

            async def run() -> None:
                await iii.connect()
                log.info("Connected. Waiting for events...")
                while True:
                    await asyncio.sleep(1)

            asyncio.run(run())
        except KeyboardInterrupt:
            log.info("Shutting down...")
    elif args.command == "build":
        step_files = discover_steps(args.dir, include_src=True)
        stream_files = discover_streams(args.dir, include_src=True)
        index_content = generate_index(step_files, stream_files)
        Path("dist").mkdir(exist_ok=True)
        output_path = Path("dist") / "index.py"
        output_path.write_text(index_content)
        log.info("Build completed: dist/index.py")
    elif args.command == "typegen":
        step_files = discover_steps(args.dir, include_src=True)
        stream_files = discover_streams(args.dir, include_src=True)
        manifest = generate_schema_manifest(step_files, stream_files)
        Path(args.output).write_text(json.dumps(manifest, indent=2))
        log.info("Typegen output written: %s", args.output)


if __name__ == "__main__":
    main()
