from __future__ import annotations

import os
import shlex
from pathlib import Path
from typing import Any, Dict, Optional, Union

from ..artifacts import RuntimeArtifactStoreAdapter, get_artifact_id, is_artifact_ref
from ..errors import AbstractVisionError
from ..vision_manager import VisionManager

_DEFAULT_LOCAL_DIFFUSERS_MODEL_ID = "runwayml/stable-diffusion-v1-5"


def _env(key: str, default: Optional[str] = None) -> Optional[str]:
    v = os.environ.get(str(key), None)
    if v is None:
        return default
    s = str(v).strip()
    return s if s else default


def _env_bool(key: str, default: bool = False) -> bool:
    v = _env(key)
    if v is None:
        return bool(default)
    return str(v).strip().lower() in {"1", "true", "yes", "on"}


def _owner_cfg(owner: Any, key: str) -> Optional[str]:
    try:
        cfg = getattr(owner, "config", None)
        if isinstance(cfg, dict):
            v = cfg.get(key)
            if v is None:
                return None
            s = str(v).strip()
            return s if s else None
    except Exception:
        return None
    return None


def _owner_cfg_bool(owner: Any, key: str, default: bool = False) -> bool:
    v = _owner_cfg(owner, key)
    if v is None:
        return bool(default)
    return str(v).strip().lower() in {"1", "true", "yes", "on"}


def _read_bytes_from_path(path: Union[str, Path]) -> bytes:
    p = Path(str(path)).expanduser()
    return p.read_bytes()


def _resolve_bytes_input(value: Union[bytes, Dict[str, Any], str], *, artifact_store: Any) -> bytes:
    if isinstance(value, (bytes, bytearray)):
        return bytes(value)
    if isinstance(value, dict):
        if not is_artifact_ref(value):
            raise ValueError("Expected an artifact ref dict like {'$artifact': '...'}")
        if artifact_store is None:
            raise ValueError("artifact_store is required to resolve artifact refs to bytes")
        store = RuntimeArtifactStoreAdapter(artifact_store)
        return store.load_bytes(get_artifact_id(value))
    if isinstance(value, str):
        p = Path(value).expanduser()
        if p.exists() and p.is_file():
            return p.read_bytes()
        raise FileNotFoundError(f"File not found: {value}")
    raise TypeError("Unsupported input type; expected bytes, artifact-ref dict, or file path")


