diffdown-app / internal / handler / handler.go

package handler

import (
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"html/template"
	"log"
	"net/http"
	"net/url"
	"regexp"
	"strings"
	"time"

	"github.com/golang-jwt/jwt/v5"
	"github.com/gorilla/websocket"

	"github.com/limeleaf/diffdown/internal/atproto"
	"github.com/limeleaf/diffdown/internal/atproto/xrpc"
	"github.com/limeleaf/diffdown/internal/auth"
	"github.com/limeleaf/diffdown/internal/collaboration"
	"github.com/limeleaf/diffdown/internal/db"
	"github.com/limeleaf/diffdown/internal/model"
	"github.com/limeleaf/diffdown/internal/render"
)

const collectionDocument = "com.diffdown.document"

func randomID() string {
	b := make([]byte, 4)
	rand.Read(b)
	return hex.EncodeToString(b)
}

type Handler struct {
	DB               *db.DB
	Tmpls            map[string]*template.Template
	BaseURL          string
	CollaborationHub *collaboration.Hub
}

func New(database *db.DB, tmpls map[string]*template.Template, baseURL string, collabHub *collaboration.Hub) *Handler {
	return &Handler{DB: database, Tmpls: tmpls, BaseURL: baseURL, CollaborationHub: collabHub}
}

// --- Template helpers ---

type PageData struct {
	Title       string
	User        *model.User
	UserHandle  string
	Content     interface{}
	Error       string
	Description string
	OGImage     string
	Next        string
}

// DocumentEditData is passed to document_edit.html.
type DocumentEditData struct {
	*model.Document
	// AccessToken is the ATProto access token for WebSocket auth.
	// Empty string if user has no ATProto session.
	AccessToken string
	// IsOwner is true when the current user owns (created) the document.
	IsOwner bool
	// IsCollaborator is true when the current user is in the collaborators list.
	IsCollaborator bool
	// OwnerDID is the document owner's ATProto DID. Empty when IsOwner is true.
	// Used by collaborators to route save requests to the owner's PDS.
	OwnerDID string
}

type DashboardContent struct {
	OwnDocs    []*model.Document
	SharedDocs []*SharedDocument
}

type SharedDocument struct {
	RKey      string
	OwnerDID  string
	Title     string
	UpdatedAt string
	CreatedAt string
}

func (h *Handler) currentUser(r *http.Request) (*model.User, string) {
	uid := auth.UserIDFromContext(r.Context())
	if uid == "" {
		return nil, ""
	}
	u, err := h.DB.GetUserByID(uid)
	if err != nil {
		return nil, ""
	}
	session, err := h.DB.GetATProtoSession(uid)
	if err == nil && session != nil && session.DID != "" {
		handle, err := atproto.ResolveHandleFromDID(session.DID)
		if err == nil && handle != "" {
			return u, handle
		}
		return u, session.DID
	}
	return u, u.DID
}

func (h *Handler) render(w http.ResponseWriter, name string, data PageData) {
	tmpl, ok := h.Tmpls[name]
	if !ok {
		log.Printf("template not found: %s", name)
		http.Error(w, "Internal error", 500)
		return
	}
	w.Header().Set("Content-Type", "text/html; charset=utf-8")
	if err := tmpl.ExecuteTemplate(w, name, data); err != nil {
		log.Printf("template error: %v", err)
		http.Error(w, "Internal error", 500)
	}
}

func (h *Handler) jsonResponse(w http.ResponseWriter, data interface{}, statusCode int) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(statusCode)
	json.NewEncoder(w).Encode(data)
}

func (h *Handler) xrpcClient(userID string) (*xrpc.Client, error) {
	return xrpc.NewClient(h.DB, userID)
}

// --- Auth handlers ---

func (h *Handler) Logout(w http.ResponseWriter, r *http.Request) {
	auth.ClearSession(w, r)
	http.Redirect(w, r, "/", http.StatusSeeOther)
}

// --- Page handlers ---

func (h *Handler) AboutPage(w http.ResponseWriter, r *http.Request) {
	user, userHandle := h.currentUser(r)
	h.render(w, "about.html", PageData{
		Title:      "About",
		User:       user,
		UserHandle: userHandle,
	})
}

// --- Dashboard ---

