from __future__ import annotations

from .shared import *

class LLMGenerateEntryMixin:
    async def generate(self, prompt: str, **kwargs: Any) -> str:
        call_kwargs = dict(kwargs)
        db = call_kwargs.pop("db", None) or call_kwargs.pop("db_session", None)
        if db is None:
            async with AsyncSessionLocal() as session:
                return await self._generate_with_db(session, prompt, **call_kwargs)
        return await self._generate_with_db(db, prompt, **call_kwargs)

    async def generate_structured_object(self, prompt: str, **kwargs: Any) -> StructuredExtractionResult:
        call_kwargs = dict(kwargs)
        db = call_kwargs.pop("db", None) or call_kwargs.pop("db_session", None)
        if db is None:
            async with AsyncSessionLocal() as session:
                return await self._generate_structured_with_db(session, prompt, **call_kwargs)
        return await self._generate_structured_with_db(db, prompt, **call_kwargs)

    async def _inject_model_context_bundle_prompt(
        self,
        *,
        db: Any,
        tenant_id: int,
        function_name: str,
        prompt: str,
    ) -> str:
        raw_prompt = str(prompt or "")
        normalized_function = str(function_name or "").strip().lower()
        if not ModelContextBundleService.should_inject_context(normalized_function):
            return raw_prompt
        # 单测里可能传入 object()/mock 作为 db 占位，此时直接跳过注入。
        if not hasattr(db, "execute"):
            return raw_prompt
        try:
            return await ModelContextBundleService.inject_prompt_with_bundle(
                db,
                tenant_id=int(tenant_id),
                function_name=normalized_function,
                prompt=raw_prompt,
            )
        except Exception as exc:
            _runtime_log("LLM上下文", f"上下文注入失败，已降级原始 prompt：function={normalized_function}, error={exc}")
            return raw_prompt