"""Optional nanobot feature discovery and enablement."""
from __future__ import annotations

import json
import os
import shutil
import subprocess
import sys
from dataclasses import dataclass
from importlib.metadata import PackageNotFoundError, distribution
from pathlib import Path
from typing import Any, cast

from loguru import logger
from packaging.requirements import Requirement
from packaging.utils import canonicalize_name

from nanobot.channels._setup import channel_setup_spec
from nanobot.channels.contracts import (
    ChannelSetupSpec,
    channel_feature_instances,
    channel_field_value,
    channel_instance_specs,
    channel_local_state_present,
    channel_set_config_enabled,
    channel_value_present,
    refresh_channel_feature_metadata,
    resolve_channel_action_target,
    stringify_channel_value,
)
from nanobot.channels.registry import channel_default_enabled
from nanobot.config.schema import Config


class OptionalFeatureError(Exception):
    def __init__(self, message: str, *, status: int = 400) -> None:
        super().__init__(message)
        self.message = message
        self.status = status


@dataclass
class InstallResult:
    ok: bool
    label: str
    pip_cmd: list[str]
    failed_cmd: list[str] | None = None
    output: str = ""


_INSTALL_TIMEOUT_SECONDS = 300
_LOG_OUTPUT_LIMIT = 4000
_HIDDEN_OPTIONAL_FEATURES = {"documents", "pdf"}
_BUNDLED_FEATURE_ALIASES = {"documents", "pdf"}


def load_pyproject(path: Path) -> dict[str, Any]:
    import tomllib

    try:
        content = path.read_text(encoding="utf-8")
    except FileNotFoundError:
        return {}
    return tomllib.loads(content)


def optional_dependency_groups_from_metadata() -> dict[str, list[str] | None]:
    from importlib.metadata import metadata, requires

    try:
        extras = metadata("nanobot-ai").get_all("Provides-Extra") or []
        raw_requirements = requires("nanobot-ai") or []
    except PackageNotFoundError:
        return {}
    groups: dict[str, list[str] | None] = {name: [] for name in extras if name != "dev"}
    for raw in raw_requirements:
        req = Requirement(raw)
        if not req.marker:
            continue
        for extra, deps in groups.items():
            if deps is not None and req.marker.evaluate({"extra": extra}):
                deps.append(raw)
    return groups


def optional_dependency_groups() -> dict[str, list[str] | None]:
    root = Path(__file__).resolve().parents[1]
    project = load_pyproject(root / "pyproject.toml").get("project", {})
    deps = project.get("optional-dependencies", {})
    if isinstance(deps, dict) and deps:
        return {
            name: list(cast(list[str], values))
            for name, values in cast(dict[str, object], deps).items()
            if name != "dev" and name not in _HIDDEN_OPTIONAL_FEATURES and isinstance(values, list)
        }
    return {
        name: values
        for name, values in optional_dependency_groups_from_metadata().items()
        if name not in _HIDDEN_OPTIONAL_FEATURES
    }


def _install_requirements_for_extra(extra: str, deps: list[str]) -> list[str]:
    install_args: list[str] = []
    for raw in deps:
        req = Requirement(raw)
        if req.marker and not req.marker.evaluate({"extra": extra}):
            continue
        req.marker = None
        install_args.append(str(req))
    return install_args


def install_args_for_extra(
    extra: str,
    deps: list[str] | None,
) -> tuple[list[str], str]:
    if deps:
        install_args = _install_requirements_for_extra(extra, deps)
        if install_args:
            return install_args, f"{extra} support"
        return [], f"{extra} support"
    target = f"nanobot-ai[{extra}]"
    return [target], f'"{target}"'


def _requirement_installed(req: Requirement, extra: str, seen: set[tuple[str, str]]) -> bool:
    if req.marker and not req.marker.evaluate({"extra": extra}):
        return True
    key = (
        canonicalize_name(req.name),
        ",".join(sorted(canonicalize_name(value) for value in req.extras)),
    )
    if key in seen:
        return True
    seen.add(key)
    try:
        dist = distribution(req.name)
    except PackageNotFoundError:
        return False
    if req.specifier and not req.specifier.contains(dist.version, prereleases=True):
        return False

    for requested_extra in req.extras:
        if not _extra_dependencies_installed(dist, requested_extra, seen):
            return False
    return True


