import type { Database } from "bun:sqlite";
import { Effect } from "effect";
import {
  makeDatabaseCloser,
  openInitializedDatabase,
  repositoryEffect,
  type RepositoryError,
} from "../../stores/sqlite";

export class PeakMetricsStore {
  private readonly db: Database;
  private readonly closeDatabase: () => Effect.Effect<void, RepositoryError>;

  public constructor(dbPath: string) {
    this.db = openInitializedDatabase(dbPath, (db) => this.migrate(db));
    this.closeDatabase = makeDatabaseCloser(this.db, "peak-metrics.close");
  }

  private migrate(db: Database): void {
    db.run(`
      CREATE TABLE IF NOT EXISTS peak_metrics (
        model_id TEXT PRIMARY KEY,
        prefill_tps REAL,
        generation_tps REAL,
        ttft_ms REAL,
        total_tokens INTEGER DEFAULT 0,
        total_requests INTEGER DEFAULT 0,
        updated_at TEXT DEFAULT CURRENT_TIMESTAMP
      )
    `);
    db.run(`
      CREATE TABLE IF NOT EXISTS peak_metric_sessions (
        session_id TEXT PRIMARY KEY,
        model_id TEXT NOT NULL,
        peak_prefill_tps REAL,
        peak_generation_tps REAL,
        best_ttft_ms REAL,
        started_at TEXT DEFAULT CURRENT_TIMESTAMP,
        updated_at TEXT DEFAULT CURRENT_TIMESTAMP
      )
    `);
    db.run(
      `CREATE INDEX IF NOT EXISTS idx_peak_metric_sessions_model_updated ON peak_metric_sessions(model_id, updated_at)`,
    );
  }

  public get(modelId: string): Record<string, unknown> | null {
    const row = this.db
      .query("SELECT * FROM peak_metrics WHERE model_id = ?")
      .get(modelId) as Record<string, unknown> | null;
    return row ? { ...row } : null;
  }

  public getEffect(
    modelId: string,
  ): Effect.Effect<Record<string, unknown> | null, RepositoryError> {
    return repositoryEffect("peak-metrics.get", () => this.get(modelId));
  }

  public updateIfBetter(
    modelId: string,
    prefillTps?: number,
    generationTps?: number,
    ttftMs?: number,
  ): Record<string, unknown> {
    const current = this.get(modelId);
    const updates: Record<string, number> = {};

    if (current) {
      if (
        prefillTps !== undefined &&
        (current["prefill_tps"] === null || Number(prefillTps) > Number(current["prefill_tps"]))
      ) {
        updates["prefill_tps"] = prefillTps;
      }
      if (
        generationTps !== undefined &&
        (current["generation_tps"] === null ||
          Number(generationTps) > Number(current["generation_tps"]))
      ) {
        updates["generation_tps"] = generationTps;
      }
      if (
        ttftMs !== undefined &&
        (current["ttft_ms"] === null || Number(ttftMs) < Number(current["ttft_ms"]))
      ) {
        updates["ttft_ms"] = ttftMs;
      }
    } else {
      if (prefillTps !== undefined) {
        updates["prefill_tps"] = prefillTps;
      }
      if (generationTps !== undefined) {
        updates["generation_tps"] = generationTps;
      }
      if (ttftMs !== undefined) {
        updates["ttft_ms"] = ttftMs;
      }
    }

    if (Object.keys(updates).length > 0) {
      if (current) {
        const setClause = Object.keys(updates)
          .map((key) => `${key} = ?`)
          .join(", ");
        this.db
          .query(
            `UPDATE peak_metrics SET ${setClause}, updated_at = CURRENT_TIMESTAMP WHERE model_id = ?`,
          )
          .run(...Object.values(updates), modelId);
      } else {
        this.db
          .query(
            `
          INSERT INTO peak_metrics (model_id, prefill_tps, generation_tps, ttft_ms)
          VALUES (?, ?, ?, ?)
        `,
          )
          .run(
            modelId,
            updates["prefill_tps"] ?? null,
            updates["generation_tps"] ?? null,
            updates["ttft_ms"] ?? null,
          );
      }
    }

    return this.get(modelId) ?? {};
  }

  public updateIfBetterEffect(
    modelId: string,
    prefillTps?: number,
    generationTps?: number,
    ttftMs?: number,
  ): Effect.Effect<Record<string, unknown>, RepositoryError> {
    return repositoryEffect("peak-metrics.update-if-better", () =>
      this.updateIfBetter(modelId, prefillTps, generationTps, ttftMs),
    );
  }

  public addTokens(modelId: string, tokens: number, requests = 1): void {
    this.db
      .query(
        `
      INSERT INTO peak_metrics (model_id, total_tokens, total_requests)
      VALUES (?, ?, ?)
      ON CONFLICT(model_id) DO UPDATE SET
        total_tokens = total_tokens + excluded.total_tokens,
        total_requests = total_requests + excluded.total_requests,
        updated_at = CURRENT_TIMESTAMP
    `,
      )
      .run(modelId, tokens, requests);
  }

