from __future__ import annotations

import os
from typing import Any

from Mama.embedding_factory import get_embeddings
from Mama.models import ExecutionContext
from Mama.ports import EmbeddingsProviderPort, LlmProviderPort, RuntimePorts
from Mama.vectorstores import DynamicVectorStoreFactory

try:
    from langchain_core.runnables import RunnableLambda  # type: ignore
except Exception:  # pragma: no cover
    RunnableLambda = None  # type: ignore

try:
    from langchain_openai import ChatOpenAI  # type: ignore
except Exception:  # pragma: no cover
    ChatOpenAI = None  # type: ignore

try:
    from langchain_ollama import ChatOllama  # type: ignore
except Exception:  # pragma: no cover
    ChatOllama = None  # type: ignore


class DefaultEmbeddingsProvider(EmbeddingsProviderPort):
    def get_embeddings(self, ctx: ExecutionContext) -> Any:
        backend = str(ctx.storage.params.get("embeddings_backend", "auto"))
        model = str(ctx.storage.params.get("embeddings_model", ""))
        base_url = str(ctx.storage.params.get("embeddings_base_url", ""))
        if model:
            os.environ["MAMA_OLLAMA_EMBED_MODEL"] = model
        if base_url:
            os.environ["MAMA_OLLAMA_BASE_URL"] = base_url
        return get_embeddings(backend)


class DefaultLlmProvider(LlmProviderPort):
    def build_llm(self, ctx: ExecutionContext) -> Any:
        provider = (ctx.llm.provider or "dummy").strip().lower()
        if provider in {"dummy", "local_dummy"}:
            return self._build_dummy()

        if provider in {"openai", "chatopenai"}:
            if ChatOpenAI is None:
                return self._build_dummy(prefix="[DUMMY-LLM] OpenAI SDK unavailable")
            params = dict(ctx.llm.parameters)
            if "model" not in params and ctx.llm.model:
                params["model"] = ctx.llm.model
            return ChatOpenAI(**params)

        if provider in {"ollama", "local_ollama"}:
            if ChatOllama is None:
                return self._build_dummy(prefix="[DUMMY-LLM] Ollama SDK unavailable")
            params = dict(ctx.llm.parameters)
            if "model" not in params and ctx.llm.model:
                params["model"] = ctx.llm.model
            params.setdefault("base_url", "http://127.0.0.1:11434")
            return ChatOllama(**params)

        return self._build_dummy(prefix=f"[DUMMY-LLM] Unsupported provider '{provider}'")

    @staticmethod
    def _build_dummy(prefix: str = "[DUMMY-LLM]"):
        def _fn(prompt: Any) -> str:
            if hasattr(prompt, "to_string"):
                prompt = prompt.to_string()  # type: ignore
            text = str(prompt)
            compact = " ".join(text.split())
            if len(compact) > 320:
                compact = compact[:317].rsplit(" ", 1)[0] + "..."
            return f"{prefix} {compact}"

        if RunnableLambda:
            return RunnableLambda(lambda x: _fn(x))  # type: ignore
        return lambda x: _fn(x)


class DefaultRuntime(RuntimePorts):
    def __init__(self):
        super().__init__(
            llm_provider=DefaultLlmProvider(),
            embeddings_provider=DefaultEmbeddingsProvider(),
            vector_store_factory=DynamicVectorStoreFactory(),
        )


def default_runtime() -> RuntimePorts:
    return DefaultRuntime()