def _extra_dependencies_installed(
    dist: Any,
    requested_extra: str,
    seen: set[tuple[str, str]],
) -> bool:
    normalized = canonicalize_name(requested_extra)
    provided = {
        canonicalize_name(value)
        for value in cast(list[str], dist.metadata.get_all("Provides-Extra") or [])
    }
    if provided and normalized not in provided:
        return False

    matched = False
    for raw in cast(list[str], dist.requires or []):
        req = Requirement(raw)
        if req.marker and not req.marker.evaluate({"extra": requested_extra}):
            continue
        matched = True
        if not _requirement_installed(req, requested_extra, seen):
            return False
    return matched or bool(provided)


def requirement_installed(raw: str, extra: str = "") -> bool:
    return _requirement_installed(Requirement(raw), extra, set())


def extra_installed(extra: str, deps: list[str] | None) -> bool:
    if deps is None:
        return True
    return all(requirement_installed(dep, extra) for dep in deps)


def run_install_command(
    argv: list[str],
    *,
    env: dict[str, str] | None = None,
) -> subprocess.CompletedProcess[str]:
    try:
        return subprocess.run(
            argv,
            capture_output=True,
            text=True,
            timeout=_INSTALL_TIMEOUT_SECONDS,
            env=env,
        )
    except subprocess.TimeoutExpired as exc:
        stdout = exc.stdout.decode(errors="replace") if isinstance(exc.stdout, bytes) else exc.stdout
        stderr = exc.stderr.decode(errors="replace") if isinstance(exc.stderr, bytes) else exc.stderr
        message = f"Timed out after {_INSTALL_TIMEOUT_SECONDS}s"
        stderr = "\n".join(part for part in ((stderr or "").rstrip(), message) if part)
        return subprocess.CompletedProcess(argv, 124, stdout=stdout or "", stderr=stderr)


def command_text(argv: list[str]) -> str:
    return subprocess.list2cmdline([str(part) for part in argv])


def _log_completed_command(label: str, proc: subprocess.CompletedProcess[str]) -> None:
    logger.info("{} exited with code {}", label, proc.returncode)
    output = (proc.stderr or proc.stdout or "").strip()
    if output:
        logger.info("{} output:\n{}", label, output[:_LOG_OUTPUT_LIMIT])


def missing_pip(proc: subprocess.CompletedProcess[str]) -> bool:
    return "no module named pip" in f"{proc.stdout}\n{proc.stderr}".lower()


def install_extra(
    extra: str,
    deps: list[str] | None,
    *,
    runner: Any = run_install_command,
) -> InstallResult:
    import importlib

    install_args, label = install_args_for_extra(extra, deps)
    pip_cmd = [sys.executable, "-m", "pip", "install", *install_args]
    if not install_args:
        logger.info("Optional feature '{}' has no installable dependencies for this platform", extra)
        return InstallResult(True, label, pip_cmd)

    logger.info("Installing optional feature '{}': {}", extra, command_text(pip_cmd))
    proc = runner(pip_cmd)
    _log_completed_command(f"Optional feature '{extra}' install", proc)
    if proc.returncode == 0:
        importlib.invalidate_caches()
        return InstallResult(True, label, pip_cmd)

    failed_cmd = pip_cmd
    failed_proc = proc
    if missing_pip(proc):
        if shutil.which("uv"):
            uv_cmd = ["uv", "pip", "install", "--python", sys.executable, *install_args]
            uv_env = os.environ.copy()
            if index_url := os.environ.get("PIP_INDEX_URL", "").strip():
                uv_env["UV_INDEX_URL"] = index_url
            logger.info("pip missing while installing '{}'; running {}", extra, command_text(uv_cmd))
            uv_proc = runner(uv_cmd, env=uv_env)
            _log_completed_command(f"Optional feature '{extra}' uv install", uv_proc)
            if uv_proc.returncode == 0:
                importlib.invalidate_caches()
                return InstallResult(True, label, pip_cmd)
            output = (uv_proc.stderr or uv_proc.stdout or "").strip()
            return InstallResult(False, label, pip_cmd, failed_cmd=uv_cmd, output=output)

        ensure_cmd = [sys.executable, "-m", "ensurepip", "--upgrade"]
        logger.info("pip missing while installing '{}'; running {}", extra, command_text(ensure_cmd))
        ensure_proc = runner(ensure_cmd)
        _log_completed_command(f"Optional feature '{extra}' ensurepip", ensure_proc)
        if ensure_proc.returncode == 0:
            logger.info("Retrying optional feature '{}': {}", extra, command_text(pip_cmd))
            proc = runner(pip_cmd)
            _log_completed_command(f"Optional feature '{extra}' install retry", proc)
            if proc.returncode == 0:
                importlib.invalidate_caches()
                return InstallResult(True, label, pip_cmd)
            failed_cmd = pip_cmd
            failed_proc = proc
        else:
            failed_cmd = ensure_cmd
            failed_proc = ensure_proc

    output = (failed_proc.stderr or failed_proc.stdout or "").strip()
    return InstallResult(False, label, pip_cmd, failed_cmd=failed_cmd, output=output)


