## SQLite database operations for worker coordination
##
## Key patterns:
## - One connection per thread (never share DbConn)
## - WAL mode for concurrent readers
## - BEGIN IMMEDIATE for write transactions
## - Explicit ack for at-least-once delivery

import std/[os, json, options, sysrand, strutils]
import tiny_sqlite
import ./types
import ./utils

# Helper for generating unique IDs
proc genOid*(): string =
  ## Generate a cryptographically random ID (16 lowercase chars)
  var bytes: array[16, byte]
  if not urandom(bytes):
    # Fallback: use timestamp + process ID if urandom fails
    let ts = epochMs()
    let pid = getCurrentProcessId()
    return "f" & toHex(ts, 12) & toHex(pid, 3)
  result = newString(16)
  for i in 0..<16:
    result[i] = chr(ord('a') + (bytes[i] mod 26))

const Schema* = """
-- Messages table (seq is rowid-based for true auto-increment)
CREATE TABLE IF NOT EXISTS messages (
    seq INTEGER PRIMARY KEY AUTOINCREMENT,
    id TEXT NOT NULL UNIQUE,
    ts_ms INTEGER NOT NULL,
    from_agent TEXT NOT NULL,
    to_agent TEXT,
    type TEXT NOT NULL,
    correlation_id TEXT,
    in_reply_to TEXT,
    payload TEXT,
    payload_ref TEXT,
    CHECK (payload IS NULL OR payload_ref IS NULL)
);

CREATE INDEX IF NOT EXISTS idx_messages_to_seq ON messages(to_agent, seq);
CREATE INDEX IF NOT EXISTS idx_messages_correlation_seq ON messages(correlation_id, seq);
CREATE INDEX IF NOT EXISTS idx_messages_type_from_ts ON messages(type, from_agent, ts_ms);

-- Cursors table (per-agent read position)
CREATE TABLE IF NOT EXISTS cursors (
    agent_id TEXT PRIMARY KEY,
    last_acked_seq INTEGER NOT NULL DEFAULT 0,
    updated_at_ms INTEGER NOT NULL
);

-- Heartbeats table (separate from messages to avoid write contention)
CREATE TABLE IF NOT EXISTS heartbeats (
    agent_id TEXT PRIMARY KEY,
    ts_ms INTEGER NOT NULL,
    status TEXT NOT NULL,
    current_task TEXT,
    progress REAL
);

-- Task claims with lease expiry
CREATE TABLE IF NOT EXISTS task_claims (
    task_id TEXT PRIMARY KEY,
    claimed_by TEXT NOT NULL,
    claimed_at_ms INTEGER NOT NULL,
    lease_until_ms INTEGER NOT NULL
);

-- Workers table (state machine)
CREATE TABLE IF NOT EXISTS workers (
    task_id TEXT PRIMARY KEY,
    state TEXT NOT NULL,
    branch TEXT NOT NULL,
    worktree TEXT NOT NULL,
    description TEXT,
    created_at_ms INTEGER NOT NULL,
    state_changed_at_ms INTEGER NOT NULL
);

-- Schema versioning
CREATE TABLE IF NOT EXISTS meta (
    key TEXT PRIMARY KEY,
    value TEXT NOT NULL
);

-- Export tracking
CREATE TABLE IF NOT EXISTS export_state (
    id INTEGER PRIMARY KEY CHECK (id = 1),
    last_seq INTEGER NOT NULL DEFAULT 0
);
"""

const
  CurrentSchemaVersion* = 1

  # Migrations are applied in order. Each migration upgrades FROM the version
  # number to the next version. Add new migrations at the end.
  # Format: (fromVersion, sql)
  Migrations*: seq[tuple[fromVersion: int, sql: string]] = @[
    # Example migration (uncomment when needed):
    # (1, "ALTER TABLE workers ADD COLUMN priority INTEGER DEFAULT 2"),
  ]

proc getSchemaVersion*(db: DbConn): int =
  ## Get current schema version from database
  let row = db.one("SELECT value FROM meta WHERE key = 'schema_version'")
  if row.isNone:
    return 0
  try:
    return parseInt(row.get[0].fromDbValue(string))
  except ValueError:
    return 0

