import gradio as gr
import torch
import random
import time
import re
import json
import threading
import collections

# 彻底屏蔽 CUDA，防止 ZeroGPU 拦截器超时
torch.cuda.is_available = lambda: False
torch.cuda.is_initialized = lambda: False

from transformers import AutoModelForCausalLM, AutoTokenizer

# 1. 载入原生希腊语因果语言模型 (Greek GPT-2)，用于物理模拟 iOS 键盘 3 候选词联想引擎
GREEK_MODEL_ID = "nikokons/gpt2-greek"
print(f"正在载入原生希腊语输入法基座模型 {GREEK_MODEL_ID}...")
greek_tokenizer = AutoTokenizer.from_pretrained(GREEK_MODEL_ID)
greek_model = AutoModelForCausalLM.from_pretrained(
    GREEK_MODEL_ID,
    torch_dtype=torch.float32,
    low_cpu_mem_usage=True
)
print("原生希腊语输入法基座模型装载完毕。")

# 2. 载入阿里工业级大语言模型 Qwen2.5-1.5B-Instruct 作为文学翻译与艺术润色中枢
LLM_MODEL_ID = "Qwen/Qwen2.5-1.5B-Instruct"
print(f"正在载入 1.5B 工业级文学理解与润色大模型 {LLM_MODEL_ID}...")
llm_tokenizer = AutoTokenizer.from_pretrained(LLM_MODEL_ID)
llm_model = AutoModelForCausalLM.from_pretrained(
    LLM_MODEL_ID,
    torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
    low_cpu_mem_usage=True
)
print("Qwen2.5-1.5B 文学润色中枢装载完毕。")

crypto_rand = random.SystemRandom()

print("正在载入希腊单字种子库 greek_64_seeds.json ...")
try:
    with open("greek_64_seeds.json", "r", encoding="utf-8") as f:
        GREEK_DICTIONARY = json.load(f)
    print(f"成功载入 {len(GREEK_DICTIONARY)} 个希腊单字种子。")
except Exception as e:
    print("加载失败，使用后备种子", e)
    GREEK_DICTIONARY = [
        {"single_word_seed": "Χάος", "zh_meaning": "原初混沌"},
        {"single_word_seed": "Ἀγάπη", "zh_meaning": "纯粹之爱"},
        {"single_word_seed": "Σοφία", "zh_meaning": "神性智慧"},
        {"single_word_seed": "Λόγος", "zh_meaning": "万物尺度"},
        {"single_word_seed": "Ψυχή", "zh_meaning": "灵质之蝶"},
        {"single_word_seed": "Ἀλήθεια", "zh_meaning": "非遗忘真实"},
        {"single_word_seed": "Τέλος", "zh_meaning": "终局因由"},
        {"single_word_seed": "Κόσμος", "zh_meaning": "宏观秩序"},
        {"single_word_seed": "Χρόνος", "zh_meaning": "流逝光阴"},
        {"single_word_seed": "Ἔρεβος", "zh_meaning": "幽冥永夜"},
        {"single_word_seed": "Μελανχολία", "zh_meaning": "蚀骨忧郁"},
        {"single_word_seed": "Αἷμα", "zh_meaning": "鲜血循环"},
        {"single_word_seed": "Ἔντερα", "zh_meaning": "腹腔深处"},
        {"single_word_seed": "Σῶμα", "zh_meaning": "肉身牢笼"},
        {"single_word_seed": "Κρύπτη", "zh_meaning": "隐秘地穴"},
        {"single_word_seed": "Νέκυια", "zh_meaning": "冥界召唤"},
        {"single_word_seed": "Ὀδύνη", "zh_meaning": "肉体剧痛"},
        {"single_word_seed": "Σιωπή", "zh_meaning": "死寂无声"},
        {"single_word_seed": "Ἄβυσσος", "zh_meaning": "无底深渊"},
        {"single_word_seed": "Φάσμα", "zh_meaning": "幽灵幻影"},
        {"single_word_seed": "Λήθη", "zh_meaning": "彻底遗忘"}
    ]

