import spaces
import gc
import os
from pathlib import Path

import gradio as gr
import numpy as np
import torch
from torch.nn import functional as F

# These must be set before importing rwkv.model. The Space is intentionally CPU-only.
os.environ["RWKV_V7_ON"] = "1"
os.environ["RWKV_JIT_ON"] = "1"
os.environ["RWKV_CUDA_ON"] = "0"

from rwkv.model import RWKV
from rwkv.rwkv_tokenizer import TRIE_TOKENIZER



BASE_DIR = Path(__file__).resolve().parent
MODEL_FILE = BASE_DIR / "RWKV-Desensitization.pth"
VOCAB_FILE = BASE_DIR / "rwkv_vocab_Desensitization.txt"
CTX_LIMIT = 4096
GEN_LIMIT = 4096
PENALTY_DECAY = 0.996

if not MODEL_FILE.is_file():
    raise FileNotFoundError(f"Model weights not found: {MODEL_FILE}")
if not VOCAB_FILE.is_file():
    raise FileNotFoundError(f"Tokenizer vocabulary not found: {VOCAB_FILE}")


########################## local RWKV model ####################################
# RWKV appends '.pth' itself when the suffix is omitted. Passing the local file
# stem also works with both the v7 and older rwkv package loaders.
model = RWKV(model=str(MODEL_FILE.with_suffix("")), strategy="cuda fp16")
tokenizer = TRIE_TOKENIZER(str(VOCAB_FILE))

model_vocab_size = int(model.args.vocab_size)
max_token_id = max(tokenizer.idx2token)
if max_token_id >= model_vocab_size:
    raise ValueError(
        f"Vocabulary token id {max_token_id} is outside the model vocabulary "
        f"size {model_vocab_size}."
    )
if tokenizer.encode("[REDACTED]") != [65530]:
    raise ValueError("The uploaded vocabulary does not contain the expected [REDACTED] token.")


def sample_logits(logits, temperature=1.0, top_p=0.7):
    """Sample from logits without constructing a second, unrelated tokenizer."""
    temperature = max(0.2, float(temperature))
    top_p = float(top_p)

    probs = F.softmax(logits.float(), dim=-1).cpu().numpy()
    if top_p <= 0:
        # A top-p value of zero is exposed as deterministic/greedy decoding.
        return int(np.argmax(probs))

    sorted_ids = np.argsort(probs)[::-1]
    sorted_probs = probs[sorted_ids]
    if top_p < 1.0:
        cumulative_probs = np.cumsum(sorted_probs)
        cutoff_index = int(np.argmax(cumulative_probs >= top_p))
        probs[probs < sorted_probs[cutoff_index]] = 0

    if temperature != 1.0:
        probs = probs ** (1.0 / temperature)
    total = probs.sum()
    if not np.isfinite(total) or total <= 0:
        return int(np.argmax(logits.float()).item())
    probs /= total
    return int(np.random.choice(len(probs), p=probs))


def evaluate(
    ctx,
    token_count=200,
    temperature=1.0,
    top_p=0.7,
    presencePenalty=0.1,
    countPenalty=0.1,
):
    ctx = ctx.strip()
    if not ctx:
        yield ""
        return

    token_count = max(1, min(int(token_count), GEN_LIMIT))
    temperature = max(0.2, float(temperature))
    top_p = max(0.0, min(float(top_p), 1.0))
    presencePenalty = max(0.0, float(presencePenalty))
    countPenalty = max(0.0, float(countPenalty))

    prompt_tokens = tokenizer.encode(ctx)
    if not prompt_tokens:
        yield ""
        return

    all_tokens = []
    out_last = 0
    out_str = ''
    occurrence = {}
    state = None
    out = None
    try:
        with torch.inference_mode():
            for i in range(token_count):
                input_ids = prompt_tokens[-CTX_LIMIT:] if i == 0 else [token]
                out, state = model.forward(input_ids, state)

                for token_id, count in occurrence.items():
                    out[token_id] -= presencePenalty + count * countPenalty

                token = sample_logits(out, temperature=temperature, top_p=top_p)
                if token == 0:  # RWKV's end-of-text token.
                    break
                all_tokens.append(token)

                for token_id in occurrence:
                    occurrence[token_id] *= PENALTY_DECAY

                token_text = tokenizer.decode([token])
                token_weight = 0 if token_text in ' \t0123456789' else 1
                occurrence[token] = occurrence.get(token, 0) + token_weight

                # Do not emit an incomplete UTF-8 sequence. Keep it in the
                # pending slice and try again after the next generated token.
                pending_text = tokenizer.decode(all_tokens[out_last:])
                if '\ufffd' not in pending_text:
                    out_str += pending_text
                    yield out_str.strip()
                    out_last = len(all_tokens)
    finally:
        del out
        del state
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

    yield out_str.strip()

