from __future__ import annotations

import os
from pathlib import Path
from typing import Optional

import oci
from litellm import CustomLLM
from litellm.types.utils import Embedding, EmbeddingResponse, Usage


class OCIEmbeddingLLM(CustomLLM):
    def __init__(self):
        super().__init__()
        self._client = None
        self._client_signature = None

    @staticmethod
    def _pick_value(optional_params: dict, param_name: str, env_name: str, default: Optional[str] = None):
        value = optional_params.get(param_name)
        if value is None:
            value = os.getenv(env_name, default)
        if isinstance(value, str):
            value = value.strip()
        return value or default

    def _build_oci_config(self, optional_params: dict) -> dict:
        # Build OCI SDK config from LiteLLM params first, then fallback to env vars.
        oci_config = {
            "user": self._pick_value(optional_params, "oci_user", "OCI_USER"),
            "fingerprint": self._pick_value(optional_params, "oci_fingerprint", "OCI_FINGERPRINT"),
            "tenancy": self._pick_value(optional_params, "oci_tenancy", "OCI_TENANCY"),
            "region": self._pick_value(optional_params, "oci_region", "OCI_REGION"),
            "key_file": self._pick_value(optional_params, "oci_key_file", "OCI_KEY_FILE"),
            "pass_phrase": self._pick_value(optional_params, "oci_pass_phrase", "OCI_PASS_PHRASE"),
        }

        if oci_config.get("key_file"):
            oci_config["key_file"] = str(Path(oci_config["key_file"]).expanduser())

        # Remove optional empty values so OCI SDK handles defaults cleanly.
        if not oci_config.get("pass_phrase"):
            oci_config.pop("pass_phrase", None)

        missing = [
            key
            for key in ("user", "fingerprint", "tenancy", "region", "key_file")
            if not oci_config.get(key)
        ]
        if missing:
            raise ValueError(
                "Missing OCI config values for embedding provider: " + ", ".join(missing)
            )

        return oci_config

    def _get_client(self, optional_params: dict):
        """Lazy init OCI client and recreate only when OCI auth/region changes."""
        config = self._build_oci_config(optional_params)
        signature = tuple(sorted(config.items()))

        if self._client is None or self._client_signature != signature:
            self._client = oci.generative_ai_inference.GenerativeAiInferenceClient(config)
            self._client_signature = signature

        return self._client

    def _build_serving_mode(self, model_id: str, optional_params: dict):
        serving_mode = self._pick_value(
            optional_params,
            "oci_serving_mode",
            "OCI_SERVING_MODE",
            default="ON_DEMAND",
        ).upper()

        if serving_mode == "DEDICATED":
            endpoint_id = self._pick_value(optional_params, "oci_endpoint_id", "OCI_ENDPOINT_ID")
            if not endpoint_id:
                raise ValueError(
                    "oci_endpoint_id/OCI_ENDPOINT_ID is required when oci_serving_mode is DEDICATED"
                )
            return oci.generative_ai_inference.models.DedicatedServingMode(endpoint_id=endpoint_id)

        return oci.generative_ai_inference.models.OnDemandServingMode(model_id=model_id)

    def embedding(
        self,
        model: str,
        input: list,
        model_response: EmbeddingResponse,
        optional_params: dict,
        encoding=None,
        api_key: Optional[str] = None,
        **kwargs,
    ) -> EmbeddingResponse:
        del encoding, api_key, kwargs  # not used by OCI embedding endpoint

        optional_params = optional_params or {}

        # model arrives as "oci-embed/cohere.embed-multilingual-v3.0"
        model_id = model.split("/", 1)[-1]

        compartment_id = self._pick_value(
            optional_params, "oci_compartment_id", "OCI_COMPARTMENT_ID"
        )
        if not compartment_id:
            raise ValueError("oci_compartment_id/OCI_COMPARTMENT_ID is required")

        input_type = self._pick_value(
            optional_params,
            "input_type",
            "OCI_INPUT_TYPE",
            default="SEARCH_DOCUMENT",
        ).upper()

        truncate = self._pick_value(
            optional_params,
            "truncate",
            "OCI_EMBED_TRUNCATE",
            default="NONE",
        ).upper()

        inputs = input if isinstance(input, list) else [input]

        detail = oci.generative_ai_inference.models.EmbedTextDetails(
            inputs=inputs,
            serving_mode=self._build_serving_mode(model_id, optional_params),
            compartment_id=compartment_id,
            input_type=input_type,
            truncate=truncate,
        )

        resp = self._get_client(optional_params).embed_text(detail)

        model_response.model = model_id
        model_response.data = [
            Embedding(object="embedding", index=i, embedding=emb)
            for i, emb in enumerate(resp.data.embeddings)
        ]
        model_response.usage = Usage(
            prompt_tokens=resp.data.usage.prompt_tokens if resp.data.usage else 0,
            completion_tokens=0,
            total_tokens=resp.data.usage.total_tokens if resp.data.usage else 0,
        )

        return model_response

    async def aembedding(self, *args, **kwargs) -> EmbeddingResponse:
        # LiteLLM proxy is async; OCI SDK call is sync.
        return self.embedding(*args, **kwargs)


oci_embedding_llm = OCIEmbeddingLLM()