func (h *Handler) Dashboard(w http.ResponseWriter, r *http.Request) {
	user, userHandle := h.currentUser(r)
	if user == nil {
		h.render(w, "landing.html", PageData{Title: "Diffdown"})
		return
	}

	empty := DashboardContent{OwnDocs: []*model.Document{}, SharedDocs: []*SharedDocument{}}

	client, err := h.xrpcClient(user.ID)
	if err != nil {
		log.Printf("Dashboard: xrpc client: %v", err)
		h.render(w, "documents.html", PageData{Title: "Documents", User: user, UserHandle: userHandle, Content: empty})
		return
	}

	// Own documents
	records, _, err := client.ListRecords(client.DID(), collectionDocument, 100, "")
	if err != nil {
		log.Printf("Dashboard: list records: %v", err)
		h.render(w, "documents.html", PageData{Title: "Documents", User: user, UserHandle: userHandle, Content: empty})
		return
	}
	ownDocs := make([]*model.Document, 0, len(records))
	for _, rec := range records {
		doc := &model.Document{}
		if err := json.Unmarshal(rec.Value, doc); err != nil {
			continue
		}
		doc.URI = rec.URI
		doc.CID = rec.CID
		doc.RKey = model.RKeyFromURI(rec.URI)
		ownDocs = append(ownDocs, doc)
	}

	// Shared documents
	session, _ := h.DB.GetATProtoSession(user.ID)
	var sharedDocs []*SharedDocument
	if session != nil {
		collabs, err := h.DB.GetCollaborations(session.DID)
		if err != nil {
			log.Printf("Dashboard: get collaborations: %v", err)
		}
		for _, c := range collabs {
			ownerUser, err := h.DB.GetUserByDID(c.OwnerDID)
			if err != nil {
				log.Printf("Dashboard: owner not found %s: %v", c.OwnerDID, err)
				continue
			}
			ownerClient, err := h.xrpcClient(ownerUser.ID)
			if err != nil {
				log.Printf("Dashboard: owner xrpc client %s: %v", c.OwnerDID, err)
				continue
			}
			value, _, err := ownerClient.GetRecord(c.OwnerDID, collectionDocument, c.DocumentRKey)
			if err != nil {
				log.Printf("Dashboard: get shared record %s/%s: %v", c.OwnerDID, c.DocumentRKey, err)
				continue
			}
			var doc model.Document
			if err := json.Unmarshal(value, &doc); err != nil {
				continue
			}
			sharedDocs = append(sharedDocs, &SharedDocument{
				RKey:      c.DocumentRKey,
				OwnerDID:  c.OwnerDID,
				Title:     doc.Title,
				UpdatedAt: doc.UpdatedAt,
				CreatedAt: doc.CreatedAt,
			})
		}
	}

	h.render(w, "documents.html", PageData{
		Title:      "Documents",
		User:       user,
		UserHandle: userHandle,
		Content:    DashboardContent{OwnDocs: ownDocs, SharedDocs: sharedDocs},
	})
}

// --- Document handlers ---

func (h *Handler) NewDocumentPage(w http.ResponseWriter, r *http.Request) {
	user, userHandle := h.currentUser(r)
	if user == nil {
		http.Redirect(w, r, "/auth/atproto", http.StatusSeeOther)
		return
	}
	h.render(w, "new_document.html", PageData{
		Title:      "New Document",
		User:       user,
		UserHandle: userHandle,
	})
}

// stripMarkdown removes basic markdown syntax to produce plain text for textContent.
var mdSyntaxRe = regexp.MustCompile(`(?m)^#{1,6}\s+|[*_~` + "`" + `\[\]()>]`)

func stripMarkdown(md string) string {
	return strings.TrimSpace(mdSyntaxRe.ReplaceAllString(md, ""))
}

func (h *Handler) NewDocumentSubmit(w http.ResponseWriter, r *http.Request) {
	user, _ := h.currentUser(r)
	if user == nil {
		http.Redirect(w, r, "/auth/atproto", http.StatusSeeOther)
		return
	}

	title := strings.TrimSpace(r.FormValue("title"))
	if title == "" {
		h.render(w, "new_document.html", PageData{Title: "New Document", User: user, Error: "Title is required"})
		return
	}

	client, err := h.xrpcClient(user.ID)
	if err != nil {
		log.Printf("NewDocumentSubmit: xrpc client: %v", err)
		h.render(w, "new_document.html", PageData{Title: "New Document", User: user, Error: "Could not connect to your PDS"})
		return
	}

	now := time.Now().UTC().Format(time.RFC3339)
	initialMD := fmt.Sprintf("# %s\n", title)
	doc := map[string]interface{}{
		"$type": "com.diffdown.document",
		"title": title,
		"content": map[string]interface{}{
			"$type":  "at.markpub.markdown",
			"flavor": "gfm",
			"text": map[string]interface{}{
				"rawMarkdown": initialMD,
			},
		},
		"textContent": stripMarkdown(initialMD),
		"createdAt":   now,
		"updatedAt":   now,
	}

	uri, _, err := client.CreateRecord(collectionDocument, doc)
	if err != nil {
		log.Printf("NewDocumentSubmit: create record: %v", err)
		h.render(w, "new_document.html", PageData{Title: "New Document", User: user, Error: "Could not create document"})
		return
	}

	rkey := model.RKeyFromURI(uri)
	http.Redirect(w, r, fmt.Sprintf("/docs/%s/edit", rkey), http.StatusSeeOther)
}