@spaces.GPU(duration=120)
def Desensitization(english_text, token_count, temperature, top_p, presence_penalty, count_penalty):
    if not english_text or not english_text.strip():
        yield "请输入包含敏感信息的文本。"
        return

    full_prompt = f"Original: {english_text}\n\nRedacted:"
    for output in evaluate(full_prompt, token_count, temperature, top_p, presence_penalty, count_penalty):
        yield output  


with gr.Blocks(title="RWKV Desensitization 0.4B (CPU backend space)") as demo:
    with gr.Tab("Desensitization"):    
        gr.HTML(f"<div style='text-align:center;'><h1>RWKV Desensitization 0.4B (CPU backend space)</h1></div>")
        with gr.Row():
            with gr.Column():
                english_input = gr.Textbox(
                    label="敏感词文本段输入(注意不能有空行)",
                    lines=20,
                    placeholder="请输入带有敏感词的文本段",
                    value="Hello, my name is David Johnson and I live in Maine.\n"
                        "My credit card number is 4095-2609-9393-4932 and my crypto wallet id is 16Yeky6GMjeNkAiNcBY7ZhrLoMSgg1BoyZ.\n"
                        "On September 18 I visited dataprivacystack.org and sent an email to test@dataprivacystack.org,  from IP address 192.168.0.1.\n"
                        "My passport: 191280342 and my phone number: (212) 555-1234.\n"
                        "My social security number is 123-45-6789 and my social insurance number is 1234567890.\n"
                        "This is a valid International Bank Account Number: IL150120690000003111111 . Can you please check the status on bank account 954567876544?\n"
                        "Kate's social security number is 078-05-1126.  Her driver license? it is 1234567A.\n"
                )

            with gr.Column():
                chinese_output = gr.Textbox(
                    label="脱敏后的文本输出",
                    lines=20,
                    placeholder="脱敏后的文本段",
                    value=""
                )

        with gr.Row():
            translate_btn = gr.Button("Desensitize", variant="primary")
            clear_btn = gr.Button("Clear", variant="secondary")
            stop_btn = gr.Button("Stop", variant="stop")

        with gr.Accordion("Advanced Settings", open=False):
            token_count = gr.Slider(10, GEN_LIMIT, label="Max Tokens", step=10, value=GEN_LIMIT)
            temperature = gr.Slider(0.2, 2.0, label="Temperature", step=0.1, value=1.0)
            top_p = gr.Slider(0.0, 1.0, label="Top P", step=0.05, value=0)
            presence_penalty = gr.Slider(0.0, 1.0, label="Presence Penalty", step=0.1, value=0)
            count_penalty = gr.Slider(0.0, 1.0, label="Count Penalty", step=0.1, value=0)

        translate_event = translate_btn.click(
            fn=Desensitization,
            inputs=[english_input, token_count, temperature, top_p, presence_penalty, count_penalty],
            outputs=[chinese_output]
        )

        clear_btn.click(
            fn=lambda: ("", ""),
            inputs=[],
            outputs=[english_input, chinese_output]
        )

        stop_btn.click(
            fn=None,
            inputs=None,
            outputs=None,
            cancels=[translate_event]  
        )

demo.queue(max_size=10, default_concurrency_limit=1)
demo.launch()