trade-assistant/backend/app/api/v1/discovery.py

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="生成失败，请稍后重试")