proc runMigrations*(db: DbConn) =
  ## Apply any pending schema migrations
  let currentVersion = db.getSchemaVersion()

  for migration in Migrations:
    if migration.fromVersion >= currentVersion:
      continue  # Skip already-applied migrations

    # Run migration in a transaction
    db.withTransaction:
      db.execScript(migration.sql)
      let newVersion = migration.fromVersion + 1
      db.exec("UPDATE meta SET value = ? WHERE key = 'schema_version'",
              $newVersion)

    logWarn("runMigrations", "Applied migration " & $migration.fromVersion &
            " → " & $(migration.fromVersion + 1))

proc initSchema*(db: DbConn) =
  ## Initialize database schema and run migrations
  db.execScript(Schema)
  # Insert meta if not exists
  db.exec("INSERT OR IGNORE INTO meta (key, value) VALUES ('schema_version', ?)",
          $CurrentSchemaVersion)
  db.exec("INSERT OR IGNORE INTO export_state (id, last_seq) VALUES (1, 0)")

  # Run any pending migrations for existing databases
  runMigrations(db)

proc openBusDb*(dbPath: string = BusDbPath): DbConn =
  ## Open database with required PRAGMAs. One connection per thread.
  try:
    createDir(parentDir(dbPath))
    result = openDatabase(dbPath)
    result.exec("PRAGMA busy_timeout = " & $BusyTimeoutMs)
    result.exec("PRAGMA foreign_keys = ON")
    result.exec("PRAGMA journal_mode = WAL")
    result.exec("PRAGMA synchronous = NORMAL")
    initSchema(result)
  except CatchableError as e:
    raise newException(DbError, "openBusDb(" & dbPath & "): " & e.msg)

proc optStr*(s: string): Option[string] =
  ## Convert empty string to none
  if s == "": none(string) else: some(s)

proc publish*(db: DbConn, agentId: string, msgType: string,
              to: string = "", correlationId: string = "",
              inReplyTo: string = "", payload: JsonNode = nil): string =
  ## Publish a message to the bus. Returns message ID.
  let msgId = genOid()
  let tsMs = epochMs()
  let payloadStr = if payload.isNil: "" else: $payload

  db.exec("""
    INSERT INTO messages (id, ts_ms, from_agent, to_agent, type,
                          correlation_id, in_reply_to, payload)
    VALUES (?, ?, ?, ?, ?, ?, ?, ?)
  """, msgId, tsMs, agentId,
      optStr(to),
      msgType,
      optStr(correlationId),
      optStr(inReplyTo),
      optStr(payloadStr))

  return msgId

type
  Message* = object
    seq*: int64
    id*: string
    tsMs*: int64
    fromAgent*: string
    toAgent*: Option[string]
    msgType*: string
    correlationId*: Option[string]
    inReplyTo*: Option[string]
    payload*: Option[JsonNode]
    payloadRef*: Option[string]

proc poll*(db: DbConn, agentId: string, limit: int = 100): seq[Message] =
  ## Fetch unacknowledged messages. Does NOT advance cursor.
  let cursor = db.one(
    "SELECT last_acked_seq FROM cursors WHERE agent_id = ?",
    agentId
  )
  let lastSeq = if cursor.isSome: cursor.get[0].fromDbValue(int64) else: 0'i64

  for row in db.iterate("""
    SELECT seq, id, ts_ms, from_agent, to_agent, type,
           correlation_id, in_reply_to, payload, payload_ref
    FROM messages
    WHERE seq > ?
      AND (to_agent IS NULL OR to_agent = ?)
    ORDER BY seq
    LIMIT ?
  """, lastSeq, agentId, limit):
    var msg = Message(
      seq: row[0].fromDbValue(int64),
      id: row[1].fromDbValue(string),
      tsMs: row[2].fromDbValue(int64),
      fromAgent: row[3].fromDbValue(string),
      msgType: row[5].fromDbValue(string),
      toAgent: optFromDb(row[4], string),
      correlationId: optFromDb(row[6], string),
      inReplyTo: optFromDb(row[7], string),
      payloadRef: optFromDb(row[9], string),
    )
    # Payload needs special handling for JSON parsing
    let payloadStr = optFromDb(row[8], string)
    if payloadStr.isSome:
      try:
        msg.payload = some(parseJson(payloadStr.get))
      except JsonParsingError as e:
        logWarn("poll", "malformed JSON in message " & msg.id & ": " & e.msg)
    result.add(msg)