  public addTokensEffect(
    modelId: string,
    tokens: number,
    requests = 1,
  ): Effect.Effect<void, RepositoryError> {
    return repositoryEffect("peak-metrics.add-tokens", () =>
      this.addTokens(modelId, tokens, requests),
    );
  }

  public updateSessionPeak(
    sessionId: string,
    modelId: string,
    prefillTps?: number,
    generationTps?: number,
    ttftMs?: number,
  ): Record<string, unknown> {
    this.db
      .query(
        `
        INSERT INTO peak_metric_sessions (
          session_id,
          model_id,
          peak_prefill_tps,
          peak_generation_tps,
          best_ttft_ms
        )
        VALUES (?, ?, ?, ?, ?)
        ON CONFLICT(session_id) DO UPDATE SET
          model_id = excluded.model_id,
          peak_prefill_tps = CASE
            WHEN excluded.peak_prefill_tps IS NULL THEN peak_metric_sessions.peak_prefill_tps
            WHEN peak_metric_sessions.peak_prefill_tps IS NULL THEN excluded.peak_prefill_tps
            WHEN excluded.peak_prefill_tps > peak_metric_sessions.peak_prefill_tps THEN excluded.peak_prefill_tps
            ELSE peak_metric_sessions.peak_prefill_tps
          END,
          peak_generation_tps = CASE
            WHEN excluded.peak_generation_tps IS NULL THEN peak_metric_sessions.peak_generation_tps
            WHEN peak_metric_sessions.peak_generation_tps IS NULL THEN excluded.peak_generation_tps
            WHEN excluded.peak_generation_tps > peak_metric_sessions.peak_generation_tps THEN excluded.peak_generation_tps
            ELSE peak_metric_sessions.peak_generation_tps
          END,
          best_ttft_ms = CASE
            WHEN excluded.best_ttft_ms IS NULL THEN peak_metric_sessions.best_ttft_ms
            WHEN peak_metric_sessions.best_ttft_ms IS NULL THEN excluded.best_ttft_ms
            WHEN excluded.best_ttft_ms < peak_metric_sessions.best_ttft_ms THEN excluded.best_ttft_ms
            ELSE peak_metric_sessions.best_ttft_ms
          END,
          updated_at = CURRENT_TIMESTAMP
      `,
      )
      .run(sessionId, modelId, prefillTps ?? null, generationTps ?? null, ttftMs ?? null);

    return this.getSession(sessionId) ?? {};
  }

  public updateSessionPeakEffect(
    sessionId: string,
    modelId: string,
    prefillTps?: number,
    generationTps?: number,
    ttftMs?: number,
  ): Effect.Effect<Record<string, unknown>, RepositoryError> {
    return repositoryEffect("peak-metric-sessions.update", () =>
      this.updateSessionPeak(sessionId, modelId, prefillTps, generationTps, ttftMs),
    );
  }

  public getSession(sessionId: string): Record<string, unknown> | null {
    const row = this.db
      .query("SELECT * FROM peak_metric_sessions WHERE session_id = ?")
      .get(sessionId) as Record<string, unknown> | null;
    return row ? { ...row } : null;
  }

  public getSessionEffect(
    sessionId: string,
  ): Effect.Effect<Record<string, unknown> | null, RepositoryError> {
    return repositoryEffect("peak-metric-sessions.get", () => this.getSession(sessionId));
  }

  public getBestSession(modelId: string): Record<string, unknown> | null {
    const row = this.db
      .query(
        `
        SELECT * FROM peak_metric_sessions
        WHERE model_id = ?
        ORDER BY
          COALESCE(peak_generation_tps, 0) DESC,
          COALESCE(peak_prefill_tps, 0) DESC,
          updated_at DESC
        LIMIT 1
      `,
      )
      .get(modelId) as Record<string, unknown> | null;
    return row ? { ...row } : null;
  }

  public getBestSessionEffect(
    modelId: string,
  ): Effect.Effect<Record<string, unknown> | null, RepositoryError> {
    return repositoryEffect("peak-metric-sessions.get-best", () => this.getBestSession(modelId));
  }

  public getAll(): Array<Record<string, unknown>> {
    const rows = this.db.query("SELECT * FROM peak_metrics ORDER BY model_id").all() as Array<
      Record<string, unknown>
    >;
    return rows.map((row) => {
      const modelId = String(row["model_id"] ?? "");
      const bestSession = modelId ? this.getBestSession(modelId) : null;
      return {
        ...row,
        best_session_id: bestSession?.["session_id"] ?? null,
        best_session_prefill_tps: bestSession?.["peak_prefill_tps"] ?? null,
        best_session_generation_tps: bestSession?.["peak_generation_tps"] ?? null,
        best_session_ttft_ms: bestSession?.["best_ttft_ms"] ?? null,
      };
    });
  }

