"""Backend-for-frontend proxy for multi-provider AI chat APIs.

This module integrates LiteLLM to support OpenAI, Anthropic, AWS Bedrock, xAI
Grok, and Google Gemini through a unified streaming interface while maintaining
PyTincture compatibility.
"""

import json
import os
import sys
import types

import litellm

# Ensure the module is registered for inspect.getfile when loaded via SourceFileLoader.
if __name__ not in sys.modules:
    module_stub = types.ModuleType(__name__)
    module_stub.__file__ = __file__
    sys.modules[__name__] = module_stub
elif not getattr(sys.modules[__name__], "__file__", None):
    sys.modules[__name__].__file__ = __file__

from pytincture.dataclass import backend_for_frontend, bff_stream

# ---------------------------------------------------------------------------
# Defaults (can be overridden when instantiating the proxy)
# ---------------------------------------------------------------------------

DEFAULT_PROVIDER_CONFIG = {
    "providers": {
        "aws_bedrock": {
            "anthropic": [
                "us.anthropic.claude-sonnet-4-20250514-v1:0",
                "us.anthropic.claude-opus-4-20250514-v1:0",
                "us.anthropic.claude-opus-4-1-20250805-v1:0",
                "us.anthropic.claude-3-5-sonnet-20241022-v2:0",
                "us.anthropic.claude-3-7-sonnet-20250219-v1:0",
            ],
            "meta": [
                "us.meta.llama3-3-70b-instruct-v1:0",
                "us.meta.llama4-maverick-17b-instruct-v1:0",
                "us.meta.llama4-scout-17b-instruct-v1:0",
            ],
            "amazon": [
                "us.amazon.nova-pro-v1:0",
                "us.amazon.nova-premier-v1:0",
                "us.amazon.nova-micro-v1:0",
                "us.amazon.nova-lite-v1:0",
            ],
        },
        "anthropic": [
            "claude-sonnet-4-20250514",
            "claude-opus-4-20250514",
            "claude-opus-4-1-20250805",
            "claude-3-7-sonnet-latest",
            "claude-3-5-sonnet-latest",
        ],
        "openai": [
            "gpt-5",
            "gpt-5-mini",
            "o4-mini",
            "gpt-4o-mini",
            "gpt-4o",
            "gpt-4.1",
            "gpt-4.1-mini",
            "o3",
            "o3-mini",
            "o1",
            "o1-mini",
            "chatgpt-4o-latest",
            "gpt-3.5-turbo",
        ],
        "xai": [
            "xai/grok-3-mini-beta",
            "xai/grok-3-beta",
            "xai/grok-2",
            "xai/grok-2-mini",
        ],
        "google": [
            "gemini-pro",
            "gemini-pro-vision",
            "gemini-1.5-pro",
            "gemini-1.5-flash",
        ],
    }
}

DEFAULT_MODEL = os.getenv("DEFAULT_MODEL", "gpt-4o")
REQUEST_TIMEOUT = float(os.getenv("REQUEST_TIMEOUT", "60"))
DEBUG_STREAM = os.getenv("MULTIPROXY_DEBUG_STREAM", "").lower() in {"1", "true", "yes"}

# ---------------------------------------------------------------------------
# LiteLLM provider helpers
# ---------------------------------------------------------------------------


