import openai
print(f"openai version: {openai.__version__}")
print(f"openai path: {openai.__file__}")
import os
import time
import pathlib
import gradio as gr
from openai import OpenAI

# ── Bootstrap ────────────────────────────────────────────────────────────────
api_key = os.environ.get("OPENAI_API_KEY", "")
print(f"API key set: {bool(api_key)}, length: {len(api_key)}")
client = OpenAI(api_key=api_key)

ASSISTANT_NAME = "Research Paper Assistant"
INSTRUCTIONS = (
    "Respond helpfully based on your vector store files or general knowledge. "
    "Make jokes with everything, Gen Z humor."
)
CORPUS_DIR = pathlib.Path(__file__).parent / "corpus"


def _create_vector_store() -> str:
    vs = client.vector_stores.create(name=f"{ASSISTANT_NAME} Store")
    corpus_files = list(CORPUS_DIR.glob("*")) if CORPUS_DIR.exists() else []
    if corpus_files:
        file_streams = [open(f, "rb") for f in corpus_files if f.is_file()]
        if file_streams:
            batch = client.vector_stores.file_batches.upload_and_poll(
                vector_store_id=vs.id, files=file_streams
            )
            for s in file_streams:
                s.close()
    return vs.id


def _create_assistant(vs_id: str) -> str:
    assistant = client.beta.assistants.create(
        name=ASSISTANT_NAME,
        instructions=INSTRUCTIONS,
        model="gpt-4o",
        tools=[{"type": "file_search"}],
        tool_resources={"file_search": {"vector_store_ids": [vs_id]}},
    )
    return assistant.id


print("🚀 Spinning up the assistant… no cap.")
VECTOR_STORE_ID = _create_vector_store()
ASSISTANT_ID = _create_assistant(VECTOR_STORE_ID)
print(f"✅ Assistant ready. ID={ASSISTANT_ID}  VS={VECTOR_STORE_ID}")

# ── Thread cache (session → thread_id) ───────────────────────────────────────
_threads: dict[str, str] = {}


def _get_thread(session_id: str) -> str:
    if session_id not in _threads:
        _threads[session_id] = client.beta.threads.create().id
    return _threads[session_id]


# ── Chat logic ────────────────────────────────────────────────────────────────

def chat(message: str, history: list, session_id: str):
    thread_id = _get_thread(session_id)
    client.beta.threads.messages.create(
        thread_id=thread_id, role="user", content=message
    )

    citations_text = ""
    collected = ""

    with client.beta.threads.runs.stream(
        thread_id=thread_id,
        assistant_id=ASSISTANT_ID,
    ) as stream:
        for event in stream:
            if hasattr(event, "data") and hasattr(event.data, "delta"):
                delta = event.data.delta
                if hasattr(delta, "content") and delta.content:
                    for block in delta.content:
                        if hasattr(block, "text") and block.text:
                            if hasattr(block.text, "value"):
                                collected += block.text.value
                            if hasattr(block.text, "annotations"):
                                for ann in block.text.annotations:
                                    if hasattr(ann, "file_citation"):
                                        try:
                                            f = client.files.retrieve(ann.file_citation.file_id)
                                            citations_text += f"\n📎 {f.filename}"
                                        except Exception:
                                            pass
            yield collected + (f"\n\n---\n**Citations:**{citations_text}" if citations_text else "")


def upload_files(files, session_state):
    if not files:
        return "No files dropped, bestie 💀"
    names = []
    for f in files:
        with open(f.name, "rb") as fh:
            uploaded = client.files.create(file=fh, purpose="assistants")
        client.vector_stores.files.create(
            vector_store_id=VECTOR_STORE_ID, file_id=uploaded.id
        )
        names.append(pathlib.Path(f.name).name)
    return f"Slay! Uploaded {len(names)} file(s): {', '.join(names)} ✨"


