# app.py
# pip install gradio==4.43.0 pandas requests unidecode

import requests
import re
import random
import gradio as gr
import pandas as pd
from difflib import SequenceMatcher
from unidecode import unidecode
import os

# -----------------------------
# 工具函数
# -----------------------------
def normalize(s: str) -> str:
    s = s.strip().lower()
    s = unidecode(s)
    s = re.sub(r"\s+", " ", s)
    return s

def dedup_words(words):
    out, seen = [], set()
    for w in words:
        k = normalize(w)
        if k not in seen:
            seen.add(k)
            out.append(w)
    return out

def similarity(a: str, b: str) -> float:
    return SequenceMatcher(None, normalize(a), normalize(b)).ratio()

# -----------------------------
# TikTok Cookie 支持
# -----------------------------
TIKTOK_COOKIE = ""

def set_cookie(cookie_text):
    global TIKTOK_COOKIE
    TIKTOK_COOKIE = cookie_text
    return "✅ Cookie 已更新"

# -----------------------------
# 1. TikTok 实时热搜抓取
# -----------------------------
def get_tiktok_trending_keywords(region="VN", query=None, topk=12):
    if query:
        url = f"https://www.tiktok.com/api/discover/search/?keyword={query}&region={region}"
    else:
        url = f"https://www.tiktok.com/api/discover/trending/?region={region}"

    headers = {
        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/114.0.0.0 Safari/537.36",
        "Referer": "https://www.tiktok.com/",
    }
    if TIKTOK_COOKIE:
        headers["Cookie"] = TIKTOK_COOKIE

    try:
        r = requests.get(url, headers=headers, timeout=8)
        r.raise_for_status()
        data = r.json()
    except Exception as e:
        print("请求失败：", e)
        return []

    keywords = []
    for item in data.get("data", []):
        title = item.get("title") or item.get("challenge_info", {}).get("title")
        views = item.get("viewCount") or 0
        if title:
            keywords.append((title.strip(), views))

    keywords.sort(key=lambda x: x[1], reverse=True)
    return [kw for kw, _ in keywords[:topk]]

# -----------------------------
# 2. 热词建议
# -----------------------------
def suggest_hotwords(product_name: str, category: str, lang: str, topk: int = 12):
    region_map = {"zh": "CN", "fr": "FR", "en": "US", "vi": "VN"}
    region = region_map.get(lang, "US")
    return get_tiktok_trending_keywords(region=region, query=product_name, topk=topk)

# -----------------------------
# 3. 标题生成
# -----------------------------
def rank_keywords(core_words, hotwords):
    items = []
    for w in core_words:
        items.append((w, 1.0))
    for i, w in enumerate(hotwords):
        items.append((w, 0.9 - i * 0.02))
    uniq = {}
    for w, sc in items:
        k = normalize(w)
        if k not in uniq or sc > uniq[k]:
            uniq[k] = sc
    ranked = sorted([(w, sc) for w, sc in uniq.items()], key=lambda x: x[1], reverse=True)
    return [w for w, _ in ranked]

def build_title(core, ranked, max_len):
    segments = dedup_words(core + ranked)
    title = " ".join(segments)
    if len(title) > max_len:
        title = title[:max_len]
    return title

def generate_titles(product_name, category, lang, n_titles, hotwords_text, max_len):
    core_words = dedup_words(re.split(r"[,\|/;，、 ]+", product_name))
    hotwords = dedup_words(re.split(r"[,\|/;，、\n]+", hotwords_text))
    ranked = rank_keywords(core_words, hotwords)

    out = []
    for i in range(n_titles):
        random.shuffle(ranked)
        title = build_title(core_words, ranked, max_len)
        out.append(title)

    df = pd.DataFrame({"序号": list(range(1, len(out) + 1)), "标题": out})
    csv_bytes = df.to_csv(index=False).encode("utf-8-sig")
    return "\n".join(out), csv_bytes

# -----------------------------
# 4. Gradio UI
# -----------------------------
LANGS = [("中文", "zh"), ("Français", "fr"), ("English", "en"), ("Tiếng Việt", "vi")]

def ui_hotwords(product_name, category, lang, topk):
    if not product_name:
        return ""
    kws = suggest_hotwords(product_name, category, lang, int(topk))
    return ", ".join(kws)

with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("## TikTok 商品标题优化工具（实时热搜 + Cookie 支持）")

    with gr.Row():
        product_name = gr.Textbox(label="商品名称与关键词", placeholder="输入商品名，例如：阻力带", lines=2)
        category = gr.Textbox(label="类目", value="健身用品/健身器材")

    with gr.Row():
        lang = gr.Dropdown(LANGS, value="fr", label="语言")
        n_titles = gr.Slider(1, 10, value=3, step=1, label="生成数量")
        max_len = gr.Slider(30, 120, value=80, step=1, label="最大标题长度")
        topk = gr.Slider(6, 30, value=12, step=1, label="热门关键词数量")

    hotwords_box = gr.Textbox(label="热门关键词（实时获取）", lines=2, placeholder="将自动填充")
    btn_refresh = gr.Button("⟳ 获取实时热门关键词")
    btn_generate = gr.Button("⚡ 生成标题")

    cookie_input = gr.Textbox(label="TikTok Cookie（可选）", placeholder="可手动粘贴 Cookie", lines=2)
    btn_set_cookie = gr.Button("✅ 更新 Cookie")

    titles_out = gr.Textbox(label="生成结果", lines=10)
    csv_file = gr.File(label="下载 CSV", file_count="single")

    product_name.change(ui_hotwords, [product_name, category, lang, topk], [hotwords_box])
    lang.change(ui_hotwords, [product_name, category, lang, topk], [hotwords_box])
    topk.change(ui_hotwords, [product_name, category, lang, topk], [hotwords_box])
    btn_refresh.click(ui_hotwords, [product_name, category, lang, topk], [hotwords_box])

    btn_generate.click(
        generate_titles,
        [product_name, category, lang, n_titles, hotwords_box, max_len],
        [titles_out, csv_file]
    )

    btn_set_cookie.click(set_cookie, [cookie_input], [cookie_input])

if __name__ == "__main__":
    demo.launch()