# mypy: disable - error - code = "no-untyped-def,misc"
from fastapi import FastAPI
from pydantic import BaseModel
from typing import Optional, List, Dict, Any
import asyncio

# Import the graph from the agent
from agent.graph import graph
# Import ChatVertexAI for translation
from langchain_google_vertexai import ChatVertexAI
import os

# Define the FastAPI app
app = FastAPI()

class ResearchRequest(BaseModel):
    """Request model for the research endpoint."""
    question: str
    initial_search_query_count: Optional[int] = 3
    max_research_loops: Optional[int] = 2
    reasoning_model: Optional[str] = "gemini-2.0-flash"


class ResearchResponse(BaseModel):
    """Response model for the research endpoint."""
    answer: str
    sources: List[Dict[str, Any]]


class TranslationRequest(BaseModel):
    """Request model for the translation endpoint."""
    text: str
    model_name: str


class TranslationResponse(BaseModel):
    """Response model for the translation endpoint."""
    original_text: str
    translated_text: str


@app.post("/api/research", response_model=ResearchResponse)
async def research(request: ResearchRequest):
    """Endpoint to perform research using the LangGraph agent.

    Args:
        request: ResearchRequest containing the question and optional parameters.

    Returns:
        ResearchResponse with the answer and sources.
   """
    # Prepare the input for the agent


    input_data = {
        "messages": [("user", request.question)]
    }

    # Prepare configuration
    config = {
        "configurable": {}
    }

    # Add optional parameters to configuration if provided
    if request.initial_search_query_count is not None:
        config["configurable"]["number_of_initial_queries"] = request.initial_search_query_count

    if request.max_research_loops is not None:
        config["configurable"]["max_research_loops"] = request.max_research_loops

    if request.reasoning_model is not None:
        config["configurable"]["answer_model"] = request.reasoning_model

    # Run the agent
    result = await asyncio.get_event_loop().run_in_executor(
        None,
        lambda: graph.invoke(input_data, config)
    )

    # Extract the answer and sources
    answer = result["messages"][-1].content if result["messages"] else ""
    sources = result.get("sources_gathered", [])

    return ResearchResponse(answer=answer, sources=sources)


@app.post("/api/translate", response_model=TranslationResponse)
async def translate(request: TranslationRequest):
    """Endpoint to translate Chinese text to English using ChatVertexAI.

    Args:
        request: TranslationRequest containing the text to translate.

    Returns:
        TranslationResponse with the original and translated text.
    """
    # Initialize ChatVertexAI
    llm = ChatVertexAI(
        model_name=request.model_name,
        temperature=0,
        max_retries=2,
        project=os.getenv("GOOGLE_CLOUD_PROJECT"),  # 必填
        location=os.getenv("GOOGLE_CLOUD_LOCATION")  # 必填
    )
    
    # Create translation prompt
    prompt = f"""你是一名跨境电商本地化专家。请根据用户输入的设备/商品名称及其描述(可选），准确提供出该设备/商品的英文名称。
            要求如下：
            1. **准确性 & 本地化**：提供符合目标市场（英语国家）消费者和电商平台习惯的标准名称，避免简单直译。
            2. **可搜索性**：提供的英文名称应是在主流电商平台（如Amazon, eBay）或专业网站上能精确检索到该商品/类目的常用名称。
            3. **输出格式**：*仅*输出一个有效的JSON对象，格式为：{{"English_names": array<string> }}。
            4. **数量**：`english_names`数组包含1个或多个最常用的英文名称。
            5. **用户输入**：{request.text}
            **重要指令**：禁止包含任何JSON之外的内容（如解释、评论、额外文本）。用户输入紧随此提示。
            """
    # Run the translation
    result = await asyncio.get_event_loop().run_in_executor(
        None,
        lambda: llm.invoke(prompt)
    )
    
    # Extract the translated text
    translated_text = result.content if hasattr(result, 'content') else str(result)
    
    return TranslationResponse(
        original_text=request.text,
        translated_text=translated_text
    )