# ── Custom CSS (futuristic glass) ─────────────────────────────────────────────
CSS = """
@import url('https://fonts.googleapis.com/css2?family=Space+Mono:wght@400;700&family=Syne:wght@400;700;800&display=swap');

:root {
  --bg: #050810;
  --glass: rgba(255,255,255,0.04);
  --glass-border: rgba(255,255,255,0.10);
  --accent: #7DF9FF;
  --accent2: #BF5AF2;
  --text: #E8EEFF;
  --muted: rgba(232,238,255,0.45);
  --radius: 18px;
  --glow: 0 0 40px rgba(125,249,255,0.15);
}

body, .gradio-container {
  background: var(--bg) !important;
  font-family: 'Syne', sans-serif !important;
  color: var(--text) !important;
  min-height: 100vh;
}

/* Animated mesh background */
.gradio-container::before {
  content: '';
  position: fixed;
  inset: 0;
  background:
    radial-gradient(ellipse 80% 60% at 20% 10%, rgba(125,249,255,0.07) 0%, transparent 60%),
    radial-gradient(ellipse 60% 50% at 80% 80%, rgba(191,90,242,0.07) 0%, transparent 60%);
  pointer-events: none;
  z-index: 0;
  animation: meshDrift 12s ease-in-out infinite alternate;
}

@keyframes meshDrift {
  from { transform: scale(1) translateY(0); }
  to   { transform: scale(1.05) translateY(-20px); }
}

/* Scanline overlay */
.gradio-container::after {
  content: '';
  position: fixed;
  inset: 0;
  background: repeating-linear-gradient(
    0deg,
    transparent,
    transparent 3px,
    rgba(255,255,255,0.012) 3px,
    rgba(255,255,255,0.012) 4px
  );
  pointer-events: none;
  z-index: 0;
}

#header-block {
  text-align: center;
  padding: 48px 0 16px;
  position: relative;
  z-index: 1;
  animation: fadeSlideDown 0.8s cubic-bezier(0.16,1,0.3,1) both;
}

@keyframes fadeSlideDown {
  from { opacity: 0; transform: translateY(-28px); }
  to   { opacity: 1; transform: translateY(0); }
}

#app-title {
  font-size: clamp(2rem, 5vw, 3.4rem) !important;
  font-weight: 800 !important;
  letter-spacing: -0.02em !important;
  background: linear-gradient(135deg, var(--accent) 0%, #fff 50%, var(--accent2) 100%) !important;
  -webkit-background-clip: text !important;
  -webkit-text-fill-color: transparent !important;
  background-clip: text !important;
  margin-bottom: 8px !important;
  text-shadow: none !important;
}

#app-title::after {
  content: '';
  display: block;
  width: 60px;
  height: 2px;
  background: linear-gradient(90deg, var(--accent), var(--accent2));
  margin: 10px auto 0;
  border-radius: 2px;
  animation: lineExpand 1s 0.5s cubic-bezier(0.16,1,0.3,1) both;
}

@keyframes lineExpand {
  from { width: 0; opacity: 0; }
  to   { width: 60px; opacity: 1; }
}

#app-subtitle {
  color: var(--muted) !important;
  font-family: 'Space Mono', monospace !important;
  font-size: 0.9rem !important;
  letter-spacing: 0.06em !important;
  animation: fadeSlideDown 0.8s 0.15s cubic-bezier(0.16,1,0.3,1) both;
}

/* Glass panels */
.glass-panel {
  background: var(--glass) !important;
  border: 1px solid var(--glass-border) !important;
  border-radius: var(--radius) !important;
  backdrop-filter: blur(20px) !important;
  -webkit-backdrop-filter: blur(20px) !important;
  box-shadow: var(--glow), inset 0 1px 0 rgba(255,255,255,0.06) !important;
  position: relative;
  z-index: 1;
  animation: fadeUp 0.7s 0.2s cubic-bezier(0.16,1,0.3,1) both;
}

@keyframes fadeUp {
  from { opacity: 0; transform: translateY(20px); }
  to   { opacity: 1; transform: translateY(0); }
}

/* Chatbot */
#chatbot {
  background: transparent !important;
  border: none !important;
}

#chatbot .message-wrap {
  padding: 0 !important;
}

#chatbot .user {
  background: linear-gradient(135deg, rgba(125,249,255,0.12), rgba(191,90,242,0.08)) !important;
  border: 1px solid rgba(125,249,255,0.2) !important;
  border-radius: 14px 14px 4px 14px !important;
  color: var(--text) !important;
  font-family: 'Syne', sans-serif !important;
}

#chatbot .bot {
  background: rgba(255,255,255,0.03) !important;
  border: 1px solid var(--glass-border) !important;
  border-radius: 14px 14px 14px 4px !important;
  color: var(--text) !important;
  font-family: 'Syne', sans-serif !important;
}

/* Input box */
#chat-input textarea {
  background: rgba(255,255,255,0.05) !important;
  border: 1px solid var(--glass-border) !important;
  border-radius: 12px !important;
  color: var(--text) !important;
  font-family: 'Space Mono', monospace !important;
  font-size: 0.9rem !important;
  transition: border-color 0.2s, box-shadow 0.2s !important;
}

#chat-input textarea:focus {
  border-color: var(--accent) !important;
  box-shadow: 0 0 0 2px rgba(125,249,255,0.15) !important;
  outline: none !important;
}

/* Buttons */
button.primary, #send-btn {
  background: linear-gradient(135deg, var(--accent), var(--accent2)) !important;
  color: #050810 !important;
  border: none !important;
  border-radius: 10px !important;
  font-family: 'Space Mono', monospace !important;
  font-weight: 700 !important;
  letter-spacing: 0.04em !important;
  transition: opacity 0.2s, transform 0.15s !important;
}

button.primary:hover, #send-btn:hover {
  opacity: 0.88 !important;
  transform: translateY(-1px) !important;
}

button.secondary {
  background: var(--glass) !important;
  border: 1px solid var(--glass-border) !important;
  color: var(--muted) !important;
  border-radius: 10px !important;
  font-family: 'Space Mono', monospace !important;
  transition: border-color 0.2s !important;
}

button.secondary:hover {
  border-color: var(--accent) !important;
  color: var(--accent) !important;
}

/* Upload area */
#upload-area {
  border: 1.5px dashed rgba(125,249,255,0.3) !important;
  border-radius: var(--radius) !important;
  background: rgba(125,249,255,0.02) !important;
  transition: border-color 0.2s, background 0.2s !important;
}

#upload-area:hover {
  border-color: var(--accent) !important;
  background: rgba(125,249,255,0.05) !important;
}

/* Upload status */
#upload-status textarea {
  background: transparent !important;
  border: none !important;
  color: var(--accent) !important;
  font-family: 'Space Mono', monospace !important;
  font-size: 0.8rem !important;
}

/* Labels */
label span {
  color: var(--muted) !important;
  font-family: 'Space Mono', monospace !important;
  font-size: 0.75rem !important;
  letter-spacing: 0.08em !important;
  text-transform: uppercase !important;
}

/* Scrollbars */
::-webkit-scrollbar { width: 4px; }
::-webkit-scrollbar-track { background: transparent; }
::-webkit-scrollbar-thumb { background: rgba(125,249,255,0.2); border-radius: 4px; }
::-webkit-scrollbar-thumb:hover { background: rgba(125,249,255,0.4); }

/* Divider glow */
.divider-glow {
  height: 1px;
  background: linear-gradient(90deg, transparent, var(--accent), var(--accent2), transparent);
  opacity: 0.3;
  margin: 8px 0;
}

/* Pulsing dot indicator */
.status-dot {
  display: inline-block;
  width: 7px;
  height: 7px;
  border-radius: 50%;
  background: var(--accent);
  box-shadow: 0 0 8px var(--accent);
  animation: pulse 2s ease-in-out infinite;
  margin-right: 8px;
  vertical-align: middle;
}

@keyframes pulse {
  0%, 100% { opacity: 1; transform: scale(1); }
  50%       { opacity: 0.5; transform: scale(0.85); }
}

/* Section labels */
.section-label {
  font-family: 'Space Mono', monospace !important;
  font-size: 0.7rem !important;
  letter-spacing: 0.14em !important;
  text-transform: uppercase !important;
  color: var(--muted) !important;
  margin-bottom: 8px !important;
}
"""

