from fastapi import APIRouter, Depends, HTTPException from typing import Optional, Dict, Any from pydantic import BaseModel from sqlalchemy.ext.asyncio import AsyncSession from app.database import get_db from app.services.discovery import DiscoveryService from app.services.credit import CreditService from app.api.v1.deps import get_current_user_id import logging logger = logging.getLogger(__name__) router = APIRouter() class SearchRequest(BaseModel): product_description: str target_market: str = "US" class AnalyzeRequest(BaseModel): company_url: str product_description: str class MarketIntelRequest(BaseModel): product_description: str target_market: str = "US" class OutreachRequest(BaseModel): company: Dict[str, Any] product: Dict[str, Any] CREDIT_COST = { "search": 10, "analyze": 5, "outreach": 3, } async def _deduct_credits(user_id: str, result_type: str, db: AsyncSession): svc = CreditService(db) ok, balance = await svc.deduct(user_id, result_type) if not ok: raise HTTPException( status_code=402, detail=f"次数不足 (剩余 {balance:.1f}, 需要 {CREDIT_COST.get(result_type, 1)})" ) return balance @router.post("/search") async def search_leads( req: SearchRequest, user_id: str = Depends(get_current_user_id), db: AsyncSession = Depends(get_db), ): if not req.product_description.strip(): raise HTTPException(status_code=400, detail="请填写产品描述") credit_svc = CreditService(db) ok, balance = await credit_svc.deduct(user_id, "lead_search") if not ok: raise HTTPException( status_code=402, detail=f"次数不足 (剩余 {balance:.1f}, 需要 10)" ) svc = DiscoveryService(db=db) try: result = await svc.search(req.product_description, req.target_market) return {"success": True, "data": result, "credits_remaining": balance - 10} except Exception as e: await credit_svc.add_credits(user_id, 10, "refund", "搜索失败退回次数") logger.error(f"Search failed: {e}") raise HTTPException(status_code=500, detail="搜索失败,请稍后重试") @router.post("/analyze") async def analyze_company( req: AnalyzeRequest, user_id: str = Depends(get_current_user_id), db: AsyncSession = Depends(get_db), ): if not req.company_url.strip(): raise HTTPException(status_code=400, detail="请填写公司网址") if not req.product_description.strip(): raise HTTPException(status_code=400, detail="请填写产品描述") credit_svc = CreditService(db) ok, balance = await credit_svc.deduct(user_id, "company_analysis") if not ok: raise HTTPException( status_code=402, detail=f"次数不足 (剩余 {balance:.1f}, 需要 5)" ) svc = DiscoveryService() try: result = await svc.analyze(req.company_url, req.product_description) return {"success": True, "data": result, "credits_remaining": balance - 5} except Exception as e: await credit_svc.add_credits(user_id, 5, "refund", "分析失败退回次数") logger.error(f"Analysis failed: {e}") raise HTTPException(status_code=500, detail="分析失败,请稍后重试") @router.post("/market-intel") async def market_intel( req: MarketIntelRequest, user_id: str = Depends(get_current_user_id), db: AsyncSession = Depends(get_db), ): if not req.product_description.strip(): raise HTTPException(status_code=400, detail="请填写产品描述") credit_svc = CreditService(db) ok, balance = await credit_svc.deduct(user_id, "market_intel") if not ok: raise HTTPException( status_code=402, detail=f"次数不足 (剩余 {balance:.1f}, 需要 20)" ) svc = DiscoveryService(db=db) try: result = await svc.market_intel(req.product_description, req.target_market) return {"success": True, "data": result, "credits_remaining": balance - 20} except Exception as e: await credit_svc.add_credits(user_id, 20, "refund", "市场分析失败退回次数") logger.error(f"Market intel failed: {e}") raise HTTPException(status_code=500, detail="分析失败,请稍后重试") @router.post("/outreach") async def generate_outreach( req: OutreachRequest, user_id: str = Depends(get_current_user_id), db: AsyncSession = Depends(get_db), ): if not req.company.get("name"): raise HTTPException(status_code=400, detail="请填写公司名称") if not req.product.get("name"): raise HTTPException(status_code=400, detail="请填写产品名称") credit_svc = CreditService(db) ok, balance = await credit_svc.deduct(user_id, "outreach") if not ok: raise HTTPException( status_code=402, detail=f"次数不足 (剩余 {balance:.1f}, 需要 3)" ) svc = DiscoveryService() try: result = await svc.generate_outreach(req.company, req.product) return {"success": True, "data": result, "credits_remaining": balance - 3} except Exception as e: await credit_svc.add_credits(user_id, 3, "refund", "生成失败退回次数") logger.error(f"Outreach generation failed: {e}") raise HTTPException(status_code=500, detail="生成失败,请稍后重试")