import std/strutils
import results
import db_connector/db_sqlite
import db_models
import nimcrypto

type SegmentationPersistence* = object
  db*: DbConn

proc removeMessageSegmentsOlderThan*(
    self: SegmentationPersistence, timestamp: int64
): Result[void, string] =
  try:
    self.db.exec(sql"DELETE FROM message_segments WHERE timestamp < ?", timestamp)
    ok()
  except DbError as e:
    err("remove message segments with error: " & e.msg)

proc removeMessageSegmentsCompletedOlderThan*(
    self: SegmentationPersistence, timestamp: int64
): Result[void, string] =
  try:
    self.db.exec(
      sql"DELETE FROM message_segments_completed WHERE timestamp < ?", timestamp
    )
    ok()
  except DbError as e:
    err("remove message segments completed with error: " & e.msg)

proc isMessageAlreadyCompleted*(
    self: SegmentationPersistence, hash: seq[byte]
): Result[bool, string] =
  try:
    let row = self.db.getRow(
      sql"SELECT COUNT(*) FROM message_segments_completed WHERE hash = ?", hash
    )
    if row.len == 0:
      return ok(false)
    let count = row[0].parseInt
    ok(count > 0)
  except CatchableError as e:
    err("check message already completed with error: " & e.msg)

proc saveMessageSegment*(
    self: SegmentationPersistence,
    segment: SegmentMessage,
    sigPubKeyBlob: seq[byte],
    timestamp: int64,
): Result[void, string] =
  try:
    self.db.exec(
      sql"""
      INSERT INTO message_segments (
        hash, segment_index, segments_count, parity_segment_index, 
        parity_segments_count, sig_pub_key, payload, timestamp
      ) VALUES (?, ?, ?, ?, ?, ?, ?, ?)""",
      segment.entireMessageHash.toHex,
      segment.index,
      segment.segmentsCount,
      segment.paritySegmentIndex,
      segment.paritySegmentsCount,
      sigPubKeyBlob,
      segment.payload.toHex,
      timestamp,
    )

    let storedPayload = self.db.getValue(
      sql"SELECT payload FROM message_segments WHERE hash = ? AND sig_pub_key = ? AND segment_index = ?",
      segment.entireMessageHash.toHex, sigPubKeyBlob, segment.index
    )
    let payloadBytes = nimcrypto.fromHex(storedPayload)
    
    ok()
  except DbError as e:
    err("save message segments with error: " & e.msg)

proc getMessageSegments*(self: SegmentationPersistence, hash: seq[byte], sigPubKeyBlob: seq[byte]): Result[seq[SegmentMessage], string] =
  var segments = newSeq[SegmentMessage]()
  let query = sql"""
    SELECT
      hash, segment_index, segments_count, parity_segment_index, parity_segments_count, payload
    FROM
      message_segments
    WHERE
      hash = ? AND sig_pub_key = ?
    ORDER BY
      (segments_count = 0) ASC,
      segment_index ASC,
      parity_segment_index ASC
  """

  try:
    for row in self.db.rows(query, hash.toHex, sigPubKeyBlob):
      let segment = SegmentMessage(
        entireMessageHash: nimcrypto.fromHex(row[0]),
        index: uint32(parseInt(row[1])),
        segmentsCount: uint32(parseInt(row[2])),
        paritySegmentIndex: uint32(parseInt(row[3])),
        paritySegmentsCount: uint32(parseInt(row[4])),
        payload: nimcrypto.fromHex(row[5])
      )
      segments.add(segment)
    return ok(segments)
  except CatchableError as e:
    return err("get Message Segments with error: " & e.msg)

proc completeMessageSegments*(
  self: SegmentationPersistence,
  hash: seq[byte],
  sigPubKeyBlob: seq[byte],
  timestamp: int64
): Result[void, string] =
  try:
    self.db.exec(sql"BEGIN")

    # Delete old message segments
    self.db.exec(sql"DELETE FROM message_segments WHERE hash = ? AND sig_pub_key = ?", hash, sigPubKeyBlob)

    # Insert completed marker
    self.db.exec(sql"INSERT INTO message_segments_completed (hash, sig_pub_key, timestamp) VALUES (?, ?, ?)",
      hash, sigPubKeyBlob, timestamp)

    self.db.exec(sql"COMMIT")
    return ok()
  except DbError as e:
    try:
      self.db.exec(sql"ROLLBACK")
    except: discard
    return err("complete segment messages with error: " & e.msg)