#!/usr/bin/env python3
"""Streamlit app to reconstruct prompts from AI assistant responses."""
import html as html_lib
import os
import re
import requests
import torch
import pandas as pd
import streamlit as st
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL_ID = "dejanseo/reverse-prompter"
SEPARATOR = "\n###\n"
CONTRASTIVE_CONFIGS = [
    {"penalty_alpha": a, "top_k": k}
    for k in [2, 4, 6, 15]
    for a in [0.1, 0.2, 0.3, 0.4, 0.5, 0.6]
]


@st.cache_resource
def load_model():
    # bfloat16 on CUDA; float32 on CPU to prevent CPU kernel incompatibilities
    if torch.cuda.is_available():
        device = "cuda"
        dtype = torch.bfloat16
    else:
        device = "cpu"
        dtype = torch.float32

    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
        torch_dtype=dtype,
        low_cpu_mem_usage=True,
    ).to(device).eval()

    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
    return model, tokenizer


def analyze_output(model, tokenizer, input_ids, prefix_len):
    with torch.no_grad():
        outputs = model(input_ids.unsqueeze(0))
        logits = outputs.logits[0, prefix_len - 1:-1]
        targets = input_ids[prefix_len:]
        log_probs = torch.log_softmax(logits, dim=-1)
        token_log_probs = log_probs[torch.arange(len(targets)), targets]
        perplexity = torch.exp(-token_log_probs.mean()).item()
        probs = torch.exp(token_log_probs).tolist()
        tokens = [tokenizer.decode([t]) for t in targets.tolist()]
    return perplexity, tokens, probs


def render_colored(tokens, probs):
    parts = []
    for token, prob in zip(tokens, probs):
        opacity = 0.15 + 0.85 * prob
        clean = html_lib.escape(token).replace("\n", " ").replace("\r", "")
        clean = re.sub(r"[^\x20-\x7E\u00A0-\uFFFF]", "", clean)
        parts.append(f'<span style="opacity:{opacity:.2f}">{clean}</span>')
    return "".join(parts)


def render_result(tokens, probs):
    colored = render_colored(tokens, probs)
    return f'<div style="margin-bottom:8px">{colored}</div>'


def fetch_url_markdown(url):
    login = os.getenv("DATAFORSEO_LOGIN", "")
    password = os.getenv("DATAFORSEO_PASSWORD", "")
    if not login or not password:
        raise ValueError("DATAFORSEO_LOGIN / DATAFORSEO_PASSWORD not set in Spaces secrets")

    resp = requests.post(
        "https://api.dataforseo.com/v3/on_page/content_parsing/live",
        json=[{"url": url, "enable_javascript": True, "markdown_view": True}],
        auth=(login, password),
        timeout=30,
    )
    if resp.status_code != 200:
        raise ValueError(f"API returned status {resp.status_code}")

    data = resp.json()
    tasks = data.get("tasks", [])
    if not tasks or tasks[0].get("status_code") != 20000:
        raise ValueError(f"API task error: {tasks[0].get('status_message') if tasks else 'no tasks'}")

    items = (tasks[0].get("result") or [{}])[0].get("items")
    if not items:
        raise ValueError("No items returned from API")

    text = items[0].get("page_as_markdown") or ""
    if len(text) < 100:
        raise ValueError("Page content too short or empty")
    return text


st.set_page_config(layout="wide")
st.logo("https://dejan.ai/media/dejan-logo.png", size="LARGE", link="https://dejan.ai/")
st.subheader("Reverse Prompting")
st.caption(
    "Paste any AI-generated text and this tool will reverse-engineer the most likely prompts that produced it. "
    "It runs the text through a fine-tuned language model across multiple decoding configurations and ranks "
    "the reconstructed prompts by model confidence."
)

model, tokenizer = load_model()

TEST_FILE = os.path.join(os.path.dirname(__file__), "test.md")
FALLBACK_TEST_TEXT = (
    "The Eiffel Tower is a wrought-iron lattice tower on the Champ de Mars in Paris, France. "
    "It is named after the engineer Gustave Eiffel, whose company designed and built the tower from 1887 to 1889."
)

if "quick_test" not in st.session_state:
    st.session_state.quick_test = False
if "fetched_text" not in st.session_state:
    st.session_state.fetched_text = ""
if "run_after_fetch" not in st.session_state:
    st.session_state.run_after_fetch = False

if st.session_state.fetched_text:
    st.session_state.text_input = st.session_state.fetched_text
    st.session_state.fetched_text = ""
    st.session_state.run_after_fetch = True

if st.session_state.quick_test:
    if os.path.exists(TEST_FILE):
        with open(TEST_FILE, encoding="utf-8") as f:
            st.session_state.text_input = f.read()
    else:
        st.session_state.text_input = FALLBACK_TEST_TEXT

tab_paste, tab_url = st.tabs(["Paste", "URL Mode \u1D49\u02E3\u1D56\u1D49\u02B3\u2071\u1D50\u1D49\u207F\u1D57\u1D43\u02E1"])

with tab_paste:
    text_input = st.text_area("Paste an AI assistant response", height=200, key="text_input")
    col_btn1, col_btn2, col_btn3, _ = st.columns([1, 1, 1, 3])
    run = col_btn1.button("Reconstruct Prompts", type="primary", width="stretch")
    quick = col_btn2.button("Quick Test", type="secondary", width="stretch")
    clear = col_btn3.button("Clear", type="secondary", width="stretch")

    if quick and not st.session_state.quick_test:
        st.session_state.quick_test = True
        st.rerun()

    if clear:
        st.session_state.text_input = ""
        st.rerun()

