Crawl/backend/app/routers/public.py

from __future__ import annotations

from datetime import datetime
from typing import List, Optional, Union
from urllib.parse import urlparse

import json
from app.utils.price_parser import parse_prices

from fastapi import APIRouter, HTTPException, Request
from pydantic import BaseModel

from app.db import get_pool
from app.services.geo import geo_resolver

router = APIRouter()


class PublicPriceOut(BaseModel):
    url: str
    model_name: str
    prices: dict
    model_info: Optional[dict] = None
    rate_limits: Optional[dict] = None
    tool_prices: Optional[list] = None
    icon: Optional[str] = None
    scraped_at: datetime


class ParsedPriceItem(BaseModel):
    url: str
    model_name: str
    tier_min: Optional[float] = None
    tier_max: Optional[float] = None
    tier_unit: Optional[str] = None
    input_price: Optional[float] = None
    output_price: Optional[float] = None
    currency: str = "CNY"
    unit: Optional[str] = None
    label: Optional[str] = None


class DiscountedPriceItem(BaseModel):
    url: str
    model_name: str
    tier_min: Optional[float] = None
    tier_max: Optional[float] = None
    tier_unit: Optional[str] = None
    input_price: Optional[float] = None
    output_price: Optional[float] = None
    currency: str = "CNY"
    unit: Optional[str] = None
    label: Optional[str] = None
    discount: Optional[float] = None  # None 表示无折扣（原价）



class ModelTypeItem(BaseModel):
    model_name: str
    type: List[str]


class PricesResponse(BaseModel):
    version: int
    models: List[PublicPriceOut]
    parsed_prices: List[ParsedPriceItem]
    discounted_prices: List[DiscountedPriceItem]
    types: List[ModelTypeItem]
    discount: float = 1.0


class UpToDateResponse(BaseModel):
    up_to_date: bool = True
    version: int


def _extract_domain(referer: Optional[str]) -> Optional[str]:
    if not referer:
        return None
    try:
        return urlparse(referer).netloc or None
    except Exception:
        return None


@router.get("/prices", response_model=Union[PricesResponse, UpToDateResponse])
async def get_public_prices(
    request: Request,
    url: Optional[str] = None,
) -> Union[PricesResponse, UpToDateResponse]:
    pool = get_pool()

    # referer 必须提供
    referer = request.headers.get("referer") or request.headers.get("origin")
    if not referer:
        raise HTTPException(status_code=400, detail="Missing Referer header")

    # version 从 Header 读取，默认 0（首次请求）
    try:
        version = int(request.headers.get("version", "0") or "0")
    except ValueError:
        version = 0

    # 记录调用来源
    ip = request.client.host if request.client else "unknown"
    geo = geo_resolver.resolve(ip)
    try:
        await pool.execute(
            "INSERT INTO price_api_logs (ip, referer, org, country, city) VALUES ($1, $2, $3, $4, $5)",
            ip, referer, geo.org, geo.country, geo.city,
        )
    except Exception:
        pass

    # 查调用方域名对应的折扣
    caller_domain = _extract_domain(referer)
    discount_rate: Optional[float] = None
    if caller_domain:
        row = await pool.fetchrow("SELECT discount FROM discounts WHERE domain = $1", caller_domain)
        if row:
            discount_rate = float(row["discount"])

    def _j(v):
        if v is None:
            return None
        return v if isinstance(v, (dict, list)) else json.loads(v)

    # 读取版本号：优先用域名专属版本，没有则回退到全局版本
    if caller_domain:
        ver_row = await pool.fetchrow(
            "SELECT version FROM domain_version WHERE domain = $1", caller_domain
        )
        if ver_row:
            current_version = int(ver_row["version"])
        else:
            ver_row = await pool.fetchrow("SELECT version FROM price_snapshot_version WHERE id = 1")
            current_version = int(ver_row["version"]) if ver_row else 0
    else:
        ver_row = await pool.fetchrow("SELECT version FROM price_snapshot_version WHERE id = 1")
        current_version = int(ver_row["version"]) if ver_row else 0

    # version != 0 且与当前一致 → 无需更新（0 视为首次请求，强制返回数据）
    if version != 0 and version == current_version:
        return UpToDateResponse(up_to_date=True, version=current_version)

    # 从 price_snapshot 读取数据
    if url is None:
        rows = await pool.fetch(
            "SELECT url, model_name, prices, model_info, rate_limits, tool_prices, icon, updated_at FROM price_snapshot ORDER BY url"
        )
    else:
        rows = await pool.fetch(
            "SELECT url, model_name, prices, model_info, rate_limits, tool_prices, icon, updated_at FROM price_snapshot WHERE url = $1",
            url,
        )
        if not rows:
            raise HTTPException(status_code=404, detail="No price snapshot found for the given URL")

    if not rows:
        raise HTTPException(status_code=503, detail="Price snapshot not yet available")

    # version != 0 且与当前一致 → 无需更新
    if version != 0 and version == current_version:
        return UpToDateResponse(up_to_date=True, version=current_version)

    def _extract_type(model_info: Optional[dict]) -> Optional[List[str]]:
        if not model_info:
            return None
        tags = model_info.get("display_tags", [])
        TYPE_TAGS = {"文本生成", "图像生成", "视频生成", "向量表示", "向量模型", "多模态向量", "语音识别", "实时语音识别", "语音合成"}
        result = [t for t in tags if t in TYPE_TAGS]
        return result if result else None

    models = [PublicPriceOut(
        url=r["url"],
        model_name=r["model_name"],
        prices=_j(r["prices"]) or {},
        model_info=_j(r["model_info"]),
        rate_limits=_j(r["rate_limits"]),
        tool_prices=_j(r["tool_prices"]),
        icon=r["icon"],
        scraped_at=r["updated_at"],
    ) for r in rows]

    parsed_prices: List[ParsedPriceItem] = []
    discounted_prices: List[DiscountedPriceItem] = []

    for r in rows:
        for item in parse_prices(_j(r["prices"]) or {}):
            parsed_prices.append(ParsedPriceItem(url=r["url"], model_name=r["model_name"], **item))
            d_item = dict(item)
            if discount_rate is not None:
                if d_item.get("input_price") is not None:
                    d_item["input_price"] = round(d_item["input_price"] * discount_rate, 6)
                if d_item.get("output_price") is not None:
                    d_item["output_price"] = round(d_item["output_price"] * discount_rate, 6)
            discounted_prices.append(DiscountedPriceItem(url=r["url"], model_name=r["model_name"], discount=discount_rate, **d_item))

    all_types = [
        ModelTypeItem(model_name=r["model_name"], type=_extract_type(_j(r["model_info"])) or [])
        for r in rows
    ]

    return PricesResponse(
        version=current_version,
        models=models,
        parsed_prices=parsed_prices,
        discounted_prices=discounted_prices,
        types=all_types,
        discount=discount_rate if discount_rate is not None else 1.0,
    )