from typing import Annotated, Any, TypedDict

from langchain_core.messages import BaseMessage
from langgraph.graph.message import add_messages
from pydantic import BaseModel

import config as global_config


class QueryRequest(BaseModel):
    """API Request model cho CuCu Note Chat"""

    user_id: str | None = None
    user_query: str
    images: list[str] | None = None
    image_analysis: dict[str, Any] | None = None


class AgentState(TypedDict):
    """Trạng thái của Agent trong LangGraph."""

    user_query: BaseMessage
    history: list[BaseMessage]
    user_id: str | None
    ai_response: BaseMessage | None
    images_embedding: list[str] | None
    messages: Annotated[list[BaseMessage], add_messages]


class AgentConfig:
    """Class chứa cấu hình runtime cho Agent."""

    def __init__(self, **kwargs):
        self.model_name = kwargs.get("model_name") or global_config.DEFAULT_MODEL
        self.openai_api_key = kwargs.get("openai_api_key")
        self.google_api_key = kwargs.get("google_api_key")
        self.groq_api_key = kwargs.get("groq_api_key")
        self.supabase_url = kwargs.get("supabase_url")
        self.supabase_key = kwargs.get("supabase_key")
        self.langfuse_public_key = kwargs.get("langfuse_public_key")
        self.langfuse_secret_key = kwargs.get("langfuse_secret_key")
        self.langfuse_base_url = kwargs.get("langfuse_base_url")


def get_config() -> AgentConfig:
    """Khởi tạo cấu hình Agent từ các biến môi trường."""
    return AgentConfig(
        model_name=global_config.DEFAULT_MODEL,
        openai_api_key=global_config.OPENAI_API_KEY,
        google_api_key=global_config.GOOGLE_API_KEY,
        groq_api_key=global_config.GROQ_API_KEY,
        supabase_url=global_config.AI_SUPABASE_URL,
        supabase_key=global_config.AI_SUPABASE_KEY,
        langfuse_public_key=global_config.LANGFUSE_PUBLIC_KEY,
        langfuse_secret_key=global_config.LANGFUSE_SECRET_KEY,
        langfuse_base_url=global_config.LANGFUSE_BASE_URL,
    )