class UnifiedAIProvider:
    """Unified AI provider using LiteLLM for multi-provider support."""

    def __init__(self, provider_config):
        self.provider_config = provider_config or {"providers": {}}
        self.setup_environment()
        self._build_model_mapping()

    @staticmethod
    def setup_environment() -> None:
        """Populate LiteLLM environment variables from existing env vars."""

        env_mappings = {
            "OPENAI_API_KEY": "OPENAI_API_KEY",
            "ANTHROPIC_API_KEY": "ANTHROPIC_API_KEY",
            "GOOGLE_KEY": "GEMINI_API_KEY",
            "XAI_API_KEY": "XAI_API_KEY",
            "AWS_ACCESS_KEY_ID": "AWS_ACCESS_KEY_ID",
            "AWS_SECRET_ACCESS_KEY": "AWS_SECRET_ACCESS_KEY",
            "AWS_DEFAULT_REGION": "AWS_DEFAULT_REGION",
        }

        for source_var, target_var in env_mappings.items():
            value = os.getenv(source_var)
            if value and not os.getenv(target_var):
                os.environ[target_var] = value

    def _build_model_mapping(self):
        providers = self.provider_config.get("providers", {})
        self.model_mapping = {}

        bedrock = providers.get("aws_bedrock", {})
        for provider_type, models in bedrock.items():
            for model in models:
                litellm_name = self._map_bedrock_model(provider_type, model)
                self.model_mapping[model] = litellm_name

        for provider, models in providers.items():
            if provider == "aws_bedrock":
                continue
            for model in models:
                self.model_mapping[model] = self._map_direct_model(provider, model)

    @staticmethod
    def _map_bedrock_model(provider_type, model):
        prefix_map = {
            "anthropic": "anthropic.",
            "meta": "meta.",
            "amazon": "amazon.",
        }
        prefix = prefix_map.get(provider_type)
        if prefix:
            clean = model.replace(f"us.{provider_type}.", prefix)
            return f"bedrock/{clean}"
        return f"bedrock/{model}"

    @staticmethod
    def _map_direct_model(provider, model):
        if provider == "openai" or provider == "xai":
            return model
        if provider == "google":
            return f"gemini/{model}"
        return f"{provider}/{model}"

    def get_litellm_model(self, model):
        return self.model_mapping.get(model, model)

    def is_xai_model(self, model):
        providers = self.provider_config.get("providers", {})
        if model.startswith("xai/"):
            return True
        return model in providers.get("xai", []) or model.startswith("grok")

    def stream_completion(self, model, messages, **kwargs):
        litellm_model = self.get_litellm_model(model)

        if self.is_xai_model(model):
            original_base_url = os.getenv("OPENAI_BASE_URL")
            os.environ["OPENAI_BASE_URL"] = "https://api.x.ai/v1"
            try:
                response = litellm.completion(
                    model=litellm_model,
                    messages=list(messages),
                    stream=True,
                    api_key=os.getenv("XAI_API_KEY"),
                    **kwargs,
                )
                for chunk in response:
                    yield chunk
            finally:
                if original_base_url is not None:
                    os.environ["OPENAI_BASE_URL"] = original_base_url
                elif "OPENAI_BASE_URL" in os.environ:
                    del os.environ["OPENAI_BASE_URL"]
        else:
            response = litellm.completion(
                model=litellm_model,
                messages=list(messages),
                stream=True,
                **kwargs,
            )
            for chunk in response:
                yield chunk

    def complete(self, model, messages, **kwargs):
        litellm_model = self.get_litellm_model(model)

        if self.is_xai_model(model):
            original_base_url = os.getenv("OPENAI_BASE_URL")
            os.environ["OPENAI_BASE_URL"] = "https://api.x.ai/v1"
            try:
                return litellm.completion(
                    model=litellm_model,
                    messages=list(messages),
                    stream=False,
                    api_key=os.getenv("XAI_API_KEY"),
                    **kwargs,
                )
            finally:
                if original_base_url is not None:
                    os.environ["OPENAI_BASE_URL"] = original_base_url
                elif "OPENAI_BASE_URL" in os.environ:
                    del os.environ["OPENAI_BASE_URL"]

        return litellm.completion(
            model=litellm_model,
            messages=list(messages),
            stream=False,
            **kwargs,
        )


# ---------------------------------------------------------------------------
# Backend-for-frontend proxy
# ---------------------------------------------------------------------------