def read_config_data(path: Path) -> dict[str, Any]:
    if not path.exists():
        return {}
    with open(path, encoding="utf-8") as f:
        return cast(dict[str, Any], json.load(f))


def write_config_data(path: Path, data: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, "w", encoding="utf-8") as f:
        json.dump(data, f, indent=2, ensure_ascii=False)


def set_channel_config_enabled(
    config_path: Path,
    channel_name: str,
    plugin: Any,
    enabled: bool,
    *,
    instance_id: str | None = "default",
) -> None:
    """Persist one instance, or the top-level plugin gate when the target is ``None``."""
    data = read_config_data(config_path)
    channels = data.setdefault("channels", {})
    existing = channels.get(channel_name, {})
    if not isinstance(existing, dict):
        existing = {}
    if instance_id is None:
        existing["enabled"] = enabled
        channels[channel_name] = existing
    else:
        try:
            channels[channel_name] = channel_set_config_enabled(
                plugin,
                existing,
                enabled,
                instance_id=instance_id,
            )
        except ValueError as exc:
            raise OptionalFeatureError(
                f"Invalid {channel_name} configuration: {exc}",
                status=400,
            ) from exc
    write_config_data(config_path, data)


def channel_enabled(
    config: Config,
    name: str,
    plugin: Any | None = None,
    *,
    default_enabled: bool | None = None,
) -> bool:
    section = getattr(config.channels, name, None)
    if default_enabled is None:
        default_enabled = plugin.default_enabled if plugin is not None else channel_default_enabled(name)
    if section is None:
        return bool(default_enabled)
    if plugin is None:
        from nanobot.channels.registry import load_channel_plugin

        plugin = load_channel_plugin(name)
    return bool(channel_instance_specs(plugin, section, enabled_only=True))


def _channel_config_snapshot(
    section: Any,
    name: str,
    spec: ChannelSetupSpec | None,
) -> tuple[dict[str, str], list[str]]:
    if hasattr(section, "model_dump"):
        section = section.model_dump(mode="json", by_alias=True)
    if not isinstance(section, dict):
        return {}, []

    if spec is None:
        return {}, []

    values: dict[str, str] = {}
    configured_fields: list[str] = []
    for field in spec.snapshot_fields:
        value = channel_field_value(section, field)
        if not channel_value_present(value):
            continue
        key = f"channels.{name}.{field}"
        configured_fields.append(key)
        if field in spec.secrets:
            continue
        values[key] = stringify_channel_value(value)
    return values, configured_fields


def _channel_has_required_setup(section: Any, spec: ChannelSetupSpec | None) -> bool:
    return bool(spec and spec.is_configured(section))