// documentView is the shared implementation for viewing a document given an ownerDID and rkey.
// isOwner should be true when the current user owns the document; it suppresses the ownerDID
// in the template so the edit button links to /docs/{rkey}/edit rather than the collaborator URL.
func (h *Handler) documentView(w http.ResponseWriter, r *http.Request, ownerUserID, ownerDID, rkey string, isOwner bool, userHandle string) {
	client, err := h.xrpcClient(ownerUserID)
	if err != nil {
		http.Error(w, "Could not connect to PDS", 500)
		return
	}

	value, _, err := client.GetRecord(ownerDID, collectionDocument, rkey)
	if err != nil {
		http.NotFound(w, r)
		return
	}

	doc := &model.Document{}
	if err := json.Unmarshal(value, doc); err != nil {
		http.Error(w, "Invalid document", 500)
		return
	}
	doc.RKey = rkey

	var rendered string
	if doc.Content != nil {
		rendered, _ = render.Markdown([]byte(doc.Content.Text.RawMarkdown))
	}

	user, _ := h.currentUser(r)
	type DocumentViewData struct {
		Doc      *model.Document
		Rendered template.HTML
		OwnerDID string // non-empty when viewing a collaborator's document
	}
	// Only set OwnerDID when the viewer is not the owner; the template uses
	// a non-empty OwnerDID to generate the collaborator edit URL.
	templateOwnerDID := ownerDID
	if isOwner {
		templateOwnerDID = ""
	}
	h.render(w, "document_view.html", PageData{
		Title:      doc.Title,
		User:       user,
		UserHandle: userHandle,
		Content:    DocumentViewData{Doc: doc, Rendered: template.HTML(rendered), OwnerDID: templateOwnerDID},
	})
}

// DocumentView renders a document as HTML (owner viewing their own doc).
func (h *Handler) DocumentView(w http.ResponseWriter, r *http.Request) {
	user, userHandle := h.currentUser(r)
	if user == nil {
		http.Redirect(w, r, "/auth/atproto", http.StatusSeeOther)
		return
	}

	rkey := r.PathValue("rkey")
	client, err := h.xrpcClient(user.ID)
	if err != nil {
		http.Error(w, "Could not connect to PDS", 500)
		return
	}
	h.documentView(w, r, user.ID, client.DID(), rkey, true, userHandle)
}

// CollaboratorDocumentView renders a document owned by another user (collaborator access).
func (h *Handler) CollaboratorDocumentView(w http.ResponseWriter, r *http.Request) {
	user, userHandle := h.currentUser(r)
	if user == nil {
		http.Redirect(w, r, "/auth/atproto", http.StatusSeeOther)
		return
	}

	ownerDID := r.PathValue("did")
	rkey := r.PathValue("rkey")

	ownerUser, err := h.DB.GetUserByDID(ownerDID)
	if err != nil {
		http.NotFound(w, r)
		return
	}

	h.documentView(w, r, ownerUser.ID, ownerDID, rkey, false, userHandle)
}

// documentEdit is the shared implementation for the edit page.
// ownerUserID/ownerDID identify whose PDS holds the document; isOwner is true for the creator.
func (h *Handler) documentEdit(w http.ResponseWriter, r *http.Request, user *model.User, ownerUserID, ownerDID, rkey string, isOwner bool, userHandle string) {
	ownerClient, err := h.xrpcClient(ownerUserID)
	if err != nil {
		http.Error(w, "Could not connect to PDS", 500)
		return
	}

	value, _, err := ownerClient.GetRecord(ownerDID, collectionDocument, rkey)
	if err != nil {
		http.NotFound(w, r)
		return
	}

	doc := &model.Document{}
	if err := json.Unmarshal(value, doc); err != nil {
		http.Error(w, "Invalid document", 500)
		return
	}
	doc.RKey = rkey

	editData := &DocumentEditData{Document: doc, IsOwner: isOwner}
	if !isOwner {
		editData.OwnerDID = ownerDID
	}
	if session, err := h.DB.GetATProtoSession(user.ID); err == nil && session != nil {
		editData.AccessToken = session.AccessToken
		userDID := session.DID
		for _, did := range doc.Collaborators {
			if did == userDID {
				editData.IsCollaborator = true
				break
			}
		}
	}

	h.render(w, "document_edit.html", PageData{
		Title:      "Edit " + doc.Title,
		User:       user,
		UserHandle: userHandle,
		Content:    editData,
	})
}