@backend_for_frontend
class multiaiproxy:
    """Multi-provider AI proxy supporting OpenAI, Anthropic, Bedrock, xAI, and Google."""

    def __init__(self, *, provider_config=None, default_model=None, timeout=None):
        config = provider_config or self._load_config_from_env() or DEFAULT_PROVIDER_CONFIG
        self._provider = UnifiedAIProvider(config)
        self._provider_config = config
        self._default_model = default_model or os.getenv("DEFAULT_MODEL", DEFAULT_MODEL)
        self._timeout = timeout or REQUEST_TIMEOUT

    @staticmethod
    def _normalize_stream_chunk(chunk):
        if chunk is None:
            return None

        payload = chunk
        if hasattr(payload, "model_dump"):
            payload = payload.model_dump(exclude_none=True)
        elif hasattr(payload, "dict"):
            payload = payload.dict(exclude_none=True)

        if isinstance(payload, (bytes, bytearray)):
            payload = payload.decode("utf-8", errors="replace")

        if isinstance(payload, str):
            text = payload.strip()
            if not text or text == "[DONE]":
                return None
            if text.startswith("data:"):
                text = text[5:].strip()
                if text == "[DONE]":
                    return None
                try:
                    payload = json.loads(text)
                except json.JSONDecodeError:
                    return {"choices": [{"delta": {"content": text}}]}
            else:
                return {"choices": [{"delta": {"content": text}}]}

        if not isinstance(payload, dict):
            return None

        if payload.get("type") == "content.delta":
            delta = payload.get("delta")
            if isinstance(delta, dict):
                text = delta.get("content") or delta.get("text") or ""
            else:
                text = delta or ""
            if text:
                return {"choices": [{"delta": {"content": text}}]}
            return None

        if "choices" not in payload:
            if "delta" in payload:
                delta = payload.get("delta")
                if isinstance(delta, dict):
                    text = delta.get("content") or delta.get("text") or ""
                else:
                    text = delta or ""
                if text:
                    return {"choices": [{"delta": {"content": text}}]}
            elif "content" in payload:
                text = payload.get("content")
                if isinstance(text, str) and text:
                    return {"choices": [{"delta": {"content": text}}]}

        return payload

    @staticmethod
    def _load_config_from_env():
        raw = os.getenv("MULTIPROXY_PROVIDER_CONFIG")
        if not raw:
            return None
        try:
            return json.loads(raw)
        except json.JSONDecodeError:
            print("Warning: MULTIPROXY_PROVIDER_CONFIG is not valid JSON; falling back to defaults.")
            return None

    def get_available_models(self):
        return self._provider_config.get("providers", {})

    def get_model_info(self, model):
        litellm_name = self._provider.get_litellm_model(model)
        provider = "unknown"
        providers = self._provider_config.get("providers", {})

        if model in providers.get("openai", []):
            provider = "openai"
        elif model in providers.get("anthropic", []):
            provider = "anthropic"
        elif model in providers.get("xai", []):
            provider = "xai"
        elif model in providers.get("google", []):
            provider = "google"
        elif any(model in group for group in providers.get("aws_bedrock", {}).values()):
            provider = "aws_bedrock"

        return {
            "original_name": model,
            "litellm_name": litellm_name,
            "provider": provider,
            "supported": model in self._provider.model_mapping,
        }

    @bff_stream()
    def chat_stream(self, messages, model=None, **extra):
        options = extra.copy()
        timeout_override = options.pop("timeout", None)
        options.pop("stream", None)
        timeout = timeout_override or self._timeout

        selected_model = model or self._default_model

        if selected_model not in self._provider.model_mapping:
            raise ValueError(
                f"Model '{selected_model}' is not supported. Available models: {list(self._provider.model_mapping.keys())}"
            )

        try:
            yielded_any = False
            for chunk in self._provider.stream_completion(
                model=selected_model,
                messages=messages,
                timeout=timeout,
                **options,
            ):
                if DEBUG_STREAM:
                    print(f"[multiproxy] raw chunk: {type(chunk)!r} {chunk!r}")
                normalized = self._normalize_stream_chunk(chunk)
                if DEBUG_STREAM:
                    print(f"[multiproxy] normalized: {normalized!r}")
                if normalized is not None:
                    yielded_any = True
                    yield normalized
            if not yielded_any:
                response = self._provider.complete(
                    model=selected_model,
                    messages=messages,
                    timeout=timeout,
                    **options,
                )
                payload = response
                if hasattr(payload, "model_dump"):
                    payload = payload.model_dump(exclude_none=True)
                elif hasattr(payload, "dict"):
                    payload = payload.dict(exclude_none=True)
                text = ""
                if isinstance(payload, dict):
                    for choice in payload.get("choices") or []:
                        message = (choice or {}).get("message") or {}
                        if isinstance(message, dict):
                            text = message.get("content") or ""
                        if text:
                            break
                if text:
                    yield {"choices": [{"delta": {"content": text}}]}
        except Exception as exc:  # pragma: no cover - provider errors
            yield {
                "error": {
                    "message": str(exc),
                    "type": "provider_error",
                    "code": "stream_error",
                }
            }

    @bff_stream()
    def chat_stream_with_provider_info(self, messages, model=None, **extra):
        model_info = self.get_model_info(model or self._default_model)
        yield {"type": "model_info", "model_info": model_info}
        for chunk in self.chat_stream(messages, model, **extra):
            if isinstance(chunk, dict) and "error" not in chunk:
                chunk["provider"] = model_info["provider"]
            yield chunk


# Backwards compatibility alias expected by existing demos
openaiproxy = multiaiproxy