class _AbstractVisionCapability:
    """AbstractCore VisionCapability backed by AbstractVision."""

    backend_id = "abstractvision:openai-compatible"

    def __init__(self, owner: Any):
        self._owner = owner
        self._backend = None

    def _get_backend(self):
        if self._backend is not None:
            return self._backend

        # Injection hook (useful for tests and advanced embedding).
        try:
            cfg = getattr(self._owner, "config", None)
            if isinstance(cfg, dict):
                inst = cfg.get("vision_backend_instance")
                if inst is not None:
                    self._backend = inst
                    return self._backend
                factory = cfg.get("vision_backend_factory")
                if callable(factory):
                    self._backend = factory(self._owner)
                    return self._backend
        except Exception:
            pass

        # Prefer AbstractCore config keys when present; fall back to AbstractVision env vars.
        backend_kind = (
            _owner_cfg(self._owner, "vision_backend")
            or _env("ABSTRACTVISION_BACKEND", "diffusers")
            or "diffusers"
        ).lower()
        if backend_kind in {"openai_compatible", "openai-compatible", "proxy"}:
            backend_kind = "openai"
        elif backend_kind in {"huggingface", "hf", "hf-diffusers"}:
            backend_kind = "diffusers"
        elif backend_kind in {
            "sd-cpp",
            "stable-diffusion.cpp",
            "stable-diffusion-cpp",
            "stable_diffusion_cpp",
        }:
            backend_kind = "sdcpp"

        if backend_kind == "diffusers":
            model_id = (
                _owner_cfg(self._owner, "vision_model_id")
                or _env("ABSTRACTVISION_DIFFUSERS_MODEL_ID")
                or _env("ABSTRACTVISION_MODEL_ID")
                or _DEFAULT_LOCAL_DIFFUSERS_MODEL_ID
            )
            device = (
                _owner_cfg(self._owner, "vision_device")
                or _env("ABSTRACTVISION_DIFFUSERS_DEVICE", "auto")
                or "auto"
            )
            torch_dtype = _owner_cfg(self._owner, "vision_torch_dtype") or _env(
                "ABSTRACTVISION_DIFFUSERS_TORCH_DTYPE"
            )
            allow_download = _owner_cfg_bool(
                self._owner,
                "vision_allow_download",
                _env_bool("ABSTRACTVISION_DIFFUSERS_ALLOW_DOWNLOAD", False),
            )
            auto_retry_fp32 = _owner_cfg_bool(
                self._owner,
                "vision_auto_retry_fp32",
                _env_bool("ABSTRACTVISION_DIFFUSERS_AUTO_RETRY_FP32", True),
            )

            from ..backends.huggingface_diffusers import (
                HuggingFaceDiffusersBackendConfig,
                HuggingFaceDiffusersVisionBackend,
            )

            cfg = HuggingFaceDiffusersBackendConfig(
                model_id=str(model_id),
                device=str(device),
                torch_dtype=str(torch_dtype) if torch_dtype else None,
                allow_download=allow_download,
                auto_retry_fp32=auto_retry_fp32,
            )
            self._backend = HuggingFaceDiffusersVisionBackend(config=cfg)
            return self._backend

        if backend_kind == "sdcpp":
            model = _owner_cfg(self._owner, "vision_sdcpp_model") or _env(
                "ABSTRACTVISION_SDCPP_MODEL"
            )
            diffusion_model = _owner_cfg(self._owner, "vision_sdcpp_diffusion_model") or _env(
                "ABSTRACTVISION_SDCPP_DIFFUSION_MODEL"
            )
            if not model and not diffusion_model:
                raise AbstractVisionError(
                    "Missing stable-diffusion.cpp model configuration. Set vision_sdcpp_model, "
                    "vision_sdcpp_diffusion_model, ABSTRACTVISION_SDCPP_MODEL, or ABSTRACTVISION_SDCPP_DIFFUSION_MODEL."
                )

            from ..backends.stable_diffusion_cpp import (
                StableDiffusionCppBackendConfig,
                StableDiffusionCppVisionBackend,
            )

            extra_args = _owner_cfg(self._owner, "vision_sdcpp_extra_args") or _env(
                "ABSTRACTVISION_SDCPP_EXTRA_ARGS"
            )
            cfg = StableDiffusionCppBackendConfig(
                sd_cli_path=_owner_cfg(self._owner, "vision_sdcpp_bin")
                or _env("ABSTRACTVISION_SDCPP_BIN", "sd-cli")
                or "sd-cli",
                model=str(model) if model else None,
                diffusion_model=str(diffusion_model) if diffusion_model else None,
                vae=_owner_cfg(self._owner, "vision_sdcpp_vae") or _env("ABSTRACTVISION_SDCPP_VAE"),
                llm=_owner_cfg(self._owner, "vision_sdcpp_llm") or _env("ABSTRACTVISION_SDCPP_LLM"),
                llm_vision=_owner_cfg(self._owner, "vision_sdcpp_llm_vision")
                or _env("ABSTRACTVISION_SDCPP_LLM_VISION"),
                clip_l=_owner_cfg(self._owner, "vision_sdcpp_clip_l")
                or _env("ABSTRACTVISION_SDCPP_CLIP_L"),
                clip_g=_owner_cfg(self._owner, "vision_sdcpp_clip_g")
                or _env("ABSTRACTVISION_SDCPP_CLIP_G"),
                t5xxl=_owner_cfg(self._owner, "vision_sdcpp_t5xxl")
                or _env("ABSTRACTVISION_SDCPP_T5XXL"),
                extra_args=tuple(shlex.split(str(extra_args))) if extra_args else (),
                timeout_s=float(
                    _owner_cfg(self._owner, "vision_timeout_s")
                    or _env("ABSTRACTVISION_TIMEOUT_S", "3600")
                    or "3600"
                ),
            )
            self._backend = StableDiffusionCppVisionBackend(config=cfg)
            return self._backend

        if backend_kind != "openai":
            raise AbstractVisionError(
                f"Unsupported AbstractVision backend for AbstractCore plugin: {backend_kind!r}. "
                "Use 'diffusers', 'sdcpp', or 'openai'."
            )

        base_url = _owner_cfg(self._owner, "vision_base_url") or _env("ABSTRACTVISION_BASE_URL")
        api_key = _owner_cfg(self._owner, "vision_api_key") or _env("ABSTRACTVISION_API_KEY")
        model_id = _owner_cfg(self._owner, "vision_model_id") or _env("ABSTRACTVISION_MODEL_ID")
        timeout_s_raw = _owner_cfg(self._owner, "vision_timeout_s") or _env(
            "ABSTRACTVISION_TIMEOUT_S"
        )
        try:
            timeout_s = float(timeout_s_raw) if timeout_s_raw else 300.0
        except Exception:
            timeout_s = 300.0

        if not base_url:
            raise AbstractVisionError(
                "Missing vision_base_url / ABSTRACTVISION_BASE_URL. "
                "Configure an OpenAI-compatible endpoint (e.g. http://localhost:8000/v1)."
            )

        # Optional video endpoints (not standardized; only enabled when configured).
        t2v_path = _owner_cfg(self._owner, "vision_text_to_video_path") or _env(
            "ABSTRACTVISION_TEXT_TO_VIDEO_PATH"
        )
        i2v_path = _owner_cfg(self._owner, "vision_image_to_video_path") or _env(
            "ABSTRACTVISION_IMAGE_TO_VIDEO_PATH"
        )
        i2v_mode = _owner_cfg(self._owner, "vision_image_to_video_mode") or _env(
            "ABSTRACTVISION_IMAGE_TO_VIDEO_MODE", "multipart"
        )

        # Import backend module lazily (keeps plugin import-light).
        from ..backends.openai_compatible import (
            OpenAICompatibleBackendConfig,
            OpenAICompatibleVisionBackend,
        )

        cfg = OpenAICompatibleBackendConfig(
            base_url=str(base_url),
            api_key=str(api_key) if api_key else None,
            model_id=str(model_id) if model_id else None,
            timeout_s=float(timeout_s),
            text_to_video_path=str(t2v_path) if t2v_path else None,
            image_to_video_path=str(i2v_path) if i2v_path else None,
            image_to_video_mode=str(i2v_mode or "multipart"),
        )
        self._backend = OpenAICompatibleVisionBackend(config=cfg)
        return self._backend

    def _make_manager(self, *, artifact_store: Any) -> VisionManager:
        store = RuntimeArtifactStoreAdapter(artifact_store) if artifact_store is not None else None
        return VisionManager(backend=self._get_backend(), store=store)

    def t2i(self, prompt: str, **kwargs: Any):
        store = kwargs.pop("artifact_store", None)
        vm = self._make_manager(artifact_store=store)
        out = vm.generate_image(str(prompt), **kwargs)
        if isinstance(out, dict):
            return out
        return bytes(getattr(out, "data", b""))

    def i2i(self, prompt: str, image: Union[bytes, Dict[str, Any], str], **kwargs: Any):
        store = kwargs.pop("artifact_store", None)
        image_b = _resolve_bytes_input(image, artifact_store=store)
        mask = kwargs.pop("mask", None)
        mask_b = None
        if mask is not None:
            mask_b = _resolve_bytes_input(mask, artifact_store=store)
        vm = self._make_manager(artifact_store=store)
        out = vm.edit_image(str(prompt), image=image_b, mask=mask_b, **kwargs)
        if isinstance(out, dict):
            return out
        return bytes(getattr(out, "data", b""))

    def t2v(self, prompt: str, **kwargs: Any):
        store = kwargs.pop("artifact_store", None)
        vm = self._make_manager(artifact_store=store)
        out = vm.generate_video(str(prompt), **kwargs)
        if isinstance(out, dict):
            return out
        return bytes(getattr(out, "data", b""))

    def i2v(self, image: Union[bytes, Dict[str, Any], str], **kwargs: Any):
        store = kwargs.pop("artifact_store", None)
        image_b = _resolve_bytes_input(image, artifact_store=store)
        vm = self._make_manager(artifact_store=store)
        out = vm.image_to_video(image=image_b, **kwargs)
        if isinstance(out, dict):
            return out
        return bytes(getattr(out, "data", b""))


def register(registry: Any) -> None:
    """Register AbstractVision as an AbstractCore capability plugin.

    This function is loaded via the `abstractcore.capabilities_plugins` entry point group.
    """

    def _factory(owner: Any) -> _AbstractVisionCapability:
        return _AbstractVisionCapability(owner)

    config_hint = (
        "Use the default local Diffusers backend with ABSTRACTVISION_MODEL_ID=runwayml/stable-diffusion-v1-5 "
        "(cache-only unless ABSTRACTVISION_DIFFUSERS_ALLOW_DOWNLOAD=1), or set ABSTRACTVISION_BACKEND=openai "
        "plus ABSTRACTVISION_BASE_URL to point to an OpenAI-compatible /v1 endpoint."
    )

    registry.register_vision_backend(
        backend_id=_AbstractVisionCapability.backend_id,
        factory=_factory,
        priority=0,
        description="AbstractVision capability plugin (Diffusers, stable-diffusion.cpp, or OpenAI-compatible HTTP; env/config-driven).",
        config_hint=config_hint,
    )
