import streamlit as st
from anthropic import Anthropic
from dotenv import load_dotenv
from docx import Document
from docx.shared import Pt
import os
import tempfile

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

# ===== Page Config =====
st.set_page_config(page_title="Claude UI-HF", layout="wide")
st.markdown("<h2 style='text-align:center;'>Claude UI-HF</h2>", unsafe_allow_html=True)

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

# ===== DOCX Creation =====
def create_docx(text, doc_name):
    cleaned_output = "\n".join(line for line in text.splitlines() if line.strip() != "---").strip()
    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)
    return docx_path

# ===== Main Streaming Function =====
def stream_claude_response(prompt, text):
    if not prompt.strip() and not text.strip():
        st.warning("⚠️ Please enter a prompt or text.")
        return ""

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

    output_text = ""
    progress_bar = st.progress(0, text="Starting...")

    # Single placeholder for output
    output_placeholder = st.empty()
    output_placeholder.text_area("Claude Output", value="", height=400, key="output_box")

    # Streaming from Claude
    with client.messages.stream(
        model="claude-sonnet-4-20250514",
        max_tokens=64000,
        messages=[{"role": "user", "content": user_content}],
    ) as stream:
        for i, event in enumerate(stream):
            if event.type == "content_block_delta" and event.delta.type == "text_delta":
                delta = event.delta.text
                output_text += delta
                # Update output placeholder
                output_placeholder.text_area("Claude Output", value=output_text, height=400)
                progress_bar.progress(min(i / 400, 0.95), text="Generating...")

    progress_bar.progress(1.0, text="✅ Done")
    return output_text

# ===== UI Inputs =====
prompt = st.text_area("Prompt", placeholder="Enter your instruction here", height=100)
text = st.text_area("Text", placeholder="Paste your large text here", height=300)
doc_name = st.text_input("📂 Document Name", placeholder="e.g. MyNotes")

run_col, cancel_col = st.columns([1, 1])
run_clicked = run_col.button("🚀 Run")
cancel_clicked = cancel_col.button("🛑 Cancel")

stats_placeholder = st.empty()
download_placeholder = st.empty()
copy_placeholder = st.empty()

# ===== Run Logic =====
if run_clicked:
    try:
        response_text = stream_claude_response(prompt, text)
        if response_text:
            # Word/char stats
            stats_placeholder.write(count_words(response_text))

            # Save DOCX & provide download
            docx_file = create_docx(response_text, doc_name)
            with open(docx_file, "rb") as f:
                download_placeholder.download_button(
                    label="⬇️ Download DOCX",
                    data=f,
                    file_name=f"{doc_name or 'Claude_Output'}.docx",
                    mime="application/vnd.openxmlformats-officedocument.wordprocessingml.document",
                    use_container_width=True
                )

            # Copy to clipboard (JS)
            copy_placeholder.markdown("""
            <button onclick="navigator.clipboard.writeText(document.querySelector('textarea[data-testid=stTextArea-input]').value); alert('✅ Copied to clipboard!');">
            📋 Copy Output
            </button>
            """, unsafe_allow_html=True)

    except Exception as e:
        st.error(f"❌ Error: {e}")

elif cancel_clicked:
    st.warning("🛑 Cancel clicked — refresh the page to restart.")