web-app/src/stores/nostr.ts

import { defineStore } from 'pinia'
import { ref, computed, watch } from 'vue'
import type { NostrEvent, NostrProfile, NostrAccount, DirectMessage } from '../types/nostr'
import { isValidPrivateKey, formatPrivateKey } from '@/lib/nostr'

declare global {
  interface Window {
    NostrTools: {
      getPublicKey: (privkey: string) => string
      generatePrivateKey: () => string
      nip04: {
        encrypt: (privkey: string, pubkey: string, content: string) => Promise<string>
        decrypt: (privkey: string, pubkey: string, content: string) => Promise<string>
      }
      getEventHash: (event: NostrEvent) => string
      signEvent: (event: NostrEvent, privkey: string) => Promise<string>
      getSignature: (event: NostrEvent, privkey: string) => string
      verifySignature: (event: NostrEvent) => boolean
      nip19: {
        decode: (str: string) => { type: string; data: string }
        npubEncode: (hex: string) => string
      }
      relayInit: (url: string) => {
        connect: () => Promise<void>
        sub: (filters: any[]) => {
          on: (event: string, callback: (event: NostrEvent) => void) => void
        }
        publish: (event: NostrEvent) => {
          on: (type: 'ok' | 'failed', cb: (msg?: string) => void) => void
        }
        close: () => void
      }
    }
  }
}

const DEFAULT_RELAYS = [
  'wss://nostr.atitlan.io'
]

// Helper functions
async function withTimeout<T>(promise: Promise<T>, timeoutMs: number = 10000): Promise<T> {
  return Promise.race([
    promise,
    new Promise<T>((_, reject) => 
      setTimeout(() => reject(new Error('Operation timed out')), timeoutMs)
    )
  ])
}

// Add to state
const connectionStatus = ref<'connected' | 'connecting' | 'disconnected'>('disconnected')
const hasUnreadMessages = ref(false)
const viewedMessageIds = ref<Set<string>>(new Set(
  JSON.parse(localStorage.getItem('nostr_viewed_messages') || '[]')
))

// Update in connect function
async function connectToRelay(url: string) {
  connectionStatus.value = 'connecting'
  console.log(`Attempting to connect to relay: ${url}`)
  const relay = window.NostrTools.relayInit(url)
  try {
    console.log(`Initializing connection to ${url}...`)
    await withTimeout(relay.connect())
    console.log(`Successfully connected to ${url}`)
    connectionStatus.value = 'connected'
    return relay
  } catch (err) {
    console.error(`Failed to connect to ${url}:`, err)
    connectionStatus.value = 'disconnected'
    return null
  }
}

async function publishEvent(event: NostrEvent, relays: { url: string }[]) {
  const promises = relays.map(async ({ url }) => {
    const relay = window.NostrTools.relayInit(url)
    try {
      await relay.connect()
      const pub = relay.publish(event)
      return new Promise((resolve, reject) => {
        pub.on('ok', () => resolve(true))
        pub.on('failed', reject)
      })
    } catch (err) {
      console.error(`Failed to publish to ${url}:`, err)
      return false
    }
  })

  await Promise.all(promises)
}

