import {
  GetObjectCommand,
  ListObjectsV2CommandOutput,
  PutObjectCommandInput,
  S3,
} from "@aws-sdk/client-s3";
import { Upload } from "@aws-sdk/lib-storage";
import { NodeHttpHandler } from "@smithy/node-http-handler";
import config from "../../config";
import { pipeline, Readable, Transform, PassThrough } from "stream";
import ArchiveStreamToS3 from "decompress-stream-to-s3";
import { Response } from "express";
import { lookup } from "mime-types";
import * as archiver from "archiver";
import { dirname, basename, join } from "path";
import AnonymousError from "../AnonymousError";
import StorageBase, { FILE_TYPE } from "./Storage";
import { IFile } from "../model/files/files.types";
import FileModel from "../model/files/files.model";
import { createLogger, serializeError } from "../logger";

const logger = createLogger("s3");

export default class S3Storage extends StorageBase {
  type = "AWS";

  constructor() {
    super();
    if (!config.S3_BUCKET)
      throw new AnonymousError("s3_config_not_provided", {
        httpStatus: 500,
      });
  }

  private client(timeout = 10000) {
    if (!config.S3_CLIENT_ID) throw new Error("S3_CLIENT_ID not set");
    if (!config.S3_CLIENT_SECRET) throw new Error("S3_CLIENT_SECRET not set");
    if (!config.S3_REGION) throw new Error("S3_REGION not set");
    if (!config.S3_ENDPOINT) throw new Error("S3_ENDPOINT not set");
    return new S3({
      credentials: {
        accessKeyId: config.S3_CLIENT_ID,
        secretAccessKey: config.S3_CLIENT_SECRET,
      },
      region: config.S3_REGION,
      endpoint: config.S3_ENDPOINT,
      requestHandler: new NodeHttpHandler({
        requestTimeout: timeout,
        connectionTimeout: timeout,
      }),
    });
  }

  /** @override */
  async exists(repoId: string, path: string = ""): Promise<FILE_TYPE> {
    if (!config.S3_BUCKET) throw new Error("S3_BUCKET not set");
    this.assertSafePath(path);
    try {
      // if we can get the file info, it is a file
      await this.fileInfo(repoId, path);
      return FILE_TYPE.FILE;
    } catch {
      // check if it is a directory
      const data = await this.client().listObjectsV2({
        Bucket: config.S3_BUCKET,
        Prefix: join(this.repoPath(repoId), path),
        MaxKeys: 1,
      });
      return (data.Contents?.length || 0) > 0
        ? FILE_TYPE.FOLDER
        : FILE_TYPE.NOT_FOUND;
    }
  }

  /** @override */
  async mk(repoId: string, _dir: string = ""): Promise<void> {
    // no need to create folder on S3
  }

  /** @override */
  async rm(repoId: string, dir: string = ""): Promise<void> {
    if (!config.S3_BUCKET) throw new Error("S3_BUCKET not set");
    this.assertSafePath(dir);
    const data = await this.client(200000).listObjectsV2({
      Bucket: config.S3_BUCKET,
      Prefix: join(this.repoPath(repoId), dir),
      MaxKeys: 100,
    });

    const params = {
      Bucket: config.S3_BUCKET,
      Delete: { Objects: new Array<{ Key: string }>() },
    };

    data.Contents?.forEach(function (content) {
      if (content.Key) {
        params.Delete.Objects.push({ Key: content.Key });
      }
    });

    if (params.Delete.Objects.length == 0) {
      // nothing to remove
      return;
    }
    await this.client(200000).deleteObjects(params);

    if (data.IsTruncated) {
      await this.rm(repoId, dir);
    }
  }

