from __future__ import annotations

import argparse

from cybervisor import __version__
from cybervisor.adapters.base import AgentAdapter
from cybervisor.adapters.registry import get_adapter
from cybervisor.config import PipelineConfig, StageConfig, load_config
import os
import sys
from pathlib import Path

from cybervisor.env import load_dotenv
from cybervisor.hooks import hook_socket_path, inject_hooks, remove_hooks
from cybervisor.logging import CybervisorLogger, HookEventListener
from cybervisor.pipeline import PipelineInterrupted, PipelineRunner
from cybervisor.preflight import run_preflight
from cybervisor.signals import SignalHandler

RUN_LOG_PATH = ".cybervisor/logs/cybervisor.log.jsonl"


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(prog="cybervisor")
    parser.add_argument("prompt", nargs="?", help="Task description to execute")
    parser.add_argument("--config", default="cybervisor.yaml", help="Path to config YAML")
    parser.add_argument("--start-stage", help="Stage name to skip to and begin execution at")
    parser.add_argument("--end-stage", help="Stage name to stop after")
    parser.add_argument("--end-before", help="Stage name to stop before")
    parser.add_argument("--version", action="store_true", help="Show version and exit")
    return parser


class StartupValidationError(Exception):
    def __init__(self, requested_stage_name: str, reason: str, valid_stage_names: list[str]) -> None:
        super().__init__(reason)
        self.requested_stage_name = requested_stage_name
        self.reason = reason
        self.valid_stage_names = valid_stage_names


def validate_stage_name_arg(requested: str | None, stages: list[StageConfig], flag_name: str) -> str | None:
    valid_names = [stage.name for stage in stages]

    seen = set()
    duplicates = set()
    for name in valid_names:
        if name in seen:
            duplicates.add(name)
        seen.add(name)

    if duplicates:
        raise StartupValidationError(
            requested or "",
            f"Duplicate stage name(s) found in config: {', '.join(sorted(duplicates))}. Cannot safely resolve starting stage.",
            valid_names
        )

    if requested is None:
        return None
    normalized = requested.strip()
    if not normalized:
        raise StartupValidationError(requested, f"{flag_name} must not be empty.", valid_names)
    if normalized not in valid_names:
        raise StartupValidationError(normalized, f"Unknown stage '{normalized}'.", valid_names)
    return normalized


def validate_start_stage(requested: str | None, stages: list[StageConfig]) -> str | None:
    return validate_stage_name_arg(requested, stages, "--start-stage")


def _select_agent(config: PipelineConfig) -> AgentAdapter:
    return get_adapter(config.agent_tool)


def _set_process_title() -> None:
    if sys.platform.startswith("win"):
        return
    try:
        from setproctitle import setproctitle
    except ImportError:
        return
    setproctitle("cybervisor")


def _runtime_dir() -> Path:
    return Path(".cybervisor")


def _process_id_file() -> Path:
    return _runtime_dir() / "process-id"


def _process_group_file() -> Path:
    return _runtime_dir() / "process-group-id"


def _process_group_id() -> int | None:
    if not hasattr(os, "getpgid"):
        return None
    return os.getpgid(os.getpid())


def _write_runtime_files() -> None:
    runtime_dir = _runtime_dir()
    runtime_dir.mkdir(parents=True, exist_ok=True)
    process_group_id = _process_group_id()
    _process_id_file().write_text(f"{os.getpid()}\n", encoding="utf-8")
    if process_group_id is not None:
        _process_group_file().write_text(f"{process_group_id}\n", encoding="utf-8")


def _remove_runtime_files() -> None:
    _process_id_file().unlink(missing_ok=True)
    _process_group_file().unlink(missing_ok=True)


def main(argv: list[str] | None = None) -> int:
    _set_process_title()
    parser = build_parser()
    args = parser.parse_args(argv)

    if args.version:
        print(__version__)
        return 0

    if not args.prompt:
        parser.error("prompt is required")

    # Support demo projects that include a .env file without requiring the user
    # to `export` variables manually. Do not override explicitly-set env vars.
    dotenv = load_dotenv(".env")
    for key, value in dotenv.items():
        if key not in {
            "CYBERVISOR_LLM_API_KEY",
            "CYBERVISOR_LLM_BASE_URL",
            "CYBERVISOR_LLM_MODEL",
        }:
            continue
        if key not in os.environ and value:
            os.environ[key] = value

    logger = CybervisorLogger()

    try:
        config = load_config(args.config)
    except FileNotFoundError as exc:
        logger.log_to_stderr(
            (
                f"Failed to load config: {exc}. "
                "Provide --config with a valid path or create cybervisor.yaml in the working directory."
            ),
            "error",
        )
        return 1
    except Exception as exc:  # pragma: no cover
        logger.log_to_stderr(f"Failed to load config: {exc}", "error")
        return 1

    try:
        start_stage_name = validate_stage_name_arg(args.start_stage, config.stages, "--start-stage")
        end_stage_name = validate_stage_name_arg(args.end_stage, config.stages, "--end-stage")
        end_before_stage_name = validate_stage_name_arg(args.end_before, config.stages, "--end-before")
    except StartupValidationError as exc:
        logger.log_to_stderr(
            f"Invalid stage selection: {exc.reason}\\nValid stages are: {', '.join(exc.valid_stage_names)}",
            "error"
        )
        return 1

    if end_stage_name is not None and end_before_stage_name is not None:
        logger.log_to_stderr(
            "Invalid stage selection: --end-stage and --end-before cannot be used together.",
            "error",
        )
        return 1

    errors = run_preflight(config)
    if errors:
        for error in errors:
            logger.log_to_stderr(error, "error")
        return 1

    agent = _select_agent(config)
    hook_listener: HookEventListener | None = None
    if agent.descriptor.capabilities.supports_hooks:
        hook_listener = HookEventListener(logger, socket_path=hook_socket_path())
        hook_listener.start()

    hooks_installed = False
    if agent.descriptor.capabilities.supports_hooks:
        try:
            inject_hooks(config.agent_tool, config.hook)
            hooks_installed = True
        except Exception as exc:
            if hook_listener is not None:
                hook_listener.stop()
            logger.log_to_stderr(f"Failed to install runtime hooks: {exc}", "error")
            return 1

    signal_handler = SignalHandler(logger, log_file=RUN_LOG_PATH)
    if hooks_installed:
        signal_handler.set_cleanup_callback(lambda: remove_hooks(config.agent_tool))
    signal_handler.register()
    _write_runtime_files()

    try:
        runner = PipelineRunner()
        try:
            ok = runner.run(
                args.prompt,
                config.stages,
                agent,
                logger,
                signal_handler=signal_handler,
                hook_listener=hook_listener,
                start_stage_name=start_stage_name,
                end_stage_name=end_stage_name,
                end_before_stage_name=end_before_stage_name,
            )
        except PipelineInterrupted:
            return 130
        return 0 if ok else 1
    finally:
        if hook_listener is not None:
            hook_listener.stop()
        if hooks_installed:
            try:
                restore_result = remove_hooks(config.agent_tool)
                logger.log_to_json(
                    logger.make_entry(
                        "system",
                        "Cleanup",
                        "Removed runtime hooks",
                        cleanup_result=restore_result,
                    ),
                    RUN_LOG_PATH,
                )
            except Exception as exc:
                logger.log_to_stderr(f"Failed to remove runtime hooks: {exc}", "warning")
        _remove_runtime_files()


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