(ns outliner.boundaries.firestore
  (:require [clojure.set :as set]))

;; --- Firestore Data Model ---
;;
;; To support large outlines and granular updates, the data is split across
;; two locations in Firestore:
;;
;; 1. The "main" document:
;;    Path: /users/{user-id}/documents/main
;;    Contains:
;;    - 'root': A vector of top-level node IDs.
;;    - 'metadata': version, lastModified timestamp.
;;    - (Legacy) 'data': May contain a 'nodes' map for older documents.
;;
;; 2. The "nodes" subcollection:
;;    Path: /users/{user-id}/documents/main/nodes/{node-id}
;;    Each document represents a single node with its text, annotations,
;;    children IDs, and parent ID.
;;
;; This split allows us to update only the modified nodes in a batch,
;; rather than re-writing the entire outline every time a single character
;; changes.

;; Data Transformation Helpers
(defn- get-data [snapshot]
  (if (and snapshot (.exists snapshot))
    (.data snapshot)
    nil))

(defn keywordize-annotation [ann]
  (cond-> ann
    (string? (:type ann)) (update :type keyword)
    (string? (:style ann)) (update :style keyword)))

(defn keywordize-node [node]
  (if (vector? (:annotations node))
    (update node :annotations #(mapv keywordize-annotation %))
    node))

(defn keywordize-doc [doc]
  (if-let [nodes (:nodes doc)]
    (assoc doc :nodes (into {} (map (fn [[id node]] [id (keywordize-node node)]) nodes)))
    doc))

(defn- fix-stored-nodes [nodes]
  (if (map? nodes)
    (into {} (map (fn [[k v]] [(name k) (keywordize-node v)]) nodes))
    {}))

(defn consistent-doc? [doc]
  (let [root (:root doc)
        nodes (:nodes doc)]
    (letfn [(check [id seen]
              (if (contains? seen id)
                true
                (if-let [node (get nodes id)]
                  (every? #(check % (conj seen id)) (:children node))
                  false)))]
      (every? #(check % #{}) root))))

(defn- get-node-diff [old-nodes new-nodes]
  (let [old-nodes (or old-nodes {})
        new-nodes (or new-nodes {})
        upsert (reduce-kv (fn [acc id node]
                            (if (not= node (get old-nodes id))
                              (assoc acc id node)
                              acc))
                          {}
                          new-nodes)
        delete (filter #(not (contains? new-nodes %)) (keys old-nodes))]
    {:upsert upsert :delete delete}))

(defn- commit-batches! [ops on-success on-error]
  (if (empty? ops)
    (when on-success (on-success))
    (let [db js/firebaseDB
          write-batch-fn (when (and js/fb (.-firestore js/fb)) (-> js/fb .-firestore .-writeBatch))]
      (if-not (and db write-batch-fn)
        (when on-error (on-error (js/Error. "Firestore not available for commit")))
        (let [chunks (partition-all 500 ops)
              promises (map (fn [chunk]
                              (let [batch (write-batch-fn db)]
                                (doseq [op chunk]
                                  (case (:type op)
                                    :set (.set batch (:ref op) (clj->js (:data op)) (or (:options op) #js {}))
                                    :delete (.delete batch (:ref op))))
                                (.commit batch)))
                            chunks)]
          (-> (.all js/Promise (to-array promises))
              (.then (fn [_] (when on-success (on-success))))
              (.catch (fn [e]
                        (js/console.error "Batch commit error:" e)
                        (when on-error (on-error e))))))))))


;; Firestore functions
(defn save-doc! [user-id new-doc old-doc & [on-success on-error]]
  (let [db js/firebaseDB
        doc-func (when (and js/fb (.-firestore js/fb)) (-> js/fb .-firestore .-doc))]
    (if-not (and db doc-func)
      (when on-error (on-error (js/Error. "Firestore not available for save-doc!")))
      (let [main-doc-ref (doc-func db "users" user-id "documents" "main")
            node-diff (get-node-diff (:nodes old-doc) (:nodes new-doc))
            metadata (or (:metadata new-doc)
                         {:lastModified (.getTime (js/Date.))
                          :version 1
                          :clientVersion "1.0.0"})
            ops [{:type :set
                  :ref main-doc-ref
                  :data {:root (:root new-doc)
                         :metadata metadata}}]]
        (let [ops (reduce-kv (fn [acc id node]
                               (conj acc {:type :set
                                          :ref (doc-func db "users" user-id "documents" "main" "nodes" (str id))
                                          :data node}))
                             ops
                             (:upsert node-diff))
              ops (reduce (fn [acc id]
                            (conj acc {:type :delete
                                       :ref (doc-func db "users" user-id "documents" "main" "nodes" (str id))}))
                          ops
                          (:delete node-diff))]
          (commit-batches! ops on-success on-error))))))

(defn- extract-load-results [results]
  (let [main-snap (aget results 0)
        nodes-snap (aget results 1)
        main-data (js->clj (get-data main-snap) :keywordize-keys true)
        old-nodes (get-in main-data [:data :nodes])
        root (or (:root main-data) (get-in main-data [:data :root]))
        sub-nodes (reduce (fn [acc doc]
                            (assoc acc (.-id doc) (keywordize-node (js->clj (.data doc) :keywordize-keys true))))
                          {}
                          (.-docs nodes-snap))]
    (if (or root (seq sub-nodes) (seq old-nodes))
      {:root root
       :nodes (merge (fix-stored-nodes old-nodes) sub-nodes)
       :metadata (:metadata main-data)}
      nil)))

(defn load-doc [user-id callback]
  (let [db js/firebaseDB
        firestore-api (when js/fb (.-firestore js/fb))
        doc-func (when firestore-api (.-doc firestore-api))
        col-func (when firestore-api (.-collection firestore-api))
        get-doc (when firestore-api (.-getDoc firestore-api))
        get-docs (when firestore-api (.-getDocs firestore-api))]
    (if-not (and db doc-func col-func get-doc get-docs)
      (do (js/console.warn "Firestore not available for load-doc")
          (callback nil))
      (let [doc-ref (doc-func db "users" user-id "documents" "main")
            col-ref (col-func db "users" user-id "documents" "main" "nodes")]
        (-> (.all js/Promise #js [(get-doc doc-ref) (get-docs col-ref)])
            (.then (fn [results] (callback (extract-load-results results))))
            (.catch (fn [error]
                      (js/console.error "Error loading from Firestore" error)
                      (callback nil))))))))

(defn- process-main-snapshot [snap state emit-fn]
  (let [pending? (.-hasPendingWrites (.-metadata snap))]
    (if-let [data (get-data snap)]
      (swap! state assoc :main (js->clj data :keywordize-keys true) :main-loaded? true :main-pending? pending?)
      (swap! state assoc :main nil :main-loaded? true :main-pending? pending?))
    (emit-fn)))

(defn- process-nodes-snapshot [snap state emit-fn]
  (let [pending? (.-hasPendingWrites (.-metadata snap))
        nodes (reduce (fn [acc doc]
                        (assoc acc (.-id doc) (keywordize-node (js->clj (.data doc) :keywordize-keys true))))
                      {}
                      (.-docs snap))]
    (swap! state assoc :nodes nodes :nodes-loaded? true :nodes-pending? pending?)
    (emit-fn)))

(defn fetch-orphan-count [user-id callback]
  (let [db js/firebaseDB
        firestore-api (when js/fb (.-firestore js/fb))
        doc-func (when firestore-api (.-doc firestore-api))
        col-func (when firestore-api (.-collection firestore-api))
        get-doc (when firestore-api (.-getDoc firestore-api))
        get-docs (when firestore-api (.-getDocs firestore-api))]
    (if-not (and db doc-func col-func get-doc get-docs)
      (do (js/console.warn "Firestore not available for fetch-orphan-count")
          (callback 0))
      (let [doc-ref (doc-func db "users" user-id "documents" "main")
            col-ref (col-func db "users" user-id "documents" "main" "nodes")]
        (-> (.all js/Promise #js [(get-doc doc-ref) (get-docs col-ref)])
            (.then (fn [results]
                     (let [main-snap (aget results 0)
                           nodes-snap (aget results 1)
                           main-data (js->clj (get-data main-snap) :keywordize-keys true)
                           root (or (:root main-data) (get-in main-data [:data :root]) [])
                           all-nodes (reduce (fn [acc doc]
                                               (assoc acc (.-id doc) (js->clj (.data doc) :keywordize-keys true)))
                                             {}
                                             (.-docs nodes-snap))
                           ;; We also consider nodes that might be in the legacy 'data/nodes' field
                           legacy-nodes (fix-stored-nodes (get-in main-data [:data :nodes]))
                           nodes (merge legacy-nodes all-nodes)]
                       (letfn [(get-reachable [ids seen]
                                 (reduce (fn [s id]
                                           (if (contains? s id)
                                             s
                                             (let [node (get nodes id)
                                                   children (:children node)]
                                               (get-reachable children (conj s id)))))
                                         seen
                                         ids))]
                         (let [reachable (get-reachable root #{})
                               all-ids (set (keys nodes))
                               orphans (set/difference all-ids reachable)]
                           (callback (count orphans)))))))
            (.catch (fn [error]
                      (js/console.error "Error fetching orphan count:" error)
                      (callback 0))))))))

(defn listen-to-doc [user-id callback]
  (let [db js/firebaseDB
        firestore-api (when js/fb (.-firestore js/fb))
        doc-func (when firestore-api (.-doc firestore-api))
        col-func (when firestore-api (.-collection firestore-api))
        on-snapshot (when firestore-api (.-onSnapshot firestore-api))]
    (if-not (and db doc-func col-func on-snapshot)
      (do (js/console.warn "Firestore not available for listen-to-doc")
          (constantly nil))
      (let [doc-ref (doc-func db "users" user-id "documents" "main")
            col-ref (col-func db "users" user-id "documents" "main" "nodes")
            state (atom {:main nil :nodes {} :main-loaded? false :nodes-loaded? false :main-pending? true :nodes-pending? true})
            emit (fn []
                   (let [{:keys [main nodes main-loaded? nodes-loaded? main-pending? nodes-pending?]} @state]
                     (when (and main-loaded? nodes-loaded? (not main-pending?) (not nodes-pending?))
                       (let [old-nodes (get-in main [:data :nodes])
                             root (or (:root main) (get-in main [:data :root]))
                             merged-nodes (merge (fix-stored-nodes old-nodes) nodes)
                             data {:root root :nodes merged-nodes :metadata (:metadata main)}]
                         (when (consistent-doc? data)
                           (callback {:type :updated :data data}))))))
            unsub-main (on-snapshot doc-ref
                                    #(process-main-snapshot % state emit)
                                    #(callback {:type :error :error %}))
            unsub-nodes (on-snapshot col-ref
                                     #(process-nodes-snapshot % state emit)
                                     #(callback {:type :error :error %}))]
        (fn []
          (unsub-main)
          (unsub-nodes))))))