  /** @override */
  async send(repoId: string, path: string, res: Response) {
    if (!config.S3_BUCKET) throw new Error("S3_BUCKET not set");
    this.assertSafePath(path);
    try {
      const command = new GetObjectCommand({
        Bucket: config.S3_BUCKET,
        Key: join(this.repoPath(repoId), path),
      });
      const s = await this.client().send(command);
      res.status(200);
      if (s.ContentType) {
        res.contentType(s.ContentType);
      }
      if (s.ContentLength) {
        res.set("Content-Length", s.ContentLength.toString());
      }
      if (s.Body) {
        const body = s.Body as Readable;
        body.on("error", (err) => {
          logger.error("S3 body stream error", { ...serializeError(err), filePath: path });
          if (!res.headersSent) res.status(502).json({ error: "storage_read_error" });
          else res.destroy();
        });
        body.pipe(res);
      } else {
        res.end();
      }
    } catch {
      try {
        res.status(500).json({ error: "file_not_found" });
      } catch (err) {
        logger.error("send failed", { ...serializeError(err), filePath: path });
      }
    }
  }

  async fileInfo(repoId: string, path: string) {
    if (!config.S3_BUCKET) throw new Error("S3_BUCKET not set");
    this.assertSafePath(path);
    const info = await this.client(3000).headObject({
      Bucket: config.S3_BUCKET,
      Key: join(this.repoPath(repoId), path),
    });
    return {
      size: info.ContentLength,
      lastModified: info.LastModified,
      contentType: info.ContentType
        ? info.ContentType
        : (lookup(path) as string),
    };
  }

  /** @override */
  async read(repoId: string, path: string): Promise<Readable> {
    if (!config.S3_BUCKET) throw new Error("S3_BUCKET not set");
    this.assertSafePath(path);
    const command = new GetObjectCommand({
      Bucket: config.S3_BUCKET,
      Key: join(this.repoPath(repoId), path),
    });
    const res = (await this.client(3000).send(command)).Body;
    if (!res) {
      throw new AnonymousError("file_not_found", {
        httpStatus: 404,
        url: join(this.repoPath(repoId), path),
      });
    }
    return res as Readable;
  }

  /** @override */
  async write(
    repoId: string,
    path: string,
    data: string | Readable,
    source?: string,
    expectedSize?: number
  ): Promise<void> {
    if (!config.S3_BUCKET) throw new Error("S3_BUCKET not set");
    this.assertSafePath(path);

    // No fire-and-forget rm on stream error: Upload uses multipart and
    // does not commit a partially-uploaded object, so there's nothing to
    // clean up. The previous handler raced with retries and could delete
    // a previously-good object on a transient source-stream hiccup. The
    // size-validated read path in GitHubStream.getFileContentCache
    // recovers from any object that does end up undersized for any
    // reason.

    // When we know the upstream byte count, count bytes through a
    // PassThrough so we can refuse to commit a short read. The Upload
    // client itself doesn't expose bytesWritten, and S3 will happily
    // PutObject a truncated body if the source stream ends cleanly with
    // fewer bytes than expected.
    let body: string | Readable = data;
    let bytesWritten = -1;
    if (typeof expectedSize === "number" && expectedSize > 0 && typeof data !== "string") {
      bytesWritten = 0;
      const counter = new PassThrough();
      counter.on("data", (chunk: Buffer) => {
        bytesWritten += chunk.length;
      });
      data.on("error", (err: Error) => counter.destroy(err));
      data.pipe(counter);
      body = counter;
    }

    const params: PutObjectCommandInput = {
      Bucket: config.S3_BUCKET,
      Key: join(this.repoPath(repoId), path),
      Body: body,
      ContentType: lookup(path).toString(),
    };
    if (source) {
      params.Tagging = `source=${source}`;
    }

    const parallelUploads3 = new Upload({
      // 30s timeout
      client: this.client(30000),
      params,
    });

    await parallelUploads3.done();

    if (
      typeof expectedSize === "number" &&
      expectedSize > 0 &&
      bytesWritten >= 0 &&
      bytesWritten < expectedSize
    ) {
      // The body completed cleanly but produced fewer bytes than the
      // tree said — delete the truncated object so the next request
      // re-fetches from GitHub instead of serving the short blob.
      await this.rm(repoId, path).catch(() => undefined);
      throw new AnonymousError("storage_write_size_mismatch", {
        httpStatus: 502,
        object: { path, expected: expectedSize, actual: bytesWritten },
      });
    }
  }

