import google.generativeai as genai
import config
from .base import AIProvider

_INLINE_SYSTEM_PROMPT = (
    "You are an inline document assistant. "
    "The user will provide a markdown document and a specific instruction. "
    "Your response must contain ONLY the text to be inserted into the document — "
    "no preamble, no explanation, no meta-commentary, no markdown code fences unless "
    "the instruction specifically asks for them. "
    "Respond as if your output will be dropped directly into the middle of a document."
)

_DOCUMENT_SYSTEM_PROMPT = (
    "You are a document editing assistant. "
    "The user will provide a markdown document and a specific instruction. "
    "Apply the instruction to the entire document and return the full rewritten document. "
    "Your response must contain ONLY the rewritten document — "
    "no preamble, no explanation, no meta-commentary, no markdown code fences. "
    "Preserve the document's structure and formatting unless the instruction says otherwise."
)


class GeminiProvider(AIProvider):
    def __init__(self) -> None:
        genai.configure(api_key=config.GEMINI_API_KEY)
        self._inline_model = genai.GenerativeModel(
            model_name="gemini-2.5-flash-lite",
            system_instruction=_INLINE_SYSTEM_PROMPT,
        )
        self._document_model = genai.GenerativeModel(
            model_name="gemini-2.5-flash-lite",
            system_instruction=_DOCUMENT_SYSTEM_PROMPT,
        )

    def complete(self, document: str, instruction: str) -> str:
        prompt = (
            f"DOCUMENT:\n\n{document}\n\n"
            f"INSTRUCTION: {instruction}"
        )
        response = self._inline_model.generate_content(prompt)
        return response.text.strip()

    def complete_document(self, document: str, instruction: str) -> str:
        prompt = (
            f"DOCUMENT:\n\n{document}\n\n"
            f"INSTRUCTION: {instruction}"
        )
        response = self._document_model.generate_content(prompt)
        return response.text.strip()