proc ack*(db: DbConn, agentId: string, seq: int64) =
  ## Acknowledge successful processing. Advances cursor.
  let tsMs = epochMs()
  db.exec("""
    INSERT INTO cursors (agent_id, last_acked_seq, updated_at_ms)
    VALUES (?, ?, ?)
    ON CONFLICT(agent_id) DO UPDATE SET
        last_acked_seq = MAX(cursors.last_acked_seq, excluded.last_acked_seq),
        updated_at_ms = excluded.updated_at_ms
  """, agentId, seq, tsMs)

proc writeHeartbeat*(db: DbConn, agentId: string, status: HeartbeatStatus,
                     currentTask: string = "", progress: float = 0.0) =
  ## Upsert heartbeat to dedicated table
  let tsMs = epochMs()
  db.exec("""
    INSERT INTO heartbeats (agent_id, ts_ms, status, current_task, progress)
    VALUES (?, ?, ?, ?, ?)
    ON CONFLICT(agent_id) DO UPDATE SET
        ts_ms = excluded.ts_ms,
        status = excluded.status,
        current_task = excluded.current_task,
        progress = excluded.progress
  """, agentId, tsMs, $status,
      optStr(currentTask),
      progress)

proc getHeartbeat*(db: DbConn, agentId: string): Option[tuple[tsMs: int64, status: string]] =
  ## Get last heartbeat for an agent
  let row = db.one(
    "SELECT ts_ms, status FROM heartbeats WHERE agent_id = ?",
    agentId
  )
  if row.isSome:
    return some((tsMs: row.get[0].fromDbValue(int64),
                 status: row.get[1].fromDbValue(string)))

proc tryClaim*(db: DbConn, taskId, agentId: string): bool =
  ## Attempt to claim a task. Returns true if successful.
  let tsMs = epochMs()
  let leaseUntil = tsMs + LeaseDurationMs

  db.exec("BEGIN IMMEDIATE")
  try:
    let existing = db.one(
      "SELECT claimed_by, lease_until_ms FROM task_claims WHERE task_id = ?",
      taskId
    )

    if existing.isSome:
      let leaseUntilMs = existing.get[1].fromDbValue(int64)
      if leaseUntilMs > tsMs:
        db.exec("ROLLBACK")
        return false
      else:
        db.exec("DELETE FROM task_claims WHERE task_id = ?", taskId)

    db.exec("""
      INSERT INTO task_claims (task_id, claimed_by, claimed_at_ms, lease_until_ms)
      VALUES (?, ?, ?, ?)
    """, taskId, agentId, tsMs, leaseUntil)

    db.exec("COMMIT")
    return true
  except CatchableError as e:
    logWarn("tryClaim", "failed for task " & taskId & ": " & e.msg)
    db.exec("ROLLBACK")
    return false

proc renewClaim*(db: DbConn, taskId, agentId: string): bool =
  ## Renew lease on a claimed task
  let tsMs = epochMs()
  let leaseUntil = tsMs + LeaseDurationMs

  db.exec("""
    UPDATE task_claims
    SET lease_until_ms = ?
    WHERE task_id = ? AND claimed_by = ?
  """, leaseUntil, taskId, agentId)

  return db.changes() > 0

proc releaseClaim*(db: DbConn, taskId, agentId: string) =
  ## Release a claim when done
  db.exec("""
    DELETE FROM task_claims
    WHERE task_id = ? AND claimed_by = ?
  """, taskId, agentId)