// DocumentEdit renders the editor for a document (owner).
func (h *Handler) DocumentEdit(w http.ResponseWriter, r *http.Request) {
	user, userHandle := h.currentUser(r)
	if user == nil {
		http.Redirect(w, r, "/auth/atproto", http.StatusSeeOther)
		return
	}

	rkey := r.PathValue("rkey")
	client, err := h.xrpcClient(user.ID)
	if err != nil {
		http.Error(w, "Could not connect to PDS", 500)
		return
	}

	h.documentEdit(w, r, user, user.ID, client.DID(), rkey, true, userHandle)
}

// CollaboratorDocumentEdit renders the editor for a document owned by another user.
func (h *Handler) CollaboratorDocumentEdit(w http.ResponseWriter, r *http.Request) {
	user, userHandle := h.currentUser(r)
	if user == nil {
		http.Redirect(w, r, "/auth/atproto", http.StatusSeeOther)
		return
	}

	ownerDID := r.PathValue("did")
	rkey := r.PathValue("rkey")

	ownerUser, err := h.DB.GetUserByDID(ownerDID)
	if err != nil {
		http.NotFound(w, r)
		return
	}

	h.documentEdit(w, r, user, ownerUser.ID, ownerDID, rkey, false, userHandle)
}

// APIDocumentSave saves a document to the PDS.
func (h *Handler) APIDocumentSave(w http.ResponseWriter, r *http.Request) {
	user, _ := h.currentUser(r)
	if user == nil {
		http.Error(w, "Unauthorized", 401)
		return
	}

	rkey := r.PathValue("rkey")
	var req struct {
		Content  string `json:"content"`
		Title    string `json:"title"`
		OwnerDID string `json:"ownerDID"` // non-empty when saving on behalf of another user (collaborator)
	}
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, "Bad request", 400)
		return
	}

	// For collaborators, save to the document owner's PDS, not the collaborator's.
	var client *xrpc.Client
	var repoDID string
	if req.OwnerDID != "" {
		ownerUser, err := h.DB.GetUserByDID(req.OwnerDID)
		if err != nil {
			log.Printf("APIDocumentSave: get owner by DID %s: %v", req.OwnerDID, err)
			http.Error(w, "Document owner not found", 404)
			return
		}
		client, err = h.xrpcClient(ownerUser.ID)
		if err != nil {
			http.Error(w, "Could not connect to owner PDS", 500)
			return
		}
		repoDID = req.OwnerDID
	} else {
		var err error
		client, err = h.xrpcClient(user.ID)
		if err != nil {
			http.Error(w, "Could not connect to PDS", 500)
			return
		}
		repoDID = client.DID()
	}

	// Fetch existing record to preserve fields
	value, _, err := client.GetRecord(repoDID, collectionDocument, rkey)
	if err != nil {
		http.Error(w, "Document not found", 404)
		return
	}

	var existing map[string]interface{}
	json.Unmarshal(value, &existing)

	// Update content
	title := req.Title
	if title == "" {
		if t, ok := existing["title"].(string); ok {
			title = t
		}
	}

	now := time.Now().UTC().Format(time.RFC3339)
	existing["title"] = title
	existing["content"] = map[string]interface{}{
		"$type":  "at.markpub.markdown",
		"flavor": "gfm",
		"text": map[string]interface{}{
			"rawMarkdown": req.Content,
		},
	}
	existing["textContent"] = stripMarkdown(req.Content)
	existing["updatedAt"] = now

	_, _, err = client.PutRecord(collectionDocument, rkey, existing)
	if err != nil {
		log.Printf("APIDocumentSave: put record: %v", err)
		http.Error(w, "Save failed", 500)
		return
	}

	h.jsonResponse(w, map[string]string{"status": "ok"}, http.StatusOK)
}

// APIDocumentAutoSave is the same as save, called on debounce from editor.
func (h *Handler) APIDocumentAutoSave(w http.ResponseWriter, r *http.Request) {
	h.APIDocumentSave(w, r)
}

// APIDocumentDelete deletes a document from the PDS.
func (h *Handler) APIDocumentDelete(w http.ResponseWriter, r *http.Request) {
	user, _ := h.currentUser(r)
	if user == nil {
		http.Error(w, "Unauthorized", 401)
		return
	}

	rkey := r.PathValue("rkey")
	client, err := h.xrpcClient(user.ID)
	if err != nil {
		http.Error(w, "Could not connect to PDS", 500)
		return
	}

	if err := client.DeleteRecord(collectionDocument, rkey); err != nil {
		log.Printf("APIDocumentDelete: %v", err)
		http.Error(w, "Delete failed", 500)
		return
	}

	h.jsonResponse(w, map[string]string{"status": "ok"}, http.StatusOK)
}

