import gradio as gr
from anthropic import Anthropic
from dotenv import load_dotenv
from docx import Document
from docx.shared import Pt
import os
import tempfile
import asyncio

# === Load API Key ===
load_dotenv()
client = Anthropic(api_key=os.getenv("ANTHROPIC_API_KEY"))

# === Async Streaming Function ===
async def query_claude_stream(prompt, text, doc_name, progress=None):
    if not prompt.strip() and not text.strip():
        yield "⚠️ Please enter a prompt or text.", None
        return

    output_text = ""
    if progress:
        progress(0)

    user_content = ""
    if prompt.strip():
        user_content += prompt.strip()
    if text.strip():
        user_content += f"\n\n{text.strip()}"

    # Blocking stream wrapped in async
    def _run_stream():
        with client.messages.stream(
            model="claude-sonnet-4-20250514",
            max_tokens=64000,
            messages=[{"role": "user", "content": user_content}],
        ) as stream:
            for event in stream:
                if event.type == "content_block_delta" and event.delta.type == "text_delta":
                    yield event.delta.text

    i = 0
    async for delta in _to_async_generator(_run_stream()):
        i += 1
        output_text += delta
        if progress:
            progress(min(0.95, i / 500))
        yield output_text, None

    if progress:
        progress(1)

    cleaned_output = "\n".join(
        line for line in output_text.splitlines() if line.strip() != "---"
    ).strip()

    # === Create DOCX ===
    if not doc_name.strip():
        doc_name = "Claude_Output"
    docx_path = os.path.join(tempfile.gettempdir(), f"{doc_name}.docx")

    document = Document()
    for line in cleaned_output.split("\n"):
        clean_line = line.strip()
        if not clean_line:
            continue
        if clean_line.startswith("# "):
            document.add_heading(clean_line[2:], level=1)
        elif clean_line.startswith("## "):
            document.add_heading(clean_line[3:], level=2)
        elif clean_line.startswith("### "):
            document.add_heading(clean_line[4:], level=3)
        else:
            para = document.add_paragraph(clean_line)
            para.style.font.size = Pt(11)

    document.save(docx_path)
    yield cleaned_output, docx_path


# === Helper: sync generator → async generator ===
async def _to_async_generator(sync_gen):
    for item in sync_gen:
        yield item
        await asyncio.sleep(0)


# === Word Counter ===
def count_words(text):
    words = len(text.split())
    chars = len(text)
    return f"📊 Words: {words:,} | Characters: {chars:,}"


# === Gradio UI (Backward-compatible: no css, no theme) ===
with gr.Blocks() as demo:

    gr.Markdown("<h2 style='text-align:center;'>Claude UI-HF</h2>")

    prompt = gr.Textbox(label="Prompt", lines=3)
    text = gr.Textbox(label="Text", lines=15)
    doc_name = gr.Textbox(label="📂 Document Name")

    with gr.Row():
        run_btn = gr.Button("🚀 Run", variant="primary")
        cancel_btn = gr.Button("🛑 Cancel")

    output = gr.Textbox(
        label="Claude Output",
        lines=15,
        interactive=True
    )

    with gr.Row():
        copy_btn = gr.Button("📋 Copy Output")
        download_btn = gr.File(label="⬇️ Download DOCX")
        stats = gr.Label(label="Word/Character Count")

    # Run -> Streaming Claude
    stream_event = run_btn.click(
        query_claude_stream,
        inputs=[prompt, text, doc_name],
        outputs=[output, download_btn],
        concurrency_id="job"
    )

    # Cancel button
    cancel_btn.click(None, None, None, cancels=stream_event)

    # Auto word counter
    output.change(count_words, inputs=output, outputs=stats)

    # Copy button
    copy_btn.click(
        lambda x: x,
        inputs=output,
        outputs=None
    ).then(
        js="(text) => {navigator.clipboard.writeText(text); alert('✅ Copied to clipboard');}"
    )


# Launch app (works on any Gradio >=3.x)
demo.queue()
demo.launch()