run = run or st.session_state.quick_test or st.session_state.run_after_fetch
if st.session_state.quick_test:
    st.session_state.quick_test = False
if st.session_state.run_after_fetch:
    st.session_state.run_after_fetch = False

with tab_url:
    st.caption(
        "This experimental feature reveals what prompts would lead to generating a page such as the one you entered. "
        "While *not the intended use of the model*, it's certainly an interesting feature for exploring semantic make-up of the page."
    )
    url_col, url_btn_col = st.columns([5, 1])
    url_input = url_col.text_input("URL to scrape", label_visibility="collapsed", placeholder="Enter a URL to scrape")
    fetch = url_btn_col.button("Fetch", type="primary", width="stretch")

    if fetch and url_input.strip():
        url = url_input.strip()
        if not url.startswith(("http://", "https://")):
            url = "https://" + url
        if "/" not in url.split("//", 1)[-1]:
            url += "/"

        with st.spinner("Fetching page..."):
            try:
                st.session_state.fetched_text = fetch_url_markdown(url)
                st.rerun()
            except Exception:
                st.warning(
                    "This page couldn't be scraped. The site may be blocking automated access or "
                    "the page has insufficient content. Try pasting the content manually in the text area."
                )

if run and text_input.strip():
    prompt = text_input.strip() + SEPARATOR
    inputs = tokenizer(prompt, return_tensors="pt", add_special_tokens=False).to(model.device)
    prefix_len = inputs["input_ids"].shape[-1]
    tables_container = st.container()

    st.subheader("Word Frequency Heatmap")
    results = []
    seen = set()
    progress = st.progress(0)
    output_container = st.empty()

    for i, config in enumerate(CONTRASTIVE_CONFIGS):
        progress.progress((i + 1) / len(CONTRASTIVE_CONFIGS))
        try:
            outputs = model.generate(
                **inputs,
                max_new_tokens=256,
                trust_remote_code=True,
                num_return_sequences=1,
                **config,
            )
            gen_ids = outputs[0][prefix_len:]
            eos_id = tokenizer.eos_token_id
            if eos_id is not None and (gen_ids == eos_id).any():
                gen_ids = gen_ids[:(gen_ids == eos_id).nonzero(as_tuple=True)[0][0]]

            text = tokenizer.decode(gen_ids, skip_special_tokens=True).strip()

            if text and text not in seen:
                seen.add(text)
                trimmed = torch.cat([outputs[0][:prefix_len], gen_ids])
                ppl, tokens, probs = analyze_output(model, tokenizer, trimmed, prefix_len)
                results.append((text, ppl, tokens, probs, config))

                top = sorted(results, key=lambda x: x[1])[:10]
                html = "".join(render_result(t, p) for _, _, t, p, c in top)
                output_container.markdown(html, unsafe_allow_html=True)
        except Exception:
            pass

    progress.empty()
    results.sort(key=lambda x: x[1])

    if results:
        top10 = results[:10]
        col_prompts, col_phrases = tables_container.columns(2)

        with col_prompts:
            st.subheader("Reconstructed Prompts")
            st.caption(
                "Top 10 most likely prompts ranked by lowest perplexity. "
                "Token opacity reflects the model's confidence in each word."
            )
            df_top = pd.DataFrame(
                [(text, round(ppl, 2), [round(p * 100) for p in probs]) for text, ppl, _, probs, _ in top10],
                columns=["Prompt", "Perplexity", "Confidence"],
            )
            st.dataframe(
                df_top,
                width="content",
                hide_index=True,
                column_config={
                    "Prompt": st.column_config.TextColumn(width=None),
                    "Perplexity": st.column_config.NumberColumn(width=None),
                    "Confidence": st.column_config.BarChartColumn(y_min=0, y_max=100),
                },
            )

        with col_phrases:
            st.subheader("Key Phrases")
            st.caption("The most important phrases scored by balance between length and frequency.")
            texts = [re.sub(r"[^\w\s]", "", text.lower()).split() for text, _, _, _, _ in top10]
            phrase_hits = {}
            for idx, words in enumerate(texts):
                for length in range(2, len(words) + 1):
                    for start in range(len(words) - length + 1):
                        phrase = " ".join(words[start:start + length])
                        if phrase not in phrase_hits:
                            phrase_hits[phrase] = set()
                        phrase_hits[phrase].add(idx)

            shared = [
                (p, len(ids), len(p.split()), len(ids) * len(p.split()))
                for p, ids in phrase_hits.items()
                if len(ids) >= 2
            ]
            shared.sort(key=lambda x: x[3], reverse=True)

            filtered = []
            for phrase, count, length, score in shared:
                if not any(phrase in longer and score <= lscore for longer, _, _, lscore in filtered):
                    filtered.append((phrase, count, length, score))

            if filtered:
                max_score = filtered[0][3]
                rows = [(p, round(s / max_score * 100)) for p, _, _, s in filtered[:20]]
                df = pd.DataFrame(rows, columns=["Phrase", "Score"])
                st.dataframe(
                    df,
                    width="content",
                    hide_index=True,
                    column_config={
                        "Score": st.column_config.ProgressColumn(format="%d%%", min_value=0, max_value=100),
                    },
                )