  public getAllEffect(): Effect.Effect<Array<Record<string, unknown>>, RepositoryError> {
    return repositoryEffect("peak-metrics.get-all", () => this.getAll());
  }

  public close(): Effect.Effect<void, RepositoryError> {
    return this.closeDatabase();
  }
}

export class LifetimeMetricsStore {
  private readonly db: Database;
  private readonly closeDatabase: () => Effect.Effect<void, RepositoryError>;

  public constructor(dbPath: string) {
    this.db = openInitializedDatabase(dbPath, (db) => this.migrate(db));
    this.closeDatabase = makeDatabaseCloser(this.db, "lifetime-metrics.close");
  }

  private migrate(db: Database): void {
    db.run(`
      CREATE TABLE IF NOT EXISTS lifetime_metrics (
        key TEXT PRIMARY KEY,
        value REAL NOT NULL DEFAULT 0,
        updated_at TEXT DEFAULT CURRENT_TIMESTAMP
      )
    `);
    const defaults: Array<[string, number]> = [
      ["tokens_total", 0],
      ["prompt_tokens_total", 0],
      ["completion_tokens_total", 0],
      ["energy_wh", 0],
      ["uptime_seconds", 0],
      ["requests_total", 0],
      ["first_started_at", 0],
    ];
    for (const [key, value] of defaults) {
      db
        .query("INSERT OR IGNORE INTO lifetime_metrics (key, value) VALUES (?, ?)")
        .run(key, value);
    }
  }

  public get(key: string): number {
    const row = this.db.query("SELECT value FROM lifetime_metrics WHERE key = ?").get(key) as {
      value?: number;
    } | null;
    return row?.value ?? 0;
  }

  public getEffect(key: string): Effect.Effect<number, RepositoryError> {
    return repositoryEffect("lifetime-metrics.get", () => this.get(key));
  }

  public getAll(): Record<string, number> {
    const rows = this.db.query("SELECT key, value FROM lifetime_metrics").all() as Array<{
      key: string;
      value: number;
    }>;
    return Object.fromEntries(rows.map((row) => [row.key, row.value]));
  }

  public getAllEffect(): Effect.Effect<Record<string, number>, RepositoryError> {
    return repositoryEffect("lifetime-metrics.get-all", () => this.getAll());
  }

  public set(key: string, value: number): void {
    this.db
      .query(
        `INSERT INTO lifetime_metrics (key, value, updated_at)
       VALUES (?, ?, CURRENT_TIMESTAMP)
       ON CONFLICT(key) DO UPDATE SET value = excluded.value, updated_at = CURRENT_TIMESTAMP`,
      )
      .run(key, value);
  }

  public setEffect(key: string, value: number): Effect.Effect<void, RepositoryError> {
    return repositoryEffect("lifetime-metrics.set", () => this.set(key, value));
  }

  public increment(key: string, delta: number): number {
    this.db
      .query(
        `INSERT INTO lifetime_metrics (key, value, updated_at)
       VALUES (?, ?, CURRENT_TIMESTAMP)
       ON CONFLICT(key) DO UPDATE SET value = value + excluded.value, updated_at = CURRENT_TIMESTAMP`,
      )
      .run(key, delta);
    return this.get(key);
  }

  public incrementEffect(key: string, delta: number): Effect.Effect<number, RepositoryError> {
    return repositoryEffect("lifetime-metrics.increment", () => this.increment(key, delta));
  }

  public ensureFirstStarted(): void {
    const current = this.get("first_started_at");
    if (current === 0) {
      this.set("first_started_at", Date.now() / 1000);
    }
  }

  public ensureFirstStartedEffect(): Effect.Effect<void, RepositoryError> {
    return repositoryEffect("lifetime-metrics.ensure-first-started", () =>
      this.ensureFirstStarted(),
    );
  }

  public addEnergy(wattHours: number): Effect.Effect<void, RepositoryError> {
    return this.incrementEffect("energy_wh", wattHours).pipe(Effect.asVoid);
  }

  public addTokens(tokens: number): Effect.Effect<void, RepositoryError> {
    return this.incrementEffect("tokens_total", tokens).pipe(Effect.asVoid);
  }

  public addPromptTokens(tokens: number): Effect.Effect<void, RepositoryError> {
    return this.incrementEffect("prompt_tokens_total", tokens).pipe(Effect.asVoid);
  }

  public addCompletionTokens(tokens: number): Effect.Effect<void, RepositoryError> {
    return this.incrementEffect("completion_tokens_total", tokens).pipe(Effect.asVoid);
  }

  public addUptime(seconds: number): Effect.Effect<void, RepositoryError> {
    return this.incrementEffect("uptime_seconds", seconds).pipe(Effect.asVoid);
  }

  public addRequests(count = 1): Effect.Effect<void, RepositoryError> {
    return this.incrementEffect("requests_total", count).pipe(Effect.asVoid);
  }

  public close(): Effect.Effect<void, RepositoryError> {
    return this.closeDatabase();
  }
}