from typing import List, Literal, cast
from pydantic import SecretStr
from _utils.google_integration.google_cloud import GCP_PROJECT, upload_to_gcs
from setup.easy_imports import ChatOpenAI, ChatGoogleGenerativeAI
import os
from langchain_core.messages import HumanMessage
from langchain_google_vertexai import ChatVertexAI

deepseek_api_key = cast(str, os.environ.get("DEEPSEEKK_API_KEY"))
google_api_key = cast(str, os.environ.get("GOOGLE_API_KEY_PEIXE"))
open_ai_token = cast(str, os.environ.get("OPENAI_API_KEY"))

Google_llms = Literal[
    "gemini-2.5-pro-preview-05-06",
    "gemini-2.0-flash",
    "gemini-2.0-flash-lite",
    "gemini-2.5-flash-preview-04-17",
]


class LLM:
    def __init__(self):
        pass

    def open_ai(self, model="gpt-4o-mini"):
        return ChatOpenAI(api_key=SecretStr(open_ai_token), model=model)

    def deepseek(self, model="deepseek-chat"):
        return ChatOpenAI(
            api_key=SecretStr(deepseek_api_key),
            base_url="https://api.deepseek.com/v1",
            model=model,
        )

    def google_gemini(self, model: Google_llms = "gemini-2.0-flash", temperature=0.4):
        return ChatGoogleGenerativeAI(
            api_key=SecretStr(google_api_key),
            model=model,
            temperature=temperature,
            max_tokens=None,
            timeout=None,
            max_retries=2,
        )

    async def google_gemini_ainvoke(
        self,
        prompt: str,
        model: Google_llms = "gemini-2.0-flash",
        max_retries: int = 3,
        temperature=0.4,
    ):
        for attempt in range(max_retries):
            try:
                response = await self.google_gemini(model, temperature).ainvoke(
                    [HumanMessage(content=prompt)]
                )

                if isinstance(response.content, list):
                    response.content = "\n".join(response.content)  # type: ignore

                return response
            except Exception as e:
                model = "gemini-2.0-flash"
                print(f"Attempt {attempt + 1} failed with error: {e}")

        # Final attempt fallback logic (optional)
        try:
            print("Final attempt with fallback model...")
            response = await self.open_ai("chat-gpt-4o-mini").ainvoke(
                [HumanMessage(content=prompt)]
            )
            return response
        except Exception as e:
            raise Exception(
                "Failed to generate the final document after 5 retries and the fallback attempt with chat-gpt-4o-mini."
            ) from e

    async def google_gemini_vertex_ainvoke(
        self,
        prompt: str,
        list_of_pdfs: List[str],
        model: Google_llms = "gemini-2.5-flash-preview-04-17",
        max_retries: int = 3,
    ) -> str | None:
        message_parts = [
            {"type": "text", "text": prompt},
        ]
        for pdf in list_of_pdfs:
            pdf_gcs_uri = upload_to_gcs(pdf)
            message_parts.append(
                {
                    # This structure is used for file references via URI
                    "type": "media",
                    "mime_type": "application/pdf",  # <-- mime_type moved up
                    "file_uri": pdf_gcs_uri,  # <-- file_uri moved up
                }
            )

        for attempt in range(max_retries):
            try:
                llm = ChatVertexAI(
                    model_name=model,
                    project=GCP_PROJECT,
                    location="us-central1",
                    temperature=0,
                )
                response = await llm.ainvoke(
                    [HumanMessage(content=message_parts)]  # type: ignore
                )

                if isinstance(response.content, list):
                    response.content = "\n".join(response.content)  # type: ignore

                return response.content  # type: ignore
            except Exception as e:
                model = "gemini-2.0-flash"
                print(f"Attempt {attempt + 1} failed with error: {e}")