export const useNostrStore = defineStore('nostr', () => {
  // State
  const account = ref<NostrAccount | null>(JSON.parse(localStorage.getItem('nostr_account') || 'null'))
  const profiles = ref<Map<string, NostrProfile>>(new Map())
  const messages = ref<Map<string, DirectMessage[]>>(new Map())
  const activeChat = ref<string | null>(null)
  const relayPool = ref<any[]>([])
  const processedMessageIds = ref(new Set<string>())
  const currentSubscription = ref<any | null>(null)
  const hasActiveSubscription = ref(false)

  // Load stored messages and IDs on initialization
  const initializeFromStorage = () => {
    try {
      const messageMap = new Map<string, DirectMessage[]>(
        JSON.parse(localStorage.getItem('nostr_messages') || '[]')
      )
      
      messageMap.forEach((msgs: DirectMessage[]) => {
        msgs.forEach(msg => {
          processedMessageIds.value.add(msg.id)
        })
      })
      
      messages.value = messageMap
    } catch (err) {
      console.error('Failed to load stored messages:', err)
      localStorage.removeItem('nostr_messages')
    }
  }

  // Call initialization
  initializeFromStorage()

  // Watch account changes and persist to localStorage
  watch(account, (newAccount) => {
    if (newAccount) {
      localStorage.setItem('nostr_account', JSON.stringify(newAccount))
    } else {
      localStorage.removeItem('nostr_account')
    }
  }, { deep: true })

  // Watch messages for changes and persist
  watch(messages, (newMessages) => {
    try {
      localStorage.setItem('nostr_messages', 
        JSON.stringify(Array.from(newMessages.entries()))
      )
    } catch (err) {
      console.error('Failed to save messages:', err)
    }
  }, { deep: true })

  // Initialize store if account exists in localStorage
  if (account.value) {
    console.log('Found existing account, initializing connection...')
    init()
  }

  // Computed
  const isLoggedIn = computed(() => !!account.value)
  const currentMessages = computed(() => 
    activeChat.value ? messages.value.get(activeChat.value) || [] : []
  )

  // Initialize connection if account exists
  async function init() {
    if (!account.value) return

    try {
      // Only clear profiles and processed IDs
      profiles.value.clear()
      processedMessageIds.value.clear()
      
      // Connect to relays
      relayPool.value = (await Promise.all(
        account.value.relays.map(async relay => {
          console.log('Connecting to relay:', relay.url)
          const connection = await connectToRelay(relay.url)
          if (!connection) {
            console.error('Failed to connect to relay:', relay.url)
          }
          return connection
        })
      )).filter((relay): relay is any => relay !== null)

      if (relayPool.value.length === 0) {
        throw new Error('Failed to connect to any relays')
      }

      // Setup visibility change handler
      setupVisibilityHandler()

      // Subscribe to messages in the background
      subscribeToMessages().catch(err => {
        console.error('Background subscription failed:', err)
      })

    } catch (err) {
      console.error('Failed to initialize:', err)
      throw err
    }
  }

  // Actions
  async function login(privkey: string) {
    if (!isValidPrivateKey(privkey)) {
      throw new Error('Invalid private key')
    }

    const formattedKey = formatPrivateKey(privkey)
    const pubkey = window.NostrTools.getPublicKey(formattedKey)
    
    account.value = {
      pubkey,
      privkey: formattedKey,
      relays: DEFAULT_RELAYS.map(url => ({ url, read: true, write: true }))
    }

    // Initialize connection in the background
    init().catch(err => {
      console.error('Background initialization failed:', err)
    })
  }

  async function logout() {
    account.value = null
    relayPool.value.forEach(relay => relay.close())
    relayPool.value = []
    messages.value.clear()
    profiles.value.clear()
    processedMessageIds.value.clear()
    activeChat.value = null
    localStorage.removeItem('nostr_messages')
    localStorage.removeItem('nostr_account')
    hasUnreadMessages.value = false
    localStorage.removeItem('nostr_unread_messages')
    viewedMessageIds.value.clear()
    localStorage.removeItem('nostr_viewed_messages')
  }

  const addMessage = async (pubkey: string, message: DirectMessage) => {
    // Skip if we've already processed this message
    if (processedMessageIds.value.has(message.id)) {
      return
    }

    processedMessageIds.value.add(message.id)
    const userMessages = messages.value.get(pubkey) || []
    
    // Check for duplicates
    const isDuplicate = userMessages.some(msg => 
      msg.content === message.content && 
      Math.abs(msg.created_at - message.created_at) < 1
    )

    if (!isDuplicate) {
      messages.value.set(pubkey, [...userMessages, message].sort((a, b) => 
        a.created_at - b.created_at
      ))

      // Only set unread if:
      // 1. Message came from websocket (not storage)
      // 2. Not from current chat
      // 3. Not sent by us
      if (!message.fromStorage && pubkey !== activeChat.value && !message.sent) {
        console.log('New unread message received:', { pubkey, messageId: message.id })
        hasUnreadMessages.value = true
      }
    }
  }

  async function sendMessage(to: string, content: string) {
    if (!account.value) return

    const encrypted = await window.NostrTools.nip04.encrypt(account.value.privkey, to, content)
    const event: NostrEvent = {
      kind: 4,
      pubkey: account.value.pubkey,
      created_at: Math.floor(Date.now() / 1000),
      tags: [['p', to]],
      content: encrypted,
      id: '',
      sig: ''
    }

    event.id = window.NostrTools.getEventHash(event)
    event.sig = await window.NostrTools.signEvent(event, account.value.privkey)

    // Add to local messages first
    const dm: DirectMessage = {
      id: event.id,
      pubkey: to,
      content,
      created_at: event.created_at,
      sent: true
    }

    await addMessage(to, dm)

    // Then publish to relays
    await publishEvent(event, account.value.relays)
  }

  async function subscribeToMessages() {
    if (!account.value || hasActiveSubscription.value) return
    
    hasActiveSubscription.value = true
    // Cleanup existing subscription
    unsubscribeFromMessages()

    // Get timestamp from 24 hours ago
    const since = Math.floor(Date.now() / 1000) - (24 * 60 * 60)
    let hasReceivedMessages = false

    try {
      const subscribeToRelay = (relay: any) => {
        return new Promise((resolve) => {
          const subs: any[] = []

          try {
            console.log('Setting up subscriptions for relay...')

            // Subscribe to received messages
            const receivedSub = relay.sub([{
              kinds: [4],
              '#p': [account.value!.pubkey],
              since,
              limit: 100 // Add limit to ensure we get historical messages
            }])
            subs.push(receivedSub)

            // Subscribe to sent messages
            const sentSub = relay.sub([{
              kinds: [4],
              authors: [account.value!.pubkey],
              since,
              limit: 100 // Add limit to ensure we get historical messages
            }])
            subs.push(sentSub)

            // Handle received messages
            receivedSub.on('event', async (event: NostrEvent) => {
              hasReceivedMessages = true
              try {
                if (processedMessageIds.value.has(event.id)) return
                
                const decrypted = await window.NostrTools.nip04.decrypt(
                  account.value!.privkey,
                  event.pubkey,
                  event.content
                )

                const dm: DirectMessage = {
                  id: event.id,
                  pubkey: event.pubkey,
                  content: decrypted,
                  created_at: event.created_at,
                  sent: false,
                  fromStorage: false // Mark as not from storage
                }

                await addMessage(event.pubkey, dm)
              } catch (err) {
                console.error('Failed to decrypt received message:', err)
              }
            })

            // Handle sent messages
            sentSub.on('event', async (event: NostrEvent) => {
              hasReceivedMessages = true
              try {
                if (processedMessageIds.value.has(event.id)) return

                const targetPubkey = event.tags.find(tag => tag[0] === 'p')?.[1]
                if (!targetPubkey) return

                const decrypted = await window.NostrTools.nip04.decrypt(
                  account.value!.privkey,
                  targetPubkey,
                  event.content
                )

                const dm: DirectMessage = {
                  id: event.id,
                  pubkey: targetPubkey,
                  content: decrypted,
                  created_at: event.created_at,
                  sent: true
                }

                await addMessage(targetPubkey, dm)
              } catch (err) {
                console.error('Failed to decrypt sent message:', err)
              }
            })

            // Handle EOSE (End of Stored Events)
            receivedSub.on('eose', () => {
              console.log('Received EOSE for received messages')
              if (!hasReceivedMessages) {
                console.log('No messages received yet, keeping subscription open')
              }
            })

            sentSub.on('eose', () => {
              console.log('Received EOSE for sent messages')
              if (!hasReceivedMessages) {
                console.log('No messages received yet, keeping subscription open')
              }
            })

            // Store subscriptions for cleanup
            currentSubscription.value = {
              unsub: () => {
                subs.forEach(sub => {
                  try {
                    if (sub && typeof sub.unsub === 'function') {
                      sub.unsub()
                    }
                  } catch (err) {
                    console.debug('Failed to unsubscribe:', err)
                  }
                })
              }
            }

            // Keep subscription open
            resolve(true)

          } catch (err) {
            console.debug('Error in subscription setup:', err)
            resolve(false)
          }
        })
      }

      // Wait for all relays to set up subscriptions
      const results = await Promise.all(
        relayPool.value.map(relay => subscribeToRelay(relay))
      )

      // Consider success if at least one relay worked
      return results.some(result => result)
    } catch (err) {
      console.debug('Subscription process failed:', err)
      return false
    }
  }

  function unsubscribeFromMessages() {
    if (currentSubscription.value && typeof currentSubscription.value.unsub === 'function') {
      try {
        currentSubscription.value.unsub()
      } catch (err) {
        console.error('Failed to unsubscribe:', err)
      }
      currentSubscription.value = null
    }
    hasActiveSubscription.value = false
  }

  async function loadProfiles() {
    if (!account.value) return

    const pubkeysToLoad = new Set<string>()
    
    // Collect all unique pubkeys from messages
    for (const [pubkey] of messages.value.entries()) {
      if (!profiles.value.has(pubkey)) {
        pubkeysToLoad.add(pubkey)
      }
    }

    if (pubkeysToLoad.size === 0) return

    try {
      const filter = {
        kinds: [0],
        authors: Array.from(pubkeysToLoad)
      }

      const loadFromRelay = (relay: any) => {
        return new Promise<void>((resolve) => {
          const sub = relay.sub([filter])
          
          sub.on('event', (event: NostrEvent) => {
            try {
              const profile = JSON.parse(event.content)
              profiles.value.set(event.pubkey, {
                pubkey: event.pubkey,
                name: profile.name,
                picture: profile.picture,
                about: profile.about,
                nip05: profile.nip05
              })
            } catch (err) {
              console.error('Failed to parse profile:', err)
            }
          })

          // Resolve after receiving EOSE (End of Stored Events)
          sub.on('eose', () => {
            resolve()
          })

          // Set a timeout in case EOSE is not received
          setTimeout(() => {
            resolve()
          }, 5000)
        })
      }

      // Load profiles from all relays concurrently
      await Promise.all(relayPool.value.map(relay => loadFromRelay(relay)))

    } catch (err) {
      console.error('Failed to load profiles:', err)
    }
  }

  // Add a reconnection function
  async function reconnectToRelays() {
    if (!account.value) return

    console.log('Attempting to reconnect to relays...')
    
    // Close existing connections
    relayPool.value.forEach(relay => {
      try {
        relay.close()
      } catch (err) {
        console.error('Error closing relay:', err)
      }
    })
    relayPool.value = []

    // Reconnect
    relayPool.value = (await Promise.all(
      account.value.relays.map(async relay => {
        console.log('Reconnecting to relay:', relay.url)
        const connection = await connectToRelay(relay.url)
        if (!connection) {
          console.error('Failed to reconnect to relay:', relay.url)
        }
        return connection
      })
    )).filter((relay): relay is any => relay !== null)

    if (relayPool.value.length === 0) {
      throw new Error('Failed to connect to any relays')
    }

    // Resubscribe to messages
    await subscribeToMessages()
  }

  // Update visibility handler
  function setupVisibilityHandler() {
    const handleVisibilityChange = async () => {
      if (document.visibilityState === 'visible' && account.value) {
        console.log('Page became visible, checking connection...')
        try {
          // Only reconnect if we don't have active connections
          if (relayPool.value.length === 0 || !hasActiveSubscription.value) {
            await reconnectToRelays()
          }
        } catch (err) {
          console.error('Failed to reconnect:', err)
        }
      }
    }

    // Remove any existing handler
    document.removeEventListener('visibilitychange', handleVisibilityChange)
    document.addEventListener('visibilitychange', handleVisibilityChange)

    // Add focus handler for mobile
    window.addEventListener('focus', handleVisibilityChange)
  }

  // Add function to clear unread state
  function clearUnreadMessages() {
    hasUnreadMessages.value = false
    localStorage.setItem('nostr_unread_messages', 'false')
    
    // Mark all current chat messages as viewed
    if (activeChat.value) {
      const chatMessages = messages.value.get(activeChat.value) || []
      chatMessages.forEach(msg => {
        viewedMessageIds.value.add(msg.id)
      })
      // Persist viewed message IDs
      localStorage.setItem('nostr_viewed_messages', 
        JSON.stringify(Array.from(viewedMessageIds.value))
      )
    }
  }

  // Add to watch section
  watch(activeChat, () => {
    // Clear unread messages when changing to a chat
    clearUnreadMessages()
  })

  return {
    account,
    profiles,
    messages,
    activeChat,
    isLoggedIn,
    currentMessages,
    init,
    login,
    logout,
    sendMessage,
    subscribeToMessages,
    unsubscribeFromMessages,
    loadProfiles,
    connectionStatus,
    hasActiveSubscription,
    hasUnreadMessages,
    clearUnreadMessages,
    reconnectToRelays,
  }
})