agent/main.py

import asyncio
import json

import uvicorn
from fastapi import FastAPI, File, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import HTMLResponse, StreamingResponse
from pydantic import BaseModel
from typing import Any, List, Optional

import blood_rag
import reports
import tools as notepad
from agent import run_agent, stream_agent
from config import PORT

app = FastAPI(title="Agentic Notepad Server")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)


# ── OpenAI-compatible request types ───────────────────────────────────────────
class Message(BaseModel):
    role:    str
    content: Any

class ChatRequest(BaseModel):
    model:       Optional[str] = None   # None → fall back to LLM_MODEL in agent config
    messages:    List[Message]
    stream:      bool = False
    temperature: Optional[float] = None


# ── /v1/chat/completions ───────────────────────────────────────────────────────
@app.post("/v1/chat/completions")
async def chat_completions(req: ChatRequest):
    msgs  = [m.model_dump() for m in req.messages]
    model = req.model or None   # pass None → agent uses LLM_MODEL from config
    print(f"[agent] ▶ using model: {model or 'default → ' + __import__('config').LLM_MODEL}", flush=True)

    if req.stream:
        async def event_stream():
            async for token in stream_agent(msgs, model=model):
                chunk = {
                    "id":      "chatcmpl-agent",
                    "object":  "chat.completion.chunk",
                    "choices": [{"delta": {"content": token}, "index": 0, "finish_reason": None}],
                }
                yield f"data: {json.dumps(chunk)}\n\n"
            done = {
                "id":      "chatcmpl-agent",
                "object":  "chat.completion.chunk",
                "choices": [{"delta": {}, "index": 0, "finish_reason": "stop"}],
            }
            yield f"data: {json.dumps(done)}\n\n"
            yield "data: [DONE]\n\n"

        return StreamingResponse(event_stream(), media_type="text/event-stream")

    text = await run_agent(msgs, model=model)
    return {
        "id":      "chatcmpl-agent",
        "object":  "chat.completion",
        "model":   model or "gpt-4o-mini",
        "choices": [{"message": {"role": "assistant", "content": text}, "finish_reason": "stop", "index": 0}],
    }


# ── Blood RAG routes ──────────────────────────────────────────────────────────
@app.post("/api/blood-rag/upload")
async def blood_rag_upload(file: UploadFile = File(...)):
    raw  = await file.read()
    name = file.filename or "document"
    try:
        text = raw.decode("utf-8", errors="replace")
    except Exception:
        return {"error": "Could not decode file"}
    meta = await blood_rag.add_document(name, text)
    return meta


@app.get("/api/blood-rag/documents")
async def blood_rag_list():
    return blood_rag.list_documents()


@app.delete("/api/blood-rag/documents/{doc_id}")
async def blood_rag_delete(doc_id: str):
    deleted = blood_rag.delete_document(doc_id)
    return {"deleted": deleted, "doc_id": doc_id}


@app.get("/api/blood-rag/search")
async def blood_rag_search(q: str, top_k: int = 4):
    result = await blood_rag.search(q, top_k=top_k)
    return {"query": q, "result": result}


# ── Blood report routes ────────────────────────────────────────────────────────
@app.post("/api/reports/upload")
async def upload_report(file: UploadFile = File(...)):
    raw = await file.read()
    return await reports.process_report_upload(raw, file.filename or "report.pdf")


@app.get("/api/reports/current")
async def current_report():
    return reports.get_current_report()


@app.delete("/api/reports/current")
async def delete_report():
    reports.clear_report()
    return {"status": "cleared"}


# ── Notepad routes ─────────────────────────────────────────────────────────────
@app.get("/notepad", response_class=HTMLResponse)
async def notepad_page():
    return _NOTEPAD_HTML


@app.get("/notepad/events")
async def notepad_events():
    async def stream():
        q = asyncio.Queue(maxsize=30)
        notepad.add_listener(q)
        try:
            yield f"data: {json.dumps(notepad.get_notes_snapshot())}\n\n"
            while True:
                try:
                    payload = await asyncio.wait_for(q.get(), timeout=25)
                    yield f"data: {payload}\n\n"
                except asyncio.TimeoutError:
                    yield ": keepalive\n\n"
        finally:
            notepad.remove_listener(q)

    return StreamingResponse(stream(), media_type="text/event-stream")


@app.post("/notepad/clear")
async def clear_notepad_route():
    notepad.clear_notepad()
    return {"status": "cleared"}


@app.get("/notepad/content")
async def notepad_content():
    return notepad.get_notes_snapshot()