// DocumentInvite creates an invite link for a document.
func (h *Handler) DocumentInvite(w http.ResponseWriter, r *http.Request) {
	user, _ := h.currentUser(r)
	if user == nil {
		http.Redirect(w, r, "/auth/atproto", http.StatusSeeOther)
		return
	}

	rkey := r.PathValue("rkey")
	if rkey == "" {
		http.Error(w, "Invalid document", http.StatusBadRequest)
		return
	}

	client, err := h.xrpcClient(user.ID)
	if err != nil {
		log.Printf("DocumentInvite: xrpc client: %v", err)
		h.render(w, "document_edit.html", PageData{Error: "Not authenticated with ATProto"})
		return
	}

	value, _, err := client.GetRecord(client.DID(), collectionDocument, rkey)
	if err != nil {
		http.Error(w, "Document not found", http.StatusNotFound)
		return
	}

	doc := &model.Document{}
	if err := json.Unmarshal(value, doc); err != nil {
		http.Error(w, "Invalid document", http.StatusInternalServerError)
		return
	}
	doc.RKey = rkey

	// The document was fetched via client.DID(), so the current user is always the owner.
	if len(doc.Collaborators) >= 5 {
		http.Error(w, "Maximum collaborators reached", http.StatusBadRequest)
		return
	}

	invite, err := collaboration.CreateInvite(h.DB, rkey, client.DID())
	if err != nil {
		log.Printf("DocumentInvite: create invite: %v", err)
		http.Error(w, "Failed to create invite", http.StatusInternalServerError)
		return
	}

	inviteLink := fmt.Sprintf("%s/docs/%s/accept?invite=%s", h.BaseURL, rkey, invite.Token)
	h.jsonResponse(w, map[string]string{"inviteLink": inviteLink}, http.StatusOK)
}

// AcceptInvite handles an invite acceptance.
func (h *Handler) AcceptInvite(w http.ResponseWriter, r *http.Request) {
	user, _ := h.currentUser(r)
	if user == nil {
		// Preserve invite token through the login redirect.
		http.Redirect(w, r, "/auth/atproto?next="+url.QueryEscape(r.URL.String()), http.StatusSeeOther)
		return
	}

	rKey := r.PathValue("rkey")
	inviteToken := r.URL.Query().Get("invite")
	if inviteToken == "" {
		http.Error(w, "Invalid invite", http.StatusBadRequest)
		return
	}

	invite, err := collaboration.ValidateInvite(h.DB, inviteToken, rKey)
	if err != nil {
		log.Printf("AcceptInvite: validate invite rkey=%s: %v", rKey, err)
		http.Error(w, "Invite not found, already used, or expired.", http.StatusBadRequest)
		return
	}

	// The collaborator's session — needed to get their DID.
	collabSession, err := h.DB.GetATProtoSession(user.ID)
	if err != nil || collabSession == nil {
		http.Redirect(w, r, "/auth/atproto?next="+url.QueryEscape(r.URL.String()), http.StatusSeeOther)
		return
	}

	// Fetch and update the document from the OWNER's PDS, not the collaborator's.
	// The invite records the owner's DID in CreatedBy.
	ownerUser, err := h.DB.GetUserByDID(invite.CreatedBy)
	if err != nil {
		log.Printf("AcceptInvite: get owner by DID %s: %v", invite.CreatedBy, err)
		http.Error(w, "Document owner not found", http.StatusInternalServerError)
		return
	}

	ownerClient, err := h.xrpcClient(ownerUser.ID)
	if err != nil {
		log.Printf("AcceptInvite: owner xrpc client: %v", err)
		http.Error(w, "Failed to connect to ATProto", http.StatusInternalServerError)
		return
	}

	doc, err := ownerClient.GetDocument(rKey)
	if err != nil {
		log.Printf("AcceptInvite: get document: %v", err)
		http.Error(w, "Document not found", http.StatusNotFound)
		return
	}

	// Already a collaborator — redirect to the owner-scoped URL.
	for _, c := range doc.Collaborators {
		if c == collabSession.DID {
			if err := h.DB.AddCollaboration(collabSession.DID, invite.CreatedBy, rKey); err != nil {
				log.Printf("AcceptInvite: record collaboration (existing): %v", err)
			}
			http.Redirect(w, r, "/docs/"+invite.CreatedBy+"/"+rKey, http.StatusSeeOther)
			return
		}
	}

	// Add collaborator DID and PUT back to owner's PDS.
	doc.Collaborators = append(doc.Collaborators, collabSession.DID)
	if _, _, err = ownerClient.PutDocument(rKey, doc); err != nil {
		log.Printf("AcceptInvite: put document: %v", err)
		http.Error(w, "Failed to add collaborator", http.StatusInternalServerError)
		return
	}

	h.DB.DeleteInvite(invite.Token)

	if err := h.DB.AddCollaboration(collabSession.DID, invite.CreatedBy, rKey); err != nil {
		log.Printf("AcceptInvite: record collaboration: %v", err)
		// Non-fatal — collaborator is already on the document, just redirect.
	}

	// Redirect to owner-scoped document URL so the view handler knows whose PDS to query.
	http.Redirect(w, r, "/docs/"+invite.CreatedBy+"/"+rKey, http.StatusSeeOther)
}