def channel_configured(
    config: Config,
    name: str,
    spec: ChannelSetupSpec | None = None,
    plugin: Any | None = None,
    *,
    default_enabled: bool | None = None,
) -> bool:
    """Return whether a channel has enough saved setup to be enabled directly."""
    section = getattr(config.channels, name, None)
    if plugin is None:
        from nanobot.channels.registry import load_channel_plugin

        plugin = load_channel_plugin(name)

    if channel_local_state_present(plugin, section):
        return True
    if section is None:
        return False

    if plugin.management.multi_instance:
        return any(
            _channel_has_required_setup(instance.config, spec)
            for instance in channel_instance_specs(
                plugin,
                section,
                enabled_only=False,
            )
        )

    if not spec or not spec.required:
        return channel_enabled(
            config,
            name,
            plugin,
            default_enabled=default_enabled,
        )
    return _channel_has_required_setup(section, spec)


def _feature_dependencies(
    name: str,
    channel_plugin: Any | None,
    extras: dict[str, list[str] | None],
) -> list[str] | None:
    if channel_plugin is not None:
        return list(channel_plugin.dependencies)
    return extras.get(name)


def optional_features_payload(
    *,
    config: Config | None = None,
    last_action: dict[str, Any] | None = None,
) -> dict[str, Any]:
    from nanobot.channels.registry import discover_plugins
    from nanobot.config.loader import load_config

    config = config or load_config()
    extras = optional_dependency_groups()
    channel_plugins = discover_plugins()
    features: list[dict[str, Any]] = []

    feature_names = set(channel_plugins) | set(extras)
    for name in sorted(feature_names):
        channel_plugin = channel_plugins.get(name)
        is_channel = channel_plugin is not None
        dependencies = _feature_dependencies(name, channel_plugin, extras)
        has_dependencies = bool(dependencies)
        installed = extra_installed(name, dependencies) if has_dependencies else True
        feature: dict[str, Any] = {
            "name": name,
            "display_name": (
                channel_plugin.display_name
                if channel_plugin is not None
                else name.replace("_", " ").title()
            ),
            "type": "channel" if is_channel else "feature",
            "installed": installed,
            "install_supported": has_dependencies or is_channel,
            "requires_restart": _feature_requires_restart(name, is_channel=is_channel),
        }
        if channel_plugin is not None:
            feature["capabilities"] = sorted(channel_plugin.capabilities)
            feature["settings_visible"] = channel_plugin.settings_visible
            if channel_plugin.webui is not None:
                feature["webui"] = channel_plugin.webui

        if not is_channel:
            feature.update({
                "enabled": installed,
                "configured": installed,
                "ready": installed,
                "status": "enabled" if installed else "missing_dependency",
            })
            features.append(feature)
            continue

        try:
            assert channel_plugin is not None
            setup_spec = channel_setup_spec(name, plugin=channel_plugin)
            if setup_spec is not None:
                feature["setup"] = setup_spec.to_public_dict(name)
            enabled = channel_enabled(
                config,
                name,
                channel_plugin,
                default_enabled=channel_plugin.default_enabled,
            )
            configured = channel_configured(
                config,
                name,
                setup_spec,
                channel_plugin,
                default_enabled=channel_plugin.default_enabled,
            )
            ready = bool(enabled and installed)
            status = "enabled" if ready else "missing_dependency" if not installed else "not_enabled"
            feature.update({
                "enabled": enabled,
                "configured": configured,
                "ready": ready,
                "status": status,
            })
            config_values, configured_fields = _channel_config_snapshot(
                getattr(config.channels, name, None),
                name,
                setup_spec,
            )
            if config_values:
                feature["config_values"] = config_values
            if configured_fields:
                feature["configured_fields"] = configured_fields
            instances = channel_feature_instances(
                channel_plugin,
                getattr(config.channels, name, None),
                setup_spec=setup_spec,
            )
            if instances is not None:
                feature["instances"] = instances
        except Exception as exc:
            logger.warning("Could not inspect {} channel configuration: {}", name, exc)
            feature.update({
                "enabled": False,
                "configured": False,
                "ready": False,
                "status": "invalid_config",
                "error": "Channel configuration could not be inspected.",
            })
        features.append(feature)

    payload: dict[str, Any] = {
        "features": features,
        "enabled_count": sum(1 for feature in features if feature["enabled"]),
    }
    if last_action:
        payload["last_action"] = last_action
    return payload