def simulate_ios_greek_keyboard(seed_word: str, max_tokens: int = 70) -> str:
    """
    物理模拟 iOS 希腊语键盘的 Top-3 联想输入引擎机制：
    1. 严格以 64 个希腊单字种子为起点；
    2. 每步在 Top 3 候选槽中以真实敲击概率（78% 中间 / 14% 左 / 8% 右）推进；
    3. 生成完全纯粹、未知的希腊语自回归联想长文本流。
    """
    device = "cpu"
    prompt = seed_word.strip()
    input_ids = greek_tokenizer.encode(prompt, return_tensors="pt").to(device)
    
    tap_bias = torch.tensor([0.78, 0.14, 0.08], device=device)
    
    for step in range(max_tokens):
        with torch.no_grad():
            outputs = greek_model(input_ids)
            logits = outputs.logits[:, -1, :].clone().squeeze(0)
            
            logits = logits / 0.95
            
            recent_ids = input_ids[0, -14:].tolist()
            for tid in set(recent_ids):
                if logits[tid] > 0:
                    logits[tid] /= 1.18
                else:
                    logits[tid] *= 1.18
            
            top3_vals, top3_indices = torch.topk(logits, k=3)
            top3_probs = torch.softmax(top3_vals, dim=-1)
            
            sample_weights = top3_probs * tap_bias
            if sample_weights.sum() == 0:
                chosen_idx = 0
            else:
                chosen_idx = torch.multinomial(sample_weights, num_samples=1).item()
            
            chosen_token_id = top3_indices[chosen_idx].view(1, 1)
            input_ids = torch.cat([input_ids, chosen_token_id], dim=1)
            
            if step >= 55:
                piece = greek_tokenizer.decode([chosen_token_id.item()])
                if any(p in piece for p in [".", ";", "!", chr(10), "…"]):
                    break
                    
    full_greek = greek_tokenizer.decode(input_ids[0].tolist(), skip_special_tokens=True)
    return full_greek.strip()

def translate_and_polish_greek(greek_text: str) -> str:
    """
    由 Qwen2.5-1.5B 忠实理解当前希腊语输入法生成的独特意象并进行现代存在主义散文重构：
    1. 严禁套用历史已知文章；
    2. 基于输入的希腊语原生内容生成 3~4 个连贯短句（65~110 字）；
    3. 自动补全自然标点符号（，。！？）；
    4. 精选 1~2 个最核心意象词加粗（**词汇**）；
    5. 纯净输出，绝对零打招呼、零AI导语。
    """
    if not greek_text or len(greek_text.strip()) < 3:
        return ""
        
    prompt = (
        "你是一位深谙存在主义与现代意识流文学的翻译家。请深入理解以下由希腊语输入法联想流生成的原始文本，以高度凝练、利落有力的现代文学散文笔触进行重构与翻译。\n"
        "要求：\n"
        "1. 严格根据当前输入的 Greek 文本原义与意象进行翻译与重构，严禁套用或抄袭任何固定历史故事；\n"
        "2. 输出 3~4 个连贯短句组成的饱满段落（字数在 65~110 字左右）；\n"
        "3. 以第一/第二人称（我、你、他们）贯穿，句式利落，充满未知探索与奇异张力；\n"
        "4. 自动补全完整的自然标点符号（，。！？）；\n"
        "5. 挑选 1~2 个最核心意象词添加 **加粗**；\n"
        "6. 绝对不要输出任何打招呼、导语、解释或前缀，直接输出纯净中文段落。\n\n"
        "Greek: Ήμουν μόνος στο σκοτάδι και περίμενα το φως. Δεν ήξερες ότι η πόλη είχε αλλάξει για πάντα. Περπατήσαμε μέχρι την άκρη της γέφυρας και κοιτάξαμε κάτω.\n"
        "Chinese: 我独自在黑暗中等待光线。你不知道这座**城市**已经永远改变了。我们走到了桥的边缘，向着深处俯瞰。风声切断了所有退路。\n\n"
        "Greek: Μου είπες να κοιτάξω τη θάλασσα και τα κύματα. Τα ίχνη μας χάνονταν στην υγρή άμμο. Δεν υπήρχε κανείς άλλος να μας απαντήσει εκείνη την ώρα.\n"
        "Chinese: 你让我注视着大海与翻滚的波浪。我们的**痕迹**正在潮湿的沙滩上消融。那个时刻，四下一片寂静，没有任何回音传来。\n\n"
        f"Greek: {greek_text.strip()}\n"
        "Chinese:"
    )
    
    inputs = llm_tokenizer(prompt, return_tensors="pt")
    
    with torch.no_grad():
        output_ids = llm_model.generate(
            **inputs,
            max_new_tokens=220,
            temperature=0.78,
            top_p=0.92,
            repetition_penalty=1.18,
            eos_token_id=llm_tokenizer.eos_token_id
        )
        
    gen_tokens = output_ids[0][inputs.input_ids.shape[1]:]
    chinese_text = llm_tokenizer.decode(gen_tokens, skip_special_tokens=True).strip()
    
    if "Greek:" in chinese_text:
        chinese_text = chinese_text.split("Greek:")[0].strip()
    if "\n\n" in chinese_text:
        chinese_text = chinese_text.split("\n\n")[0].strip()
        
    return chinese_text