// --- API: Comments ---

func (h *Handler) CommentCreate(w http.ResponseWriter, r *http.Request) {
	user, _ := h.currentUser(r)
	if user == nil {
		http.Error(w, "Unauthorized", http.StatusUnauthorized)
		return
	}

	rKey := r.PathValue("rkey")
	if rKey == "" {
		http.Error(w, "Invalid document", http.StatusBadRequest)
		return
	}

	var req struct {
		OwnerDID    string `json:"ownerDID"`
		ParagraphID string `json:"paragraphId"`
		Text        string `json:"text"`
	}
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, "Invalid request", http.StatusBadRequest)
		return
	}
	if req.Text == "" {
		http.Error(w, "Comment text required", http.StatusBadRequest)
		return
	}

	session, err := h.DB.GetATProtoSession(user.ID)
	if err != nil || session == nil {
		http.Error(w, "Not authenticated with ATProto", http.StatusUnauthorized)
		return
	}

	ownerUserID := user.ID
	ownerDID := session.DID
	if req.OwnerDID != "" && req.OwnerDID != session.DID {
		ownerUser, err := h.DB.GetUserByDID(req.OwnerDID)
		if err != nil {
			http.Error(w, "Owner not found", http.StatusBadRequest)
			return
		}
		ownerUserID = ownerUser.ID
		ownerDID = req.OwnerDID
	}

	ownerClient, err := h.xrpcClient(ownerUserID)
	if err != nil {
		http.Error(w, "Failed to connect to owner PDS", http.StatusInternalServerError)
		return
	}

	value, _, err := ownerClient.GetRecord(ownerDID, collectionDocument, rKey)
	if err != nil {
		log.Printf("CommentCreate: GetRecord: %v", err)
		http.Error(w, "Document not found", http.StatusNotFound)
		return
	}
	var doc model.Document
	if err := json.Unmarshal(value, &doc); err != nil {
		http.Error(w, "Failed to parse document", http.StatusInternalServerError)
		return
	}

	authorHandle, _ := atproto.ResolveHandleFromDID(session.DID)

	paragraphID := req.ParagraphID
	if paragraphID == "" {
		paragraphID = "general"
	}

	comment := model.EmbeddedComment{
		ID:           randomID(),
		ParagraphID:  paragraphID,
		Text:         req.Text,
		Author:       session.DID,
		AuthorHandle: authorHandle,
		CreatedAt:    time.Now().UTC().Format(time.RFC3339),
	}
	doc.Comments = append(doc.Comments, comment)

	if _, _, err := ownerClient.PutDocument(rKey, &doc); err != nil {
		log.Printf("CommentCreate: PutDocument: %v", err)
		http.Error(w, "Failed to save comment", http.StatusInternalServerError)
		return
	}

	h.jsonResponse(w, comment, http.StatusCreated)
}

func (h *Handler) CommentList(w http.ResponseWriter, r *http.Request) {
	rKey := r.PathValue("rkey")
	if rKey == "" {
		http.Error(w, "Invalid document", http.StatusBadRequest)
		return
	}

	user, _ := h.currentUser(r)
	if user == nil {
		http.Error(w, "Unauthorized", http.StatusUnauthorized)
		return
	}

	session, err := h.DB.GetATProtoSession(user.ID)
	if err != nil || session == nil {
		http.Error(w, "Not authenticated with ATProto", http.StatusUnauthorized)
		return
	}

	ownerUserID := user.ID
	ownerDID := session.DID
	if qOwner := r.URL.Query().Get("ownerDID"); qOwner != "" && qOwner != session.DID {
		ownerUser, err := h.DB.GetUserByDID(qOwner)
		if err != nil {
			http.Error(w, "Owner not found", http.StatusBadRequest)
			return
		}
		ownerUserID = ownerUser.ID
		ownerDID = qOwner
	}

	ownerClient, err := h.xrpcClient(ownerUserID)
	if err != nil {
		http.Error(w, "Failed to connect to owner PDS", http.StatusInternalServerError)
		return
	}

	value, _, err := ownerClient.GetRecord(ownerDID, collectionDocument, rKey)
	if err != nil {
		log.Printf("CommentList: GetRecord: %v", err)
		http.Error(w, "Document not found", http.StatusNotFound)
		return
	}
	var doc model.Document
	if err := json.Unmarshal(value, &doc); err != nil {
		http.Error(w, "Failed to parse document", http.StatusInternalServerError)
		return
	}

	comments := doc.Comments
	if comments == nil {
		comments = []model.EmbeddedComment{}
	}
	h.jsonResponse(w, comments, http.StatusOK)
}

