import { BaseCheckpointSaver } from '@langchain/langgraph';
import type { Checkpoint, CheckpointMetadata, CheckpointTuple } from '@langchain/langgraph';
import type { RunnableConfig } from '@langchain/core/runnables';
import type Redis from 'ioredis';
import type { FastifyBaseLogger } from 'fastify';

/**
 * Tiered checkpoint saver: Redis (hot) + Iceberg (cold)
 *
 * Hot path: Active checkpoints stored in Redis with TTL
 * Cold path: Durable storage in Iceberg for long-term retention
 *
 * Based on architecture discussion: Redis for active sessions,
 * Iceberg for durable storage with time-travel capabilities.
 */
export class TieredCheckpointSaver extends BaseCheckpointSaver<number> {
  private readonly HOT_TTL_SECONDS = 86400; // 24 hours
  private readonly KEY_PREFIX = 'ckpt:';

  constructor(
    private redis: Redis,
    private logger: FastifyBaseLogger,
    // Note: Iceberg writes are handled via Kafka + Flink for consistency
    // Reads can be implemented when needed using IcebergClient
    // private iceberg?: IcebergClient
  ) {
    super();
  }

  /**
   * Get checkpoint from Redis (hot) or Iceberg (cold)
   */
  async getTuple(config: RunnableConfig): Promise<CheckpointTuple | undefined> {
    const threadId = config.configurable?.thread_id as string;
    if (!threadId) {
      throw new Error('thread_id required in config.configurable');
    }

    const checkpointId = config.configurable?.checkpoint_id as string | undefined;

    this.logger.debug({ threadId, checkpointId }, 'Getting checkpoint');

    // Hot path: Try Redis first
    const key = this.getRedisKey(threadId, checkpointId);
    const cached = await this.redis.get(key);

    if (cached) {
      this.logger.debug({ threadId, checkpointId }, 'Checkpoint found in Redis (hot)');
      return this.deserialize(cached);
    }

    // Cold path: Load from Iceberg (if needed)
    // Note: Implement when Iceberg query is required
    // Can use IcebergClient to query gateway.checkpoints table
    // or set up a Kafka topic for checkpoint persistence

    this.logger.debug({ threadId, checkpointId }, 'Checkpoint not in Redis, Iceberg cold storage not yet implemented');
    return undefined;
  }

  /**
   * Save checkpoint to Redis (hot) and async flush to Iceberg (cold)
   */
  async put(
    config: RunnableConfig,
    checkpoint: Checkpoint,
    metadata: CheckpointMetadata
  ): Promise<RunnableConfig> {
    const threadId = config.configurable?.thread_id as string;
    if (!threadId) {
      throw new Error('thread_id required in config.configurable');
    }

    this.logger.debug({ threadId, checkpointId: checkpoint.id }, 'Saving checkpoint');

    const serialized = this.serialize(checkpoint, metadata);

    // Hot: Redis with TTL
    const key = this.getRedisKey(threadId, checkpoint.id);
    await this.redis.set(key, serialized, 'EX', this.HOT_TTL_SECONDS);

    // Also store latest checkpoint pointer
    const latestKey = this.getRedisKey(threadId);
    await this.redis.set(latestKey, serialized, 'EX', this.HOT_TTL_SECONDS);

    // Cold: Async flush to Iceberg (fire and forget)
    this.flushToIceberg(threadId, checkpoint, metadata).catch((error) => {
      this.logger.error({ error, threadId }, 'Failed to flush checkpoint to Iceberg');
    });

    return {
      configurable: {
        ...config.configurable,
        thread_id: threadId,
        checkpoint_id: checkpoint.id,
      },
    };
  }

  /**
   * List all checkpoints for a thread
   */
  async *list(
    config: RunnableConfig
  ): AsyncGenerator<CheckpointTuple> {
    const threadId = config.configurable?.thread_id as string;
    if (!threadId) {
      throw new Error('thread_id required in config.configurable');
    }

    // Try to get from Redis first
    const pattern = `${this.KEY_PREFIX}${threadId}:*`;
    const keys = await this.redis.keys(pattern);

    for (const key of keys) {
      const data = await this.redis.get(key);
      if (data) {
        const tuple = this.deserialize(data);
        if (tuple) {
          yield tuple;
        }
      }
    }

    // TODO: Also scan Iceberg for historical checkpoints
  }

  /**
   * Delete thread (for GDPR compliance)
   */
  async deleteThread(threadId: string): Promise<void> {
    this.logger.info({ threadId }, 'Deleting thread');

    const pattern = `${this.KEY_PREFIX}${threadId}*`;
    const keys = await this.redis.keys(pattern);
    if (keys.length > 0) {
      await this.redis.del(...keys);
    }

    // TODO: Also delete from Iceberg
    // await this.deleteFromIceberg(threadId);
  }

  /**
   * Put writes (required by BaseCheckpointSaver)
   */
  async putWrites(
    config: RunnableConfig,
    writes: [string, unknown][],
    taskId: string
  ): Promise<void> {
    // For this simple implementation, we just log writes
    // In a full implementation, you'd store pending writes separately
    const threadId = config.configurable?.thread_id;
    this.logger.debug({ threadId, taskId, writes }, 'Put writes called');
  }



  /**
   * Generate Redis key for checkpoint
   */
  private getRedisKey(threadId: string, checkpointId?: string): string {
    if (checkpointId) {
      return `${this.KEY_PREFIX}${threadId}:${checkpointId}`;
    }
    return `${this.KEY_PREFIX}${threadId}:latest`;
  }

  /**
   * Serialize checkpoint to JSON string
   */
  private serialize(checkpoint: Checkpoint, metadata: CheckpointMetadata): string {
    return JSON.stringify({
      checkpoint,
      metadata,
      savedAt: new Date().toISOString(),
    });
  }

  /**
   * Deserialize checkpoint from JSON string
   */
  private deserialize(data: string): CheckpointTuple | undefined {
    try {
      const parsed = JSON.parse(data);
      return {
        config: {
          configurable: {
            thread_id: parsed.checkpoint.id,
            checkpoint_id: parsed.checkpoint.id,
          },
        },
        checkpoint: parsed.checkpoint,
        metadata: parsed.metadata,
        parentConfig: undefined,
      };
    } catch (error) {
      this.logger.error({ error }, 'Failed to deserialize checkpoint');
      return undefined;
    }
  }

  /**
   * Async flush checkpoint to Iceberg for durable storage
   *
   * Note: For production, send to Kafka topic that Flink consumes:
   * - Topic: gateway_checkpoints
   * - Flink job writes to gateway.checkpoints Iceberg table
   * - Ensures consistent write pattern with rest of system
   */
  private async flushToIceberg(
    _threadId: string,
    checkpoint: Checkpoint,
    _metadata: CheckpointMetadata
  ): Promise<void> {
    // TODO: Send to Kafka topic for Flink processing
    // const kafkaMessage = {
    //   user_id: metadata.userId || '0',
    //   session_id: threadId,
    //   checkpoint_id: checkpoint.id,
    //   checkpoint_data: JSON.stringify(checkpoint),
    //   metadata: JSON.stringify(metadata),
    //   timestamp: Date.now() * 1000, // microseconds
    // };
    // await this.kafkaProducer.send({
    //   topic: 'gateway_checkpoints',
    //   messages: [{ value: JSON.stringify(kafkaMessage) }]
    // });

    this.logger.debug({ threadId: _threadId, checkpointId: checkpoint.id },
      'Checkpoint flush to Iceberg (via Kafka) not yet implemented');
  }


}