def clean_output(text: str) -> str:
    """去除 HTML 标签与格式标记"""
    text = re.sub(r"</?[a-zA-Z0-9]+[^>]*>", "", text)
    text = text.replace("<|endoftext|>", "").replace("<|im_end|>", "").replace("<|im_start|>", "")
    text = re.sub(r"\[[0-9a-zA-Z\s,\.\-_:\x27\"]*\]", "", text)
    text = re.sub(r"\n{3,}", "\n\n", text)
    return text.strip()

def clean_chinese(text: str) -> str:
    """清洗希腊字母、外文字符与无意义符号，保留 Markdown **重点加粗** 与合法标点"""
    if not text:
        return ""
        
    # 1. 彻底清除所有残留的希腊字母与外文字母
    text = re.sub(r"[\u0370-\u03ff\u1f00-\u1fff]+", "", text)
    text = re.sub(r"[a-zA-Z]+", "", text)
    
    # 2. 清除方括号、花括号、分号、项目符号与无意义乱码符号（严格保留 * 用于 Markdown 重点词加粗）
    text = re.sub(r"\[[^\]]*\]", "", text)
    text = re.sub(r"【[^】]*】", "", text)
    text = re.sub(r"[\[\]【】\{\}\(\)\;\:\=\+\-\/\<\>\&\|\$\#\\•·~_\`]+", "", text)
    
    # 3. 将换行合并为连贯的长句
    text = text.replace("\n", " ")
    text = re.sub(r"\s+", "", text)
    text = re.sub(r"[。\.]{2,}", "。", text)
    text = re.sub(r"[，,]{2,}", "，", text)
    text = re.sub(r"[！!]{2,}", "！", text)
    text = re.sub(r"[？?]{2,}", "？", text)
    text = re.sub(r"^[，。！？、；：\s]+", "", text)
    text = text.strip()
    if text and text[-1] not in ["。", "！", "？", "”", "’"]:
        text += "。"
    return text

# 多线程预生成滑动缓冲池配置
TARGET_BUFFER_SIZE = 8
BUFFER_POOL = collections.deque(maxlen=16)
buffer_lock = threading.Lock()

# 全局同步状态
GLOBAL_STATE = {
    "server_start": time.time(),
    "blocks": [],
    "is_generating": False,
    "current_stage": "INITIALIZING",
    "progress": 0,
    "buffer_count": 0,
    "logs": ["[00:00:00] 系统启动，开启纯自回归希腊单字种子生成管线。"]
}
state_lock = threading.Lock()

def add_server_log(msg: str):
    timestamp_str = time.strftime("%H:%M:%S", time.localtime())
    entry = f"[{timestamp_str}] {msg}"
    print(entry)
    with state_lock:
        GLOBAL_STATE["logs"].append(entry)
        if len(GLOBAL_STATE["logs"]) > 25:
            GLOBAL_STATE["logs"] = GLOBAL_STATE["logs"][-25:]