// --- API: Render markdown ---

func (h *Handler) APIRender(w http.ResponseWriter, r *http.Request) {
	var req struct {
		Content string `json:"content"`
	}
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, "Bad request", 400)
		return
	}

	rendered, err := render.Markdown([]byte(req.Content))
	if err != nil {
		http.Error(w, "Render error", 500)
		return
	}

	h.jsonResponse(w, map[string]string{"html": rendered}, http.StatusOK)
}

var upgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool { return true },
}

func (h *Handler) CollaboratorWebSocket(w http.ResponseWriter, r *http.Request) {
	rKey := r.PathValue("rkey")
	if rKey == "" {
		http.Error(w, "Invalid document", http.StatusBadRequest)
		return
	}

	accessToken := r.URL.Query().Get("access_token")
	dpopProof := r.URL.Query().Get("dpop_proof")
	if accessToken == "" || dpopProof == "" {
		http.Error(w, "Missing auth tokens", http.StatusUnauthorized)
		return
	}

	did, name, err := h.validateWSToken(accessToken, dpopProof)
	if err != nil {
		http.Error(w, "Invalid tokens", http.StatusUnauthorized)
		return
	}

	user, err := h.DB.GetUserByDID(did)
	if err != nil {
		http.Error(w, "No user found", http.StatusUnauthorized)
		return
	}

	session, err := h.DB.GetATProtoSession(user.ID)
	if err != nil || session == nil {
		http.Error(w, "No ATProto session", http.StatusUnauthorized)
		return
	}

	// If owner_did is provided, fetch the document from the owner's PDS
	// (used by collaborators whose copy lives on a different PDS).
	ownerDID := r.URL.Query().Get("owner_did")
	var docClient *xrpc.Client
	var docRepoDID string
	if ownerDID != "" {
		ownerUser, err := h.DB.GetUserByDID(ownerDID)
		if err != nil {
			log.Printf("CollaboratorWebSocket: get owner by DID %s: %v", ownerDID, err)
			http.Error(w, "Document owner not found", http.StatusForbidden)
			return
		}
		docClient, err = h.xrpcClient(ownerUser.ID)
		if err != nil {
			http.Error(w, "Failed to connect to owner PDS", http.StatusInternalServerError)
			return
		}
		docRepoDID = ownerDID
	} else {
		docClient, err = h.xrpcClient(session.UserID)
		if err != nil {
			http.Error(w, "Failed to connect to ATProto", http.StatusInternalServerError)
			return
		}
		docRepoDID = did
	}

	value, _, err := docClient.GetRecord(docRepoDID, collectionDocument, rKey)
	if err != nil {
		http.Error(w, "Document not found", http.StatusNotFound)
		return
	}
	doc := &model.Document{}
	if err := json.Unmarshal(value, doc); err != nil {
		http.Error(w, "Invalid document", http.StatusInternalServerError)
		return
	}

	// Owner always has access; collaborators must be in the collaborators list.
	isCollaborator := did == docRepoDID
	for _, c := range doc.Collaborators {
		if c == did {
			isCollaborator = true
			break
		}
	}
	if !isCollaborator {
		http.Error(w, "Not a collaborator", http.StatusForbidden)
		return
	}

	color := colorFromDID(did)

	// Fetch handle and avatar from ATProto profile; best-effort, empty on failure.
	var avatar, handle string
	if profile, err := atproto.ResolveProfile(did); err == nil && profile != nil {
		avatar = profile.Avatar
		handle = profile.Handle
	}

	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Printf("WebSocket upgrade failed: %v", err)
		return
	}

	room := h.CollaborationHub.GetOrCreateRoom(rKey)
	wsClient := collaboration.NewClient(h.CollaborationHub, conn, did, name, handle, color, avatar, rKey)
	room.RegisterClient(wsClient)

	go wsClient.WritePump()
	wsClient.ReadPump()
}

