# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Single transport for all Artemis models: OpenAI Responses API."""

from enum import StrEnum
import hashlib
import os
from typing import Any

from langchain_core.language_models.chat_models import BaseChatModel
from pydantic import BaseModel, Field

from artemis.config.settings import settings


class ModelProvider(StrEnum):
    OPENAI = "openai"

    @classmethod
    def from_string(cls, val: Any) -> "ModelProvider":
        if val is None or str(val).strip().lower() in ("", "openai"):
            return cls.OPENAI
        raise ValueError(f"Unsupported LLM provider {val!r}; only OpenAI Responses is supported")


class ModelEndpoint(BaseModel):
    provider: ModelProvider = ModelProvider.OPENAI
    model_name: str = Field(default_factory=lambda: settings.OPENAI_MODEL)
    api_key: str | None = Field(default=None, repr=False)
    api_base: str | None = None
    temperature: float | None = None
    max_tokens: int | None = Field(default=None, gt=0)
    timeout_seconds: float = Field(default_factory=lambda: settings.OPENAI_TIMEOUT_SECONDS, gt=0)
    max_retries: int = Field(default_factory=lambda: settings.OPENAI_MAX_RETRIES, ge=0)
    is_multimodal: bool = True
    reasoning_effort: str | None = Field(default_factory=lambda: settings.OPENAI_REASONING_EFFORT)
    enable_grounding: bool = False

    def cache_key(self) -> tuple:
        key = self.api_key or (
            settings.OPENAI_API_KEY.get_secret_value() if settings.OPENAI_API_KEY else ""
        )
        return (
            self.provider.value,
            self.model_name,
            self.temperature,
            self.max_tokens,
            self.timeout_seconds,
            self.max_retries,
            self.reasoning_effort,
            self.enable_grounding,
            self.api_base or settings.OPENAI_BASE_URL,
            settings.OPENAI_USE_SYSTEM_PROXY,
            hashlib.sha256(key.encode()).hexdigest(),
        )


class ModelFactory:
    _cache: dict[tuple, BaseChatModel] = {}

    @classmethod
    def get_model(cls, endpoint: ModelEndpoint) -> BaseChatModel:
        key = endpoint.cache_key()
        if key not in cls._cache:
            cls._cache[key] = cls.create_model(endpoint)
        return cls._cache[key]

    @classmethod
    def create_model(cls, endpoint: ModelEndpoint) -> BaseChatModel:
        ModelProvider.from_string(endpoint.provider)
        if os.environ.get("ARTEMIS_FAKE_LLM") == "1":
            from artemis.llm.fake_model import FakeChatModel

            return FakeChatModel(
                delay_s=float(os.environ.get("ARTEMIS_FAKE_LLM_DELAY_S", "0") or 0)
            )

        from langchain_openai import ChatOpenAI
        from openai import DefaultAsyncHttpxClient, DefaultHttpxClient

        from artemis.llm.diagnostics import observe_request_headers, observe_response_headers
        from artemis.llm.http import model_http_options

        api_key = endpoint.api_key or (
            settings.OPENAI_API_KEY.get_secret_value() if settings.OPENAI_API_KEY else None
        )
        if not api_key:
            raise ValueError("请在环境变量或 .env 中配置 OPENAI_API_KEY")
        kwargs = {
            "model": endpoint.model_name,
            "api_key": api_key,
            "base_url": endpoint.api_base or settings.OPENAI_BASE_URL or None,
            "timeout": endpoint.timeout_seconds,
            "stream_chunk_timeout": endpoint.timeout_seconds,
            # Artemis owns the retry budget; SDK retries would multiply it.
            "max_retries": 0,
            "temperature": endpoint.temperature,
            "max_tokens": endpoint.max_tokens,
            "reasoning_effort": None
            if endpoint.reasoning_effort == "default"
            else endpoint.reasoning_effort,
            "use_responses_api": True,
            "output_version": "responses/v1",
            "store": False,
            # The explicit clients own proxy policy, including the synchronous path.
            "openai_proxy": "",
            "http_client": DefaultHttpxClient(
                timeout=endpoint.timeout_seconds, **model_http_options()
            ),
            "http_async_client": DefaultAsyncHttpxClient(
                timeout=endpoint.timeout_seconds,
                **model_http_options(),
                event_hooks={
                    "request": [observe_request_headers],
                    "response": [observe_response_headers],
                },
            ),
        }
        return ChatOpenAI(**{k: v for k, v in kwargs.items() if v is not None})