def with_channel_runtime_status(
    payload: dict[str, Any],
    runtime_status: dict[str, Any],
) -> dict[str, Any]:
    """Overlay live ChannelManager state on configuration-derived features."""
    statuses_by_owner: dict[str, list[dict[str, Any]]] = {}
    for status in runtime_status.values():
        if not isinstance(status, dict):
            continue
        status_object = cast(dict[str, Any], status)
        owner = status_object.get("owner")
        if isinstance(owner, str):
            statuses_by_owner.setdefault(owner, []).append(status_object)

    features: list[dict[str, Any]] = []
    for raw_feature in cast(list[object], payload.get("features", [])):
        if not isinstance(raw_feature, dict):
            continue
        feature = cast(dict[str, Any], raw_feature).copy()
        if feature.get("type") != "channel":
            features.append(feature)
            continue

        desired_enabled = bool(feature.get("enabled"))
        owner_statuses = statuses_by_owner.get(str(feature.get("name")), [])
        if desired_enabled and not owner_statuses:
            owner_statuses = [{
                "state": "failed",
                "running": False,
                "error": "Enabled channel has no runtime. Check gateway logs.",
            }]

        instances = feature.get("instances")
        if isinstance(instances, list):
            by_instance = {
                str(status.get("instance_id", "default")): status
                for status in owner_statuses
            }
            decorated_instances: list[dict[str, Any]] = []
            for original_instance in cast(list[object], instances):
                if not isinstance(original_instance, dict):
                    continue
                instance = cast(dict[str, Any], original_instance).copy()
                desired_instance = bool(instance.get("enabled"))
                status = by_instance.get(str(instance.get("id", "default")))
                if desired_instance and status is None:
                    status = {
                        "state": "failed",
                        "running": False,
                        "error": "Enabled channel instance has no runtime. Check gateway logs.",
                    }
                    owner_statuses.append(status)
                state = str(status.get("state", "stopped")) if status else "stopped"
                instance["runtime_status"] = state
                instance["running"] = state == "running"
                if status and status.get("error"):
                    instance["runtime_error"] = str(status["error"])
                decorated_instances.append(instance)
            feature["instances"] = decorated_instances

        state = _combined_channel_runtime_state(owner_statuses, desired_enabled)
        feature["runtime_status"] = state
        feature["running"] = state == "running"
        feature["ready"] = state == "running"
        feature["status"] = "enabled" if state == "running" else state
        error = next(
            (
                str(status["error"])
                for status in owner_statuses
                if status.get("error")
            ),
            None,
        )
        if error:
            feature["runtime_error"] = error
        features.append(feature)

    decorated = dict(payload)
    decorated["features"] = features
    decorated["enabled_count"] = sum(
        1
        for feature in features
        if (
            feature.get("running")
            if feature.get("type") == "channel"
            else feature.get("enabled")
        )
    )
    return decorated


def _combined_channel_runtime_state(
    statuses: list[dict[str, Any]],
    desired_enabled: bool,
) -> str:
    if not desired_enabled:
        return "stopped"
    states = {str(status.get("state", "stopped")) for status in statuses}
    if "failed" in states:
        return "failed"
    if "running" in states:
        return "running"
    if "starting" in states:
        return "starting"
    return "stopped"


