"""
Jina Search Agent - 使用 Jina Reader API 提取网页内容
"""
import os
import sys
import logging
import aiohttp
from typing import Optional, List
from datetime import datetime

from fastapi import FastAPI, HTTPException, Query
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
import uvicorn

# 添加 common 模块路径
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# 导入回调工具
try:
    from common.agent_callback_utils import AgentCallbackHandler, CallbackContextManager
    CALLBACK_ENABLED = True
except ImportError:
    CALLBACK_ENABLED = False
    AgentCallbackHandler = None
    CallbackContextManager = None

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# 环境变量
SERVICE_HOST = os.getenv("SERVICE_HOST", "0.0.0.0")
SERVICE_PORT = int(os.getenv("SERVICE_PORT", "8080"))
POD_NAME = os.getenv("POD_NAME", "jina-search-agent")
USER_ID = os.getenv("USER_ID", "")
JINA_API_KEY = os.getenv("JINA_API_KEY", "")

# Jina Reader API
JINA_READER_URL = "https://r.jina.ai/"

# FastAPI 应用
app = FastAPI(
    title="Jina Search Agent",
    description="使用 Jina Reader API 提取网页内容",
    version="1.0.0"
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 回调处理器
callback_handler: Optional[AgentCallbackHandler] = None


# ==================== 请求/响应模型 ====================

class SearchRequest(BaseModel):
    """搜索/提取请求"""
    url: str = Field(..., description="要提取内容的 URL")
    user_id: Optional[str] = Field(None, description="用户ID（用于计费回调）")
    jina_api_key: Optional[str] = Field(None, description="Jina API 密钥（可选，覆盖环境变量）")


class SearchResponse(BaseModel):
    """搜索/提取响应"""
    url: str
    title: Optional[str] = None
    content: str
    timestamp: str


class BatchSearchRequest(BaseModel):
    """批量搜索请求"""
    urls: List[str] = Field(..., description="要提取内容的 URL 列表")
    user_id: Optional[str] = Field(None, description="用户ID（用于计费回调）")
    jina_api_key: Optional[str] = Field(None, description="Jina API 密钥（可选）")


class BatchSearchResponse(BaseModel):
    """批量搜索响应"""
    results: List[SearchResponse]
    success_count: int
    failed_count: int
    timestamp: str


class HealthResponse(BaseModel):
    """健康检查响应"""
    status: str
    pod_name: str
    jina_api_configured: bool
    callback_enabled: bool
    timestamp: str


# ==================== 生命周期 ====================

@app.on_event("startup")
async def startup_event():
    """应用启动时初始化回调处理器"""
    global callback_handler
    
    if CALLBACK_ENABLED:
        callback_handler = AgentCallbackHandler(
            agent_name=POD_NAME,
            user_id=USER_ID
        )
        logger.info(f"回调处理器已初始化: callback_url={callback_handler.callback_url}")
    else:
        logger.warning("回调模块未加载，计费回调功能不可用")


# ==================== 辅助函数 ====================

async def fetch_url_content(url: str, api_key: str) -> dict:
    """使用 Jina Reader API 提取 URL 内容"""
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Accept": "application/json"
    }
    
    jina_url = f"{JINA_READER_URL}{url}"
    
    try:
        async with aiohttp.ClientSession() as session:
            async with session.get(jina_url, headers=headers, timeout=30) as response:
                if response.status == 200:
                    text = await response.text()
                    return {
                        "success": True,
                        "url": url,
                        "content": text,
                        "title": None  # Jina Reader 返回的是纯文本
                    }
                else:
                    error_text = await response.text()
                    logger.warning(f"Jina Reader 请求失败 [{response.status}]: {url}")
                    return {
                        "success": False,
                        "url": url,
                        "error": f"HTTP {response.status}: {error_text[:200]}"
                    }
    except Exception as e:
        logger.error(f"提取内容失败: {url} - {e}")
        return {
            "success": False,
            "url": url,
            "error": str(e)
        }


# ==================== API 端点 ====================

@app.get("/health", response_model=HealthResponse)
@app.get("/", response_model=HealthResponse)
async def health_check():
    """健康检查"""
    return HealthResponse(
        status="healthy",
        pod_name=POD_NAME,
        jina_api_configured=bool(JINA_API_KEY),
        callback_enabled=CALLBACK_ENABLED,
        timestamp=datetime.utcnow().isoformat()
    )