func (h *Handler) validateWSToken(accessToken, dpopProof string) (string, string, error) {
	claims := &jwt.MapClaims{}
	parser := jwt.Parser{}
	_, _, err := parser.ParseUnverified(accessToken, claims)
	if err != nil {
		return "", "", fmt.Errorf("parse token: %w", err)
	}

	did, ok := (*claims)["sub"].(string)
	if !ok {
		return "", "", fmt.Errorf("no sub in token")
	}

	user, err := h.DB.GetUserByDID(did)
	if err != nil {
		return "", "", fmt.Errorf("user not found: %w", err)
	}

	session, err := h.DB.GetATProtoSession(user.ID)
	if err != nil {
		return "", "", fmt.Errorf("session not found: %w", err)
	}

	if time.Now().After(session.ExpiresAt) {
		return "", "", fmt.Errorf("session expired")
	}

	name, _ := (*claims)["name"].(string)
	return did, name, nil
}

func colorFromDID(did string) string {
	colors := []string{"#e74c3c", "#3498db", "#2ecc71", "#9b59b6", "#f39c12"}
	hash := 0
	for _, c := range did {
		hash += int(c)
	}
	return colors[hash%len(colors)]
}

// SubmitSteps receives ProseMirror steps from a collaborator, appends them
// to the step log, and broadcasts confirmed steps to the room.
//
// POST /api/docs/{rkey}/steps
// Body: {"clientVersion": N, "steps": ["...json..."], "clientID": "did:..."}
// Response 200: {"version": N}
// Response 409: {"version": N, "steps": ["...json..."]} — client must rebase
func (h *Handler) SubmitSteps(w http.ResponseWriter, r *http.Request) {
	user, _ := h.currentUser(r)
	if user == nil {
		http.Error(w, "Unauthorized", http.StatusUnauthorized)
		return
	}
	rkey := r.PathValue("rkey")

	var body struct {
		ClientVersion int      `json:"clientVersion"`
		Steps         []string `json:"steps"`
		ClientID      string   `json:"clientID"`
	}
	if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
		http.Error(w, "Bad request", http.StatusBadRequest)
		return
	}
	if len(body.Steps) == 0 {
		http.Error(w, "No steps", http.StatusBadRequest)
		return
	}

	newVersion, err := h.DB.AppendSteps(rkey, body.ClientVersion, body.Steps, body.ClientID)
	if err != nil {
		// Version conflict — return steps the client missed.
		missed, dbErr := h.DB.GetStepsSince(rkey, body.ClientVersion)
		if dbErr != nil {
			log.Printf("SubmitSteps: GetStepsSince: %v", dbErr)
			http.Error(w, "Internal error", http.StatusInternalServerError)
			return
		}
		currentVersion, _ := h.DB.GetDocVersion(rkey)
		stepJSONs := make([]string, len(missed))
		for i, s := range missed {
			stepJSONs[i] = s.JSON
		}
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusConflict)
		json.NewEncoder(w).Encode(map[string]interface{}{
			"version": currentVersion,
			"steps":   stepJSONs,
		})
		return
	}

	// Broadcast to other room members via WebSocket.
	if room := h.CollaborationHub.GetRoom(rkey); room != nil {
		type stepsMsg struct {
			Type     string   `json:"type"`
			Steps    []string `json:"steps"`
			Version  int      `json:"version"`
			ClientID string   `json:"clientID"`
		}
		data, _ := json.Marshal(stepsMsg{
			Type:     "steps",
			Steps:    body.Steps,
			Version:  newVersion,
			ClientID: body.ClientID,
		})
		room.Broadcast(data)
	}

	h.jsonResponse(w, map[string]int{"version": newVersion}, http.StatusOK)
}

// GetSteps returns all steps since the given version.
//
// GET /api/docs/{rkey}/steps?since={v}
// Response 200: {"version": N, "steps": ["...json..."]}
func (h *Handler) GetSteps(w http.ResponseWriter, r *http.Request) {
	user, _ := h.currentUser(r)
	if user == nil {
		http.Error(w, "Unauthorized", http.StatusUnauthorized)
		return
	}
	rkey := r.PathValue("rkey")

	sinceStr := r.URL.Query().Get("since")
	var since int
	if sinceStr != "" {
		fmt.Sscanf(sinceStr, "%d", &since)
	}

	rows, err := h.DB.GetStepsSince(rkey, since)
	if err != nil {
		log.Printf("GetSteps: %v", err)
		http.Error(w, "Internal error", http.StatusInternalServerError)
		return
	}
	version, _ := h.DB.GetDocVersion(rkey)
	stepJSONs := make([]string, len(rows))
	for i, s := range rows {
		stepJSONs[i] = s.JSON
	}
	h.jsonResponse(w, map[string]interface{}{
		"version": version,
		"steps":   stepJSONs,
	}, http.StatusOK)
}