import { getSandbox, Sandbox } from '@cloudflare/sandbox';
import { DurableObject } from 'cloudflare:workers';
import { readReader } from 'common-utils/streams/readReader';
import { Deferred, makeDeferred } from 'common-utils/promises/makeDeferred';
import { withTimeout } from 'common-utils/promises/withTimeout.ts';

export type StageName =
  | 'previewServerStarted'
  | 'previewServerReadyForRequests'
  | 'nameSet'
  | 'portExposed';

class WorkspaceLifecycleCoordinator {
  private stages: Map<StageName, Deferred<void>>;

  constructor() {
    this.stages = new Map([
      ['previewServerStarted', makeDeferred<void>()],
      ['previewServerReadyForRequests', makeDeferred<void>()],
      ['nameSet', makeDeferred<void>()],
      ['portExposed', makeDeferred<void>()],
    ]);
  }

  markStageComplete(stage: StageName): void {
    const deferred = this.stages.get(stage);
    if (!deferred) {
      throw new Error(`Unknown stage: ${stage}`);
    }

    deferred.resolve();
  }

  markStageFailed(stage: StageName, error: Error): void {
    const deferred = this.stages.get(stage);
    if (!deferred) {
      throw new Error(`Unknown stage: ${stage}`);
    }

    deferred.reject(error);
  }

  async whenStageReady(stage: StageName, timeoutMs?: number): Promise<void> {
    const deferred = this.stages.get(stage);
    if (!deferred) {
      throw new Error(`Unknown stage: ${stage}`);
    }
    const promise = deferred.promise;
    if (timeoutMs !== undefined) {
      return withTimeout(promise, timeoutMs);
    }
    return promise;
  }
}

type Preview = {
  url: string;
  port: number;
  name: string | undefined;
};

const PREVIEW_SERVER_STARTUP_TIMEOUT_MS = 30 * 1000; // 30s
export class Workspace extends Sandbox<Env> {
  sleepAfter = '10m';
  static async provision(
    namespace: DurableObjectNamespace<Workspace>,
    id: number,
  ) {
    const sandboxNamespace = namespace as unknown as DurableObjectNamespace<
      Sandbox<unknown>
    >;
    const sandbox = getSandbox(
      sandboxNamespace,
      id.toString(),
    ) as unknown as DurableObjectStub<Workspace>;
    // This is a bit of a hack but the workspace name is only set after calling getSandbox so we need to mark it as set here.
    sandbox.markNameAsSet();
    return sandbox;
  }
  private preview: Preview | undefined;
  private previewServerProcessId: string | undefined;
  private lifecycleCoordinator = new WorkspaceLifecycleCoordinator();

  constructor(ctx: DurableObject['ctx'], env: Env) {
    super(ctx, env);
    console.log('Creating new Workspace instance');
    this.startPreviewServer();
    this.exposePreviewPort();
  }

  private isPreviewServerLogReady(chunk: string): boolean {
    return Boolean(chunk.match(/VITE\s+(v[\d.]+)\s+ready in/));
  }

  private async startPreviewServer(): Promise<void> {
    const serverSession = await this.createSession({
      id: 'serverSession',
      cwd: '/workspace',
    });
    const process = await serverSession.startProcess(
      'pnpm --filter react-sandbox-renderer start',
    );
    this.previewServerProcessId = process.id;
    if (!this.previewServerProcessId) {
      this.lifecycleCoordinator.markStageFailed(
        'previewServerReadyForRequests',
        new Error('Failed to start preview server process'),
      );
      console.error('Failed to start preview server process');
      return;
    }
    this.lifecycleCoordinator.markStageComplete('previewServerStarted');
    console.log(
      `Vite dev server process started with id: ${this.previewServerProcessId}`,
    );
    const logStream = await serverSession.streamProcessLogs(
      this.previewServerProcessId,
    );
    const reader = logStream.getReader();

    readReader({
      reader,
      onChunk: chunk => {
        if (this.isPreviewServerLogReady(chunk)) {
          console.log('✓ Vite dev server is ready');
          this.lifecycleCoordinator.markStageComplete(
            'previewServerReadyForRequests',
          );
          return false;
        }
        return true;
      },
      onError: error => {
        this.lifecycleCoordinator.markStageFailed(
          'previewServerReadyForRequests',
          error as Error,
        );
      },
      onEnd: () => {
        this.lifecycleCoordinator.markStageFailed(
          'previewServerReadyForRequests',
          new Error('Preview server process exited'),
        );
      },
    });
  }

  public markNameAsSet(): void {
    this.lifecycleCoordinator.markStageComplete('nameSet');
  }

  // needs to be called after getSandbox is called since the sandbox needs a name to expose its port
  async exposePreviewPort(): Promise<void> {
    await Promise.all([
      this.lifecycleCoordinator.whenStageReady('previewServerStarted'),
      this.lifecycleCoordinator.whenStageReady('nameSet'),
    ]);
    this.preview = await this.exposePort(8000, {
      hostname: this.env.WORKER_HOSTNAME,
    });
    this.lifecycleCoordinator.markStageComplete('portExposed');
  }

  getPreviewUrl(): string {
    if (!this.preview) {
      throw new Error('Preview port not exposed');
    }
    return `${this.preview.url}`;
  }

  async getPreviewServerProcessId(): Promise<string> {
    if (!this.previewServerProcessId) {
      throw new Error('Preview server process not started');
    }
    return this.previewServerProcessId;
  }

  async isPreviewReady(): Promise<void> {
    await Promise.all([
      this.lifecycleCoordinator.whenStageReady(
        'previewServerReadyForRequests',
        PREVIEW_SERVER_STARTUP_TIMEOUT_MS,
      ),
      this.lifecycleCoordinator.whenStageReady(
        'portExposed',
        PREVIEW_SERVER_STARTUP_TIMEOUT_MS,
      ),
    ]);
  }
}