import { appendArchiveEntry } from "../archive-entry";
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 } 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 clients = new Map<number, S3>();

  private client(timeout = 10000) {
    const existing = this.clients.get(timeout);
    if (existing) return existing;
    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");
    const client = new S3({
      credentials: {
        accessKeyId: config.S3_CLIENT_ID,
        secretAccessKey: config.S3_CLIENT_SECRET,
      },
      region: config.S3_REGION,
      endpoint: config.S3_ENDPOINT,
      forcePathStyle: true,
      requestHandler: new NodeHttpHandler({
        requestTimeout: timeout,
        connectionTimeout: timeout,
      }),
    });
    this.clients.set(timeout, client);
    return client;
  }

  /** @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 prefix = join(this.repoPath(repoId), dir).replace(/\/$/, "");
    await this.removeMatching(prefix, key => key === prefix || key.startsWith(prefix + "/"));
  }

  async removeLegacyContent(repoId: string): Promise<void> {
    if (!config.S3_BUCKET) throw new Error("S3_BUCKET not set");
    const prefix = this.repoPath(repoId);
    await this.removeMatching(prefix, key => key.startsWith(prefix) &&
      !/^__content\/(?:[a-f0-9]{64}|__retired)(?:\/|$)/.test(key.slice(prefix.length)));
  }

  private async removeMatching(prefix: string, shouldRemove: (key: string) => boolean) {
    if (!config.S3_BUCKET) throw new Error("S3_BUCKET not set");
    const client = this.client(200000);
    let continuationToken: string | undefined;
    do {
      const data = await client.listObjectsV2({
        Bucket: config.S3_BUCKET,
        Prefix: prefix,
        MaxKeys: 1000,
        ContinuationToken: continuationToken,
      });
      const objects = (data.Contents || [])
        .filter(({ Key }) => Key && shouldRemove(Key))
        .map(({ Key }) => ({ Key: Key! }));
      if (objects.length) {
        const result = await client.deleteObjects({
          Bucket: config.S3_BUCKET,
          Delete: { Objects: objects },
        });
        if (result.Errors?.length) {
          throw new AnonymousError("storage_delete_failed", {
            httpStatus: 502,
            object: result.Errors,
          });
        }
      }
      continuationToken = data.IsTruncated ? data.NextContinuationToken : undefined;
      if (data.IsTruncated && !continuationToken) {
        throw new Error("S3 returned a truncated listing without a continuation token");
      }
    } while (continuationToken);
  }

  /** @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 Transform({
        transform(chunk, _encoding, callback) { bytesWritten += chunk.length; callback(null, chunk); },
        flush(callback) {
          callback(bytesWritten < expectedSize! ? new AnonymousError("storage_write_size_mismatch", { httpStatus: 502 }) : undefined);
        },
      });
      data.on("error", (err: Error) => counter.destroy(err));
      counter.once("close", () => (data as Readable).destroy());
      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 active: Readable | undefined;
    let stopped = false;
    archive.once("close", () => { stopped = true; active?.destroy(); });
    void (async () => {
      let next: string | undefined;
      do {
        const listing = await this.client(30000).listObjectsV2({ Bucket: config.S3_BUCKET!,
          Prefix: join(this.repoPath(repoId), dir), MaxKeys: 250, ContinuationToken: next });
        next = listing.IsTruncated ? listing.NextContinuationToken : undefined;
        for (const file of listing.Contents || []) {
          if (stopped) return;
          if (!file.Key || file.Key.endsWith("/")) continue;
          const filename = file.Key.slice(this.repoPath(repoId).length);
          const source = await this.read(repoId, filename);
          if (stopped) { source.destroy(); return; }
          active = source;
          let output = source;
          if (opt?.fileTransformer) {
            const transformer = opt.fileTransformer(filename);
            source.on("error", error => transformer.destroy(error));
            transformer.once("close", () => source.destroy());
            output = source.pipe(transformer);
            active = output;
          }
          await appendArchiveEntry(archive, output, { name: filename.slice(dir.length) });
        }
        if (listing.IsTruncated && !next) throw new Error("Invalid S3 continuation token");
      } while (next && !stopped);
      if (!stopped) await archive.finalize();
    })().catch(error => archive.destroy(error));
    return archive;
  }
}