import gradio as gr
import requests
import os

# -----------------------------
# GROK SETUP (xAI)
# -----------------------------

API_KEY = os.getenv("XAI_API_KEY")
API_URL = "https://api.x.ai/v1/chat/completions"

if not API_KEY:
    raise ValueError("❌ XAI_API_KEY not set in Hugging Face Secrets")


def call_llm(prompt):
    try:
        headers = {
            "Authorization": f"Bearer {API_KEY}",
            "Content-Type": "application/json"
        }

        payload = {
            "model": "grok-4-1-fast-non-reasoning",
            "messages": [
                {"role": "system", "content": "You are a helpful AI assistant."},
                {"role": "user", "content": prompt}
            ],
            "temperature": 0.7,
            "max_tokens": 400
        }

        response = requests.post(
            API_URL,
            headers=headers,
            json=payload,
            timeout=30
        )

        data = response.json()

        if "choices" not in data:
            return f"❌ API Error:\n{data}"

        return data["choices"][0]["message"]["content"]

    except Exception as e:
        return f"❌ Request Failed: {str(e)}"


# -----------------------------
# BLOCK LOGIC
# -----------------------------

def create_blocks(text):
    lines = [line.strip() for line in text.split("\n") if line.strip()]
    if not lines:
        lines = ["EMPTY"]

    blocks = [f"{i}: {line}" for i, line in enumerate(lines)]
    return blocks, gr.update(choices=blocks, value=blocks[0])


def select_block(blocks, selected):
    if not blocks or not selected:
        return ""

    idx = int(selected.split(":")[0])
    return blocks[idx].split(": ", 1)[1]


def update_block(blocks, selected, new_text):
    if not blocks or not selected:
        return blocks, gr.update()

    idx = int(selected.split(":")[0])
    blocks[idx] = f"{idx}: {new_text}"

    return blocks, gr.update(choices=blocks, value=blocks[idx])


# -----------------------------
# AI FUNCTIONS
# -----------------------------

# ✨ Rewrite selected block
def ai_rewrite(blocks, selected):
    if not blocks or not selected:
        return blocks, gr.update()

    idx = int(selected.split(":")[0])
    original = blocks[idx].split(": ", 1)[1]

    prompt = f"Rewrite this more clearly and professionally:\n\n{original}"

    new_text = call_llm(prompt)

    blocks[idx] = f"{idx}: {new_text}"

    return blocks, gr.update(choices=blocks, value=blocks[idx])


def ai_rewrite_all(blocks):
    if not blocks:
        return "⚠️ No blocks to rewrite."

    full_text = "\n".join([b.split(": ", 1)[1] for b in blocks])

    prompt = f"""
Rewrite the following text to be clearer, more professional, and well-structured:

{full_text}
"""

    return call_llm(prompt)


# -----------------------------
# UI
# -----------------------------

with gr.Blocks(title="Blocksmith + Grok") as app:

    gr.Markdown("# 🧱 Blocksmith (Powered by Grok)")

    state_blocks = gr.State([])

    input_text = gr.Textbox(label="Paste Text", lines=8)
    create_btn = gr.Button("Create Blocks")

    block_selector = gr.Dropdown(choices=[], label="Blocks", interactive=True)

    block_editor = gr.Textbox(label="Edit Block")

    with gr.Row():
        update_btn = gr.Button("Update")
        rewrite_btn = gr.Button("✨ Rewrite (AI)")
        #rewrite_all_btn = gr.Button("🚀 Rewrite All (AI)")

    full_output = gr.Textbox(
        label="Rewritten Full Text",
        lines=10
    )
    rewrite_all_btn = gr.Button("🚀 Rewrite All (AI)")
    
    # -----------------------------
    # EVENTS
    # -----------------------------

    create_btn.click(
        fn=create_blocks,
        inputs=input_text,
        outputs=[state_blocks, block_selector]
    )

    block_selector.change(
        fn=select_block,
        inputs=[state_blocks, block_selector],
        outputs=block_editor
    )

    update_btn.click(
        fn=update_block,
        inputs=[state_blocks, block_selector, block_editor],
        outputs=[state_blocks, block_selector]
    )

    rewrite_btn.click(
        fn=ai_rewrite,
        inputs=[state_blocks, block_selector],
        outputs=[state_blocks, block_selector]
    )

    rewrite_all_btn.click(
        fn=ai_rewrite_all,
        inputs=state_blocks,
        outputs=full_output
    )

app.launch()