BrainBox/src/brain/index.ts

import { randomUUID } from "node:crypto";
import { config } from "@/config";
import { IdentityDB, type Space } from "identitydb";
import { llm } from "@/openrouter";
import { loadPrompt } from "@/openrouter/promptLoader";
import {
  availabilitySchema,
  dailyScheduleSchema,
  monthlyScheduleSchema,
} from "@/openrouter/schema";
import { logger } from "@/utils/logger";
import { factExtractor } from "./factExtractor";
import { brainManager, type BrainItem } from "./manager";
import {
  type Availability,
  type DailySchedule,
  type MonthlySchedule,
  formatDateKey,
  formatMonthKey,
  nextDay,
  nextMonth,
  pad2,
} from "./schedule";
import { BadRequestResponseError } from "@openrouter/sdk/models/errors";

export interface DebugOptions {
  personality: string;
}

export class Brain {
  private availabilityCache: Map<string, Availability[]> = new Map();

  constructor(
    public db: IdentityDB,
    public space: Space,
    public brainbase: BrainItem,
    public debug: boolean = false,
  ) {}

  async createDailySchedule(
    datetime: Date,
    message: string,
  ): Promise<DailySchedule | null> {
    try {
      const target = nextDay(datetime);
      const dateKey = formatDateKey(target);
      const topicName = `daily-schedule:${dateKey}`;

      const monthlySummary = await this.getMonthlySummaryForDay(target);
      const history = await this.getHistoryFacts();

      const instruction = await loadPrompt("DAILY_SCHEDULE");
      const promptMessage = [
        `Target date: ${dateKey} (${target.toLocaleDateString("en-US", { weekday: "long" })})`,
        `Personality: ${this.brainbase.baseSystemPrompt}`,
        monthlySummary
          ? `Monthly summary for this day: ${monthlySummary}`
          : "(no monthly summary available for this date)",
        `Recent history (facts):`,
        history,
        `User direction: ${message}`,
      ].join("\n\n");

      const schedule = await llm.call<DailySchedule>(llm.models.identity, {
        instruction,
        message: promptMessage,
        jsonSchemaName: "daily-schedule",
        jsonSchema: dailyScheduleSchema,
      });

      if (!this.debug) {
        await this.db.addFact({
          spaceName: this.space.name,
          statement: JSON.stringify(schedule),
          summary: `Daily schedule for ${dateKey} (${schedule.length} slots)`,
          source: "createDailySchedule",
          confidence: 1.0,
          topics: [
            {
              name: topicName,
              category: "temporal",
              granularity: "concrete",
              role: "schedule",
            },
            {
              name: "daily-schedule",
              category: "concept",
              granularity: "abstract",
              role: "schedule",
            },
            {
              name: dateKey,
              category: "temporal",
              granularity: "concrete",
              role: "date",
            },
          ],
        });
      }

      return schedule;
    } catch (error) {
      let reason =
        error instanceof Error
          ? error.message + `(${error.name})`
          : String(error);
      if (error instanceof BadRequestResponseError)
        reason = reason + `${error.body}`;
      logger.error(`createDailySchedule failed: ${reason}`);
      return null;
    }
  }

  async createMonthlySchedule(
    datetime: Date,
    message: string,
  ): Promise<MonthlySchedule | null> {
    try {
      const next = nextMonth(datetime);
      const monthKey = `${next.year}-${pad2(next.month + 1)}`;
      const topicName = `monthly-schedule:${monthKey}`;

      const history = await this.getHistoryFacts();

      const instruction = await loadPrompt("MONTHLY_SCHEDULE");
      const promptMessage = [
        `Target month: ${monthKey} (${next.daysInMonth} days)`,
        `Personality: ${this.brainbase.baseSystemPrompt}`,
        `Recent history (facts):`,
        history,
        `User direction: ${message}`,
      ].join("\n\n");

      const schedule = await llm.call<MonthlySchedule>(llm.models.identity, {
        instruction,
        message: promptMessage,
        jsonSchemaName: "monthly-schedule",
        jsonSchema: monthlyScheduleSchema,
      });

      if (!this.debug) {
        await this.db.addFact({
          spaceName: this.space.name,
          statement: JSON.stringify(schedule),
          summary: `Monthly schedule for ${monthKey} (${schedule.length} days)`,
          source: "createMonthlySchedule",
          confidence: 1.0,
          topics: [
            {
              name: topicName,
              category: "temporal",
              granularity: "concrete",
              role: "schedule",
            },
            {
              name: "monthly-schedule",
              category: "concept",
              granularity: "abstract",
              role: "schedule",
            },
            {
              name: monthKey,
              category: "temporal",
              granularity: "concrete",
              role: "period",
            },
          ],
        });
      }

      return schedule;
    } catch (error) {
      const reason = error instanceof Error ? error.message : String(error);
      logger.error(`createMonthlySchedule failed: ${reason}`);
      return null;
    }
  }