# ── Build the UI ──────────────────────────────────────────────────────────────
with gr.Blocks(css=CSS, title="Cool again.") as demo:

    session_id = gr.State(lambda: str(time.time_ns()))

    with gr.Column(elem_id="header-block"):
        gr.HTML('<h1 id="app-title">Cool again.</h1>')
        gr.HTML('<p id="app-subtitle"><span class="status-dot"></span>I can help. Ask away!</p>')
        gr.HTML('<div class="divider-glow"></div>')

    with gr.Row():
        # ── Main chat column ──────────────────────────────────────────────────
        with gr.Column(scale=3, elem_classes="glass-panel"):
            gr.HTML('<p class="section-label">// chat interface</p>')
            chatbot = gr.Chatbot(
                elem_id="chatbot",
                label="",
                height=520,
                show_label=False,
                bubble_full_width=False,
                render_markdown=True,
            )
            with gr.Row():
                msg_box = gr.Textbox(
                    elem_id="chat-input",
                    placeholder="drop your research question here, no cap… ✨",
                    show_label=False,
                    lines=1,
                    scale=5,
                )
                send_btn = gr.Button("Send →", variant="primary", scale=1, elem_id="send-btn")
            clear_btn = gr.Button("Clear chat 🧹", variant="secondary", size="sm")

        # ── Side panel ───────────────────────────────────────────────────────
        with gr.Column(scale=1, elem_classes="glass-panel"):
            gr.HTML('<p class="section-label">// knowledge drop</p>')
            gr.HTML(
                '<p style="color:rgba(232,238,255,0.5);font-size:0.8rem;font-family:\'Space Mono\',monospace;line-height:1.6;">'
                'Upload your own PDFs, docs, or text files and the assistant will slay with that context fr fr 💅'
                '</p>'
            )
            file_uploader = gr.File(
                label="Upload files",
                file_types=[".pdf", ".txt", ".md", ".docx"],
                file_count="multiple",
                elem_id="upload-area",
            )
            upload_btn = gr.Button("Upload to knowledge base 🚀", variant="primary")
            upload_status = gr.Textbox(
                label="Status",
                interactive=False,
                elem_id="upload-status",
                lines=2,
            )

            gr.HTML('<div class="divider-glow" style="margin:16px 0;"></div>')
            gr.HTML(
                '<p class="section-label">// about</p>'
                '<p style="color:rgba(232,238,255,0.4);font-size:0.75rem;font-family:\'Space Mono\',monospace;line-height:1.7;">'
                'Powered by GPT-4o + file_search.<br>'
                'Citations shown below responses.<br>'
                'Context persists per session.'
                '</p>'
            )

    # ── Wire events ───────────────────────────────────────────────────────────

    def _respond(message, history, sid):
        if not message.strip():
            yield history
            return
        history = history or []
        history.append([message, ""])
        for token in chat(message, history[:-1], sid):
            history[-1][1] = token
            yield history

    def _clear():
        return []

    msg_box.submit(_respond, [msg_box, chatbot, session_id], chatbot).then(
        lambda: "", None, msg_box
    )
    send_btn.click(_respond, [msg_box, chatbot, session_id], chatbot).then(
        lambda: "", None, msg_box
    )
    clear_btn.click(_clear, outputs=chatbot)
    upload_btn.click(upload_files, inputs=[file_uploader], outputs=[upload_status])

demo.launch()