  /** @override */
  async listFiles(repoId: string, dir: string = ""): Promise<IFile[]> {
    if (!config.S3_BUCKET) throw new Error("S3_BUCKET not set");
    this.assertSafePath(dir);
    if (dir && dir[dir.length - 1] != "/") dir = dir + "/";
    const out: IFile[] = [];
    let req: ListObjectsV2CommandOutput;
    let nextContinuationToken: string | undefined;
    do {
      req = await this.client(30000).listObjectsV2({
        Bucket: config.S3_BUCKET,
        Prefix: join(this.repoPath(repoId), dir),
        MaxKeys: 250,
        ContinuationToken: nextContinuationToken,
      });
      if (!req.Contents) return out;
      nextContinuationToken = req.NextContinuationToken;

      for (const f of req.Contents) {
        if (!f.Key) continue;
        f.Key = f.Key.replace(join(this.repoPath(repoId), dir), "");
        out.push(
          new FileModel({
            name: basename(f.Key),
            path: dirname(f.Key),
            repoId,
            size: f.Size,
            sha: f.ETag,
          })
        );
      }
    } while (req && req.Contents && req.IsTruncated);
    return out;
  }

  /** @override */
  async extractZip(
    repoId: string,
    path: string,
    data: Readable,
    source?: string
  ): Promise<void> {
    this.assertSafePath(path);
    let toS3: ArchiveStreamToS3;
    return new Promise((resolve, reject) => {
      if (!config.S3_BUCKET) return reject("S3_BUCKET not set");
      toS3 = new ArchiveStreamToS3({
        bucket: config.S3_BUCKET,
        prefix: join(this.repoPath(repoId), path),
        s3: this.client(2 * 60 * 60 * 1000), // 2h timeout
        type: "zip",
        onEntry: (header) => {
          // Strip the wrapping top-level dir, then drop any "../" / absolute
          // segments so a crafted entry name cannot write objects outside
          // the repo key prefix (zip-slip, CWE-23).
          const stripped = header.name.substring(
            header.name.indexOf("/") + 1
          );
          header.name = this.sanitizeZipEntryName(stripped);
          if (source) {
            header.Tagging = `source=${source}`;
            header.Metadata = {
              source: source,
            };
          }
        },
        maxParallel: 10,
      });
      pipeline(data, toS3, (err) => {
        if (err) {
          return reject(err);
        }
        resolve();
      })
        .on("finish", () => {
          resolve();
        })
        .on("error", reject);
    });
  }

  /** @override */
  async archive(
    repoId: string,
    dir: string = "",
    opt?: {
      format?: "zip" | "tar";
      fileTransformer?: (p: string) => Transform;
    }
  ) {
    if (!config.S3_BUCKET) throw new Error("S3_BUCKET not set");
    this.assertSafePath(dir);
    const archive = archiver(opt?.format || "zip", {});
    if (dir && dir[dir.length - 1] != "/") dir = dir + "/";

    let req: ListObjectsV2CommandOutput;
    let nextContinuationToken: string | undefined;
    do {
      req = await this.client(30000).listObjectsV2({
        Bucket: config.S3_BUCKET,
        Prefix: join(this.repoPath(repoId), dir),
        MaxKeys: 250,
        ContinuationToken: nextContinuationToken,
      });

      nextContinuationToken = req.NextContinuationToken;
      for (const f of req.Contents || []) {
        if (!f.Key) continue;
        const filename = basename(f.Key);
        const prefix = dirname(
          f.Key.replace(join(this.repoPath(repoId), dir), "")
        );

        let rs: Readable = await this.read(repoId, f.Key);
        if (opt?.fileTransformer) {
          const src = rs;
          const transformer = opt.fileTransformer(f.Key);
          src.on("error", (err) => transformer.destroy(err));
          rs = src.pipe(transformer);
        }

        archive.append(rs, {
          name: filename,
          prefix,
        });
      }
    } while (req && req.Contents?.length && req.IsTruncated);
    archive.finalize();
    return archive;
  }
}