  async getTodayScheduledAvailability(
    datetime: Date,
  ): Promise<Availability[] | null> {
    try {
      const dateKey = formatDateKey(datetime);
      const cached = this.availabilityCache.get(dateKey);
      if (cached) return cached;

      if (this.debug) {
        logger.warn(
          "getTodayScheduledAvailability requires a persisted daily schedule; debug brains have no DB. Use deriveAvailabilityFromSchedule(schedule) instead.",
        );
        return null;
      }

      const topicName = `daily-schedule:${dateKey}`;
      const facts = await this.db.getTopicFacts(topicName, {
        spaceName: this.space.name,
      });
      if (facts.length === 0) return null;

      const dailySchedule = JSON.parse(facts[0]!.statement) as DailySchedule;
      const availability =
        await this.deriveAvailabilityFromSchedule(dailySchedule);

      this.availabilityCache.set(dateKey, availability);
      return availability;
    } catch (error) {
      const reason = error instanceof Error ? error.message : String(error);
      logger.error(`getTodayScheduledAvailability failed: ${reason}`);
      return null;
    }
  }

  async deriveAvailabilityFromSchedule(
    schedule: DailySchedule,
  ): Promise<Availability[]> {
    try {
      const instruction = await loadPrompt("SCHEDULE_AVAILABILITY");
      const promptMessage = JSON.stringify({
        schedule,
        personality: this.brainbase.baseSystemPrompt,
      });

      return await llm.call<Availability[]>(llm.models.identity, {
        instruction,
        message: promptMessage,
        jsonSchemaName: "availability",
        jsonSchema: availabilitySchema,
      });
    } catch (error) {
      const reason = error instanceof Error ? error.message : String(error);
      logger.error(`deriveAvailabilityFromSchedule failed: ${reason}`);
      throw error;
    }
  }

  removeScheduledAvailability(): void {
    this.availabilityCache.clear();
  }

  private async getMonthlySummaryForDay(target: Date): Promise<string | null> {
    if (this.debug) return null;
    try {
      const monthKey = formatMonthKey(target);
      const topicName = `monthly-schedule:${monthKey}`;
      const facts = await this.db.getTopicFacts(topicName, {
        spaceName: this.space.name,
      });
      if (facts.length === 0) return null;

      const monthly = JSON.parse(facts[0]!.statement) as MonthlySchedule;
      const day = target.getDate();
      const entry = monthly.find((d) => d.day === day);
      return entry?.summary ?? null;
    } catch {
      return null;
    }
  }

  private async getHistoryFacts(): Promise<string> {
    if (this.debug) return "";
    try {
      const topics = await this.db.listTopics({
        spaceName: this.space.name,
        includeFacts: true,
      });
      const statements: string[] = [];
      for (const topic of topics) {
        const t = topic as { facts?: Array<{ statement: string }> };
        if (t.facts) {
          for (const f of t.facts) statements.push(f.statement);
        }
      }
      return statements.slice(-30).join("\n");
    } catch {
      return "";
    }
  }

  static async create(
    displayName: string,
    seed: string,
  ): Promise<Brain | null> {
    try {
      const personaInitInstruction = await loadPrompt("PERSONA_INIT");
      const description = await llm.call<string>(llm.models.identity, {
        instruction: personaInitInstruction,
        message: seed,
      });

      const personaSystemInstruction = await loadPrompt(
        "PERSONA_BASE_SYSTEM_PROMPT",
      );
      const baseSystemPrompt = await llm.call<string>(llm.models.identity, {
        instruction: personaSystemInstruction,
        message: description,
      });

      const db = await IdentityDB.connect({
        client: "sqlite",
        filename: config.dbPath,
      });
      await db.initialize();
      const brainId = randomUUID();
      const spaceName = `brain:${brainId}`;
      const space = await db.upsertSpace({
        name: spaceName,
        description: displayName,
      });

      await db.ingestStatement(description, {
        extractor: factExtractor,
        spaceName,
      });

      const brainbase: BrainItem = {
        brainId,
        spaceName,
        displayName,
        baseSystemPrompt,
      };
      await brainManager.saveBrain(brainId, brainbase);

      return new Brain(db, space, brainbase);
    } catch (error) {
      const reason = error instanceof Error ? error.message : String(error);
      logger.error(`Failed to create brain "${displayName}": ${reason}`);
      return null;
    }
  }

  static async load(brainId: string): Promise<Brain | null> {
    const brain = await brainManager.loadBrain(brainId);
    if (!brain) return null;

    const db = await IdentityDB.connect({
      client: "sqlite",
      filename: config.dbPath,
    });

    const space = await db.getSpaceByName(brain.spaceName);
    if (!space) return null;

    return new Brain(db, space, brain);
  }

  static async createDebug(options: DebugOptions): Promise<Brain> {
    const db = await IdentityDB.connect({
      client: "sqlite",
      filename: ":memory:",
    });
    await db.initialize();
    const space = await db.upsertSpace({
      name: "debug",
      description: "Debug Brain",
    });

    const brainbase: BrainItem = {
      brainId: "debug",
      spaceName: "debug",
      displayName: "Debug Brain",
      baseSystemPrompt: options.personality,
    };

    return new Brain(db, space, brainbase, true);
  }
}