@app.post("/search", response_model=SearchResponse)
@app.post("/fetch", response_model=SearchResponse)
async def fetch_content(request: SearchRequest):
    """提取单个 URL 的内容"""
    api_key = request.jina_api_key or JINA_API_KEY
    
    if not api_key:
        raise HTTPException(
            status_code=400,
            detail="Jina API key 未设置。请在请求中传入 jina_api_key 或设置环境变量 JINA_API_KEY"
        )
    
    # 使用回调上下文管理器
    if CALLBACK_ENABLED and callback_handler and request.user_id:
        with CallbackContextManager(
            handler=callback_handler,
            user_id=request.user_id,
            request_id=f"jina-fetch-{int(datetime.utcnow().timestamp())}"
        ) as ctx:
            ctx.add_tool("jina_reader")
            ctx.add_tool("web_content_extraction")
            
            result = await fetch_url_content(request.url, api_key)
            
            if not result["success"]:
                raise HTTPException(
                    status_code=500,
                    detail=f"提取内容失败: {result.get('error', '未知错误')}"
                )
            
            return SearchResponse(
                url=result["url"],
                title=result.get("title"),
                content=result["content"],
                timestamp=datetime.utcnow().isoformat()
            )
    else:
        result = await fetch_url_content(request.url, api_key)
        
        if not result["success"]:
            raise HTTPException(
                status_code=500,
                detail=f"提取内容失败: {result.get('error', '未知错误')}"
            )
        
        return SearchResponse(
            url=result["url"],
            title=result.get("title"),
            content=result["content"],
            timestamp=datetime.utcnow().isoformat()
        )


@app.get("/search")
@app.get("/fetch")
async def fetch_content_get(
    url: str = Query(..., description="要提取内容的 URL"),
    user_id: Optional[str] = Query(None, description="用户ID（用于计费回调）"),
    jina_api_key: Optional[str] = Query(None, description="Jina API 密钥")
):
    """GET 方式提取内容"""
    request = SearchRequest(url=url, user_id=user_id, jina_api_key=jina_api_key)
    return await fetch_content(request)


@app.post("/batch", response_model=BatchSearchResponse)
async def batch_fetch_content(request: BatchSearchRequest):
    """批量提取多个 URL 的内容"""
    api_key = request.jina_api_key or JINA_API_KEY
    
    if not api_key:
        raise HTTPException(
            status_code=400,
            detail="Jina API key 未设置"
        )
    
    # 使用回调上下文管理器
    if CALLBACK_ENABLED and callback_handler and request.user_id:
        with CallbackContextManager(
            handler=callback_handler,
            user_id=request.user_id,
            request_id=f"jina-batch-{int(datetime.utcnow().timestamp())}"
        ) as ctx:
            ctx.add_tool("jina_reader")
            ctx.add_tool("batch_web_extraction")
            
            results = []
            success_count = 0
            failed_count = 0
            
            for url in request.urls:
                result = await fetch_url_content(url, api_key)
                if result["success"]:
                    results.append(SearchResponse(
                        url=result["url"],
                        title=result.get("title"),
                        content=result["content"],
                        timestamp=datetime.utcnow().isoformat()
                    ))
                    success_count += 1
                else:
                    failed_count += 1
            
            return BatchSearchResponse(
                results=results,
                success_count=success_count,
                failed_count=failed_count,
                timestamp=datetime.utcnow().isoformat()
            )
    else:
        results = []
        success_count = 0
        failed_count = 0
        
        for url in request.urls:
            result = await fetch_url_content(url, api_key)
            if result["success"]:
                results.append(SearchResponse(
                    url=result["url"],
                    title=result.get("title"),
                    content=result["content"],
                    timestamp=datetime.utcnow().isoformat()
                ))
                success_count += 1
            else:
                failed_count += 1
        
        return BatchSearchResponse(
            results=results,
            success_count=success_count,
            failed_count=failed_count,
            timestamp=datetime.utcnow().isoformat()
        )


# ==================== 主入口 ====================

def main():
    """主函数"""
    logger.info(f"启动 Jina Search Agent - {POD_NAME}")
    logger.info(f"Jina API Key: {'已配置' if JINA_API_KEY else '未配置'}")
    logger.info(f"回调功能: {'已启用' if CALLBACK_ENABLED else '未启用'}")
    
    uvicorn.run(app, host=SERVICE_HOST, port=SERVICE_PORT, log_level="info")


if __name__ == "__main__":
    main()