def producer_worker():
    """
    后台独立预生成生产者线程：
    常驻从 64 个希腊单字种子库中随机抽取词源，
    通过 iOS 键盘 3 候选词自回归机制持续生成完全未知的原生段落。
    """
    add_server_log("预生成生产者线程启动，常驻从 64 希腊种子池自回归生成...")
    
    while True:
        with buffer_lock:
            current_pool_len = len(BUFFER_POOL)
            
        with state_lock:
            GLOBAL_STATE["buffer_count"] = current_pool_len
            
        if current_pool_len < TARGET_BUFFER_SIZE:
            with state_lock:
                GLOBAL_STATE["is_generating"] = True
                GLOBAL_STATE["current_stage"] = f"PREGENERATING ({current_pool_len}/{TARGET_BUFFER_SIZE})"
                GLOBAL_STATE["progress"] = 30

            try:
                entry = crypto_rand.choice(GREEK_DICTIONARY)
                seed = entry.get("single_word_seed", "Χάος")
                zh_m = entry.get("zh_meaning", "")
                
                # 1. iOS 键盘 3 候选词长文本纯自回归模拟
                raw_greek = simulate_ios_greek_keyboard(seed, max_tokens=70)
                raw_greek = clean_output(raw_greek)

                with state_lock:
                    GLOBAL_STATE["progress"] = 70

                # 2. Qwen2.5 基于真实输入的文学理解与意象重构
                translated = translate_and_polish_greek(raw_greek)
                cleaned = clean_chinese(translated)

                if cleaned and len(cleaned) >= 25:
                    new_block = {
                        "id": int(time.time() * 1000),
                        "seed": seed,
                        "text": cleaned,
                        "raw_text": raw_greek,
                        "timestamp": time.time()
                    }
                    with buffer_lock:
                        BUFFER_POOL.append(new_block)
                        new_pool_len = len(BUFFER_POOL)
                    add_server_log(f"自回归段落就绪: 《{seed}》 (池: {new_pool_len}/{TARGET_BUFFER_SIZE}, {len(cleaned)}字) -> {cleaned[:22]}...")
            except Exception as e:
                add_server_log(f"自回归生成异常: {e}")
            finally:
                with state_lock:
                    GLOBAL_STATE["is_generating"] = False
                    GLOBAL_STATE["progress"] = 100
            time.sleep(0.2)
        else:
            with state_lock:
                GLOBAL_STATE["current_stage"] = f"BUFFER_FULL ({current_pool_len}/{TARGET_BUFFER_SIZE})"
            time.sleep(0.4)

def publisher_worker():
    """
    平滑流式分发调度器：
    从预生成缓冲池中按 2.2 秒节奏向长卷供稿，彻底消灭前端等待。
    """
    add_server_log("流式分发调度器启动，以高频节奏向全局长卷供应自回归段落...")
    
    while True:
        with buffer_lock:
            has_buffer = len(BUFFER_POOL) > 0
            if has_buffer:
                ready_block = BUFFER_POOL.popleft()
                remaining = len(BUFFER_POOL)
            else:
                ready_block = None
                remaining = 0
                
        if ready_block:
            with state_lock:
                GLOBAL_STATE["blocks"].append(ready_block)
                if len(GLOBAL_STATE["blocks"]) > 35:
                    GLOBAL_STATE["blocks"] = GLOBAL_STATE["blocks"][-35:]
                GLOBAL_STATE["buffer_count"] = remaining
                GLOBAL_STATE["current_stage"] = f"STREAM_DISPATCHED (Buffer: {remaining})"
            add_server_log(f"分发新鲜自回归段落 ({len(ready_block['text'])}字): {ready_block['text'][:24]}...")
            time.sleep(2.2)
        else:
            time.sleep(0.2)

# 启动后台多线程流水线
t_producer = threading.Thread(target=producer_worker, daemon=True)
t_publisher = threading.Thread(target=publisher_worker, daemon=True)
t_producer.start()
t_publisher.start()

def get_state():
    with state_lock:
        clean_blocks = [
            {k: v for k, v in b.items() if k != "raw_text"}
            for b in GLOBAL_STATE["blocks"]
        ]
        return json.dumps({
            "server_time": time.time(),
            "blocks": clean_blocks,
            "is_generating": GLOBAL_STATE["is_generating"],
            "current_stage": GLOBAL_STATE["current_stage"],
            "progress": GLOBAL_STATE["progress"],
            "buffer_count": GLOBAL_STATE["buffer_count"],
            "logs": GLOBAL_STATE["logs"]
        }, ensure_ascii=False)

def manual_trigger():
    return json.dumps({"status": "running"})

import spaces
@spaces.GPU
def dummy_gpu_func():
    pass

with gr.Blocks(title="24H Global Sync Stream") as demo:
    btn_trigger = gr.Button("Trigger", visible=False)
    out_trigger = gr.Textbox(visible=False)
    btn_trigger.click(fn=manual_trigger, inputs=[], outputs=[out_trigger], api_name="trigger")

    btn_state = gr.Button("State", visible=False)
    out_state = gr.Textbox(visible=False)
    btn_state.click(fn=get_state, inputs=[], outputs=[out_state], api_name="state")

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