# ── Notepad HTML ───────────────────────────────────────────────────────────────
_NOTEPAD_HTML = """<!DOCTYPE html>
<html lang="en">
<head>
  <meta charset="UTF-8"/>
  <meta name="viewport" content="width=device-width, initial-scale=1.0"/>
  <title>Live Notepad</title>
  <style>
    *, *::before, *::after { box-sizing: border-box; margin: 0; padding: 0; }
    body {
      font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif;
      background: #0d0d0d; color: #e0e0e0;
      min-height: 100vh; display: flex; flex-direction: column;
      align-items: center; padding: 40px 16px; gap: 20px;
    }
    header { text-align: center; }
    header h1 { font-size: 26px; font-weight: 700; letter-spacing: -0.5px; }
    header p  { font-size: 13px; color: #666; margin-top: 4px; }

    #status-bar {
      display: flex; align-items: center; gap: 8px;
      font-size: 13px; color: #888;
    }
    #status-dot {
      width: 8px; height: 8px; border-radius: 50%;
      background: #444; transition: background 0.3s;
    }
    #status-dot.live  { background: #4CAF50; }
    #status-dot.error { background: #f44336; }

    #notepad {
      width: 660px; max-width: 100%;
      background: #141414; border: 1px solid #222;
      border-radius: 14px; min-height: 320px;
      padding: 20px; display: flex; flex-direction: column; gap: 10px;
    }

    #empty-hint {
      color: #3a3a3a; font-size: 14px;
      text-align: center; margin: auto;
      line-height: 1.6;
    }

    .note-item {
      display: flex; flex-direction: column; gap: 3px;
      padding: 12px 14px; background: #1a1a1a;
      border-radius: 8px; border-left: 3px solid #4CAF50;
      animation: fadeIn 0.25s ease;
    }
    @keyframes fadeIn {
      from { opacity: 0; transform: translateY(5px); }
      to   { opacity: 1; transform: none; }
    }
    .note-time { font-size: 11px; color: #555; font-family: monospace; }
    .note-text { font-size: 14px; color: #e0e0e0; line-height: 1.5; }

    #clear-btn {
      padding: 10px 24px; border-radius: 8px; border: none;
      font-size: 13px; font-weight: 600; cursor: pointer;
      background: #2a2a2a; color: #e0e0e0;
      transition: background 0.2s;
    }
    #clear-btn:hover { background: #f44336; color: #fff; }

    #note-count {
      font-size: 12px; color: #555;
    }
  </style>
</head>
<body>
  <header>
    <h1>Live Notepad</h1>
    <p>Notes appear here in real time as you speak to the avatar</p>
  </header>

  <div id="status-bar">
    <div id="status-dot"></div>
    <span id="status-text">Connecting…</span>
    <span style="margin-left:8px" id="note-count"></span>
  </div>

  <div id="notepad">
    <div id="notes-list"></div>
    <div id="empty-hint">
      No notes yet.<br/>
      Tell the avatar: <em>"Note that…"</em> or <em>"Write down that…"</em>
    </div>
  </div>

  <button id="clear-btn">Clear All Notes</button>

  <script>
    const notesList  = document.getElementById('notes-list');
    const emptyHint  = document.getElementById('empty-hint');
    const statusDot  = document.getElementById('status-dot');
    const statusText = document.getElementById('status-text');
    const noteCount  = document.getElementById('note-count');

    function esc(s) {
      return String(s)
        .replace(/&/g, '&amp;')
        .replace(/</g, '&lt;')
        .replace(/>/g, '&gt;');
    }

    function renderNotes(notes) {
      noteCount.textContent = notes.length ? `${notes.length} note${notes.length > 1 ? 's' : ''}` : '';
      if (!notes.length) {
        notesList.innerHTML  = '';
        emptyHint.style.display = 'block';
        return;
      }
      emptyHint.style.display = 'none';
      notesList.innerHTML = notes.map(n =>
        `<div class="note-item">
          <span class="note-time">${esc(n.time)}</span>
          <span class="note-text">${esc(n.text)}</span>
        </div>`
      ).join('');
      notesList.lastElementChild?.scrollIntoView({ behavior: 'smooth' });
    }

    const es = new EventSource('/notepad/events');
    es.onopen    = () => { statusDot.className = 'live';  statusText.textContent = 'Live'; };
    es.onmessage = (e) => { const d = JSON.parse(e.data); if (d.notes !== undefined) renderNotes(d.notes); };
    es.onerror   = () => { statusDot.className = 'error'; statusText.textContent = 'Reconnecting…'; };

    document.getElementById('clear-btn').addEventListener('click', () => {
      fetch('/notepad/clear', { method: 'POST' });
    });
  </script>
</body>
</html>"""


if __name__ == "__main__":
    uvicorn.run("main:app", host="0.0.0.0", port=PORT, reload=True)