def enable_optional_feature(
    name: str,
    *,
    config_path: Path | None = None,
    allow_install: bool = True,
    instance_id: str | None = None,
    runner: Any = run_install_command,
) -> dict[str, Any]:
    from nanobot.channels.registry import discover_plugins
    from nanobot.config.loader import get_config_path

    if name in _BUNDLED_FEATURE_ALIASES:
        payload = optional_features_payload(
            last_action={
                "ok": True,
                "message": f"Feature '{name}' is included with nanobot",
                "enabled": True,
            }
        )
        payload["requires_restart"] = False
        return payload
    config_path = config_path or get_config_path()
    requested_instance_id = (instance_id or "").strip() or None
    extras = optional_dependency_groups()
    channel_plugins = discover_plugins()
    known = set(channel_plugins) | set(extras)
    if name not in known:
        available = ", ".join(sorted(known))
        raise OptionalFeatureError(f"Unknown feature: {name}. Available: {available}", status=404)

    channel_plugin = channel_plugins.get(name)
    dependencies = _feature_dependencies(name, channel_plugin, extras)
    if dependencies and not extra_installed(name, dependencies):
        if not allow_install:
            raise OptionalFeatureError(
                "Installing optional features from a remote WebUI is disabled. "
                "Run this action from localhost or set tools.webuiAllowRemotePackageInstall to true.",
                status=403,
            )
        result = install_extra(
            name,
            dependencies,
            runner=runner,
        )
        if not result.ok:
            failed = command_text(result.failed_cmd or result.pip_cmd)
            detail = f": {result.output}" if result.output else ""
            raise OptionalFeatureError(f"Failed: {failed}{detail}", status=500)

    channel_cls: Any | None = None
    target_instance_id: str | None = None
    if channel_plugin is not None:
        try:
            channel_cls = channel_plugin.load_channel_class()
        except Exception as exc:
            raise OptionalFeatureError(
                f"Channel '{name}' is not importable after enable: {exc}",
                status=500,
            ) from exc
        target_instance_id = resolve_channel_action_target(
            requested_instance_id,
        )
        set_channel_config_enabled(
            config_path,
            name,
            channel_plugin,
            True,
            instance_id=target_instance_id,
        )
        message = f"Enabled channel '{name}'"
    else:
        message = f"Enabled feature '{name}'"

    if channel_cls is not None and target_instance_id is not None:
        try:
            refresh_channel_feature_metadata(
                channel_cls,
                config_path,
                instance_id=target_instance_id,
            )
        except Exception as exc:
            logger.warning("Could not refresh {} channel metadata: {}", name, exc)

    from nanobot.config.loader import load_config

    payload = optional_features_payload(
        config=load_config(config_path),
        last_action={"ok": True, "message": message, "enabled": True},
    )
    payload["requires_restart"] = _feature_requires_restart(
        name,
        is_channel=channel_plugin is not None,
    )
    return payload


def ensure_enabled_channel_dependencies(
    enabled_names: set[str],
    plugins: dict[str, Any],
    *,
    runner: Any = run_install_command,
) -> dict[str, str]:
    """Install requirements declared by enabled channel manifests.

    Returns user-safe errors keyed by channel name. Detailed installer output
    remains in gateway logs.
    """
    failures: dict[str, str] = {}
    for name in sorted(enabled_names):
        plugin = plugins.get(name)
        if plugin is None:
            continue
        dependencies = list(plugin.dependencies)
        if not dependencies or extra_installed(name, dependencies):
            continue
        result = install_extra(name, dependencies, runner=runner)
        if result.ok and extra_installed(name, dependencies):
            continue
        failures[name] = "Channel dependencies could not be installed. Check gateway logs."
        logger.error("Could not prepare dependencies for enabled channel '{}'", name)
    return failures


def _feature_requires_restart(name: str, *, is_channel: bool) -> bool:
    """Return whether an installed feature needs the running engine rebuilt."""
    if is_channel:
        return True
    # These libraries are imported lazily or used by a newly spawned service.
    return name not in {"api", "documents", "pdf", "olostep"}


def disable_optional_feature(
    name: str,
    *,
    config_path: Path | None = None,
    instance_id: str | None = None,
) -> dict[str, Any]:
    from nanobot.channels.registry import discover_plugins
    from nanobot.config.loader import get_config_path, load_config

    config_path = config_path or get_config_path()
    requested_instance_id = (instance_id or "").strip() or None
    extras = optional_dependency_groups()
    channel_plugins = discover_plugins()
    known_channels = set(channel_plugins)
    known = known_channels | set(extras)
    if name not in known:
        available = ", ".join(sorted(known))
        raise OptionalFeatureError(f"Unknown feature: {name}. Available: {available}", status=404)
    if name not in known_channels:
        raise OptionalFeatureError(f"Feature '{name}' cannot be disabled", status=400)
    channel_plugin = channel_plugins[name]
    target_instance_id = resolve_channel_action_target(requested_instance_id)
    set_channel_config_enabled(
        config_path,
        name,
        channel_plugin,
        False,
        instance_id=target_instance_id,
    )
    payload = optional_features_payload(
        config=load_config(config_path),
        last_action={"ok": True, "message": f"Disabled channel '{name}'", "enabled": False}
    )
    payload["requires_restart"] = True
    return payload
