import SwiftUI
import FoundationModels
import Observation

@Observable class ApplicationData {
   var response: AttributedString = ""
   var prompt: String = ""
   @ObservationIgnored var model: SystemLanguageModel
   @ObservationIgnored var session: LanguageModelSession

   static let shared: ApplicationData = ApplicationData()
   private init() {
      model = SystemLanguageModel.default
      session = LanguageModelSession(model: model)
   }
   func sendPrompt() async {
      do {
         let answer = try await session.respond(to: prompt)
         var newResponse = AttributedString("\(answer.content)\n\n")
         newResponse.font = .system(size: 16, weight: .regular)
         response.append(newResponse)
         
         calculateTokens()
      } catch {
         response = AttributedString("Error accessing the model: \(error)")
      }
      prompt = ""
   }
   func calculateTokens() {
      let transcript = session.transcript
      var count: Int = 0

      for entry in transcript {
         switch entry {
         case .prompt(let prompt):
            if let text = prompt.segments.first {
               count += text.description.count / 4
            }
         case .response(let response):
            count += response.description.count / 4
         default:
            break
         }
      }
      print("Tokens: \(count)")
   }
}