// Copyright (c) 2019-2026 Alexander Medvednikov. All rights reserved.
// Use of this source code is governed by an MIT license
// that can be found in the LICENSE file.
module s3

import os
import strings
import time

// Multipart upload constants — defaults aligned with the S3 protocol limits.
pub const min_part_size = i64(5 * 1024 * 1024) // 5 MiB
pub const max_part_size = i64(5 * 1024 * 1024 * 1024) // 5 GiB
pub const max_parts = 10000 // hard S3 limit

// upload_file streams a local file to S3, choosing single-part or multipart
// based on size. Use this for files larger than ~50 MiB or anything you
// don't want to slurp into memory.
//
// `key` is the destination object key. The local file is read in
// `client.part_size` chunks (default 5 MiB).
pub fn (c &Client) upload_file(key string, local_path string, opts PutOptions) ! {
    stat_local := os.stat(local_path) or {
        return new_error('LocalFileNotFound',
            'Cannot stat local file: ${local_path} (${err.msg()})')
    }
    size := i64(stat_local.size)
    if size <= min_part_size {
        data := os.read_bytes(local_path) or {
            return new_error('LocalReadFailed', 'Cannot read ${local_path}: ${err.msg()}')
        }
        return c.put(key, data, opts)
    }
    c.upload_file_multipart(key, local_path, size, opts)!
}

// MultipartUploader is a stateful handle to an in-flight multipart upload,
// produced by `Client.start_multipart`. Use it when you want to push parts
// generated on-the-fly (network sources, decompressors, anything you don't
// want to materialise on disk):
//
//   mut up := c.start_multipart('key', s3.PutOptions{ content_type: 'application/octet-stream' })!
//   for chunk in chunks { up.upload(chunk)! }
//   up.complete()!
//
// On error, `complete()` / `upload()` return without aborting. The caller
// must invoke `abort()` themselves so that `defer { up.abort() or {} }`
// remains an explicit, visible cleanup hook.
pub struct MultipartUploader {
mut:
    client      &Client
    key         string
    upload_id   string
    opts        PutOptions
    parts       []PartRef
    part_number int
    completed   bool
    aborted     bool
}

// start_multipart begins a multipart upload and returns a MultipartUploader.
// Memory cost: zero — each `upload(chunk)` call streams the chunk to S3 and
// returns when it is acknowledged.
pub fn (c &Client) start_multipart(key string, opts PutOptions) !MultipartUploader {
    upload_id := c.initiate_multipart(key, opts)!
    return MultipartUploader{
        // `unsafe { c }` is required to store the receiver's pointer in the
        // returned struct — V's borrow checker forbids this implicitly, even
        // though `Client` is `@[heap]` so the lifetime is safe.
        client:      unsafe { c }
        key:         key
        upload_id:   upload_id
        opts:        opts
        parts:       []PartRef{}
        part_number: 1
    }
}

// upload pushes one chunk as the next part. Each chunk MUST be at least
// `min_part_size` (5 MiB) except the last one — that's an S3 invariant
// the server enforces; we do not buffer for you.
pub fn (mut u MultipartUploader) upload(data []u8) ! {
    if u.completed || u.aborted {
        return new_error('InvalidState', 'multipart upload is already finalised')
    }
    if u.part_number > max_parts {
        return new_error('TooManyParts', 'Exceeded ${max_parts} parts — increase part_size')
    }
    etag := u.client.upload_part(u.key, u.upload_id, u.part_number, data, u.opts)!
    u.parts << PartRef{
        part_number: u.part_number
        etag:        etag
    }
    u.part_number++
}

// complete finalises the upload. After this call the object is visible.
pub fn (mut u MultipartUploader) complete() ! {
    if u.completed {
        return
    }
    u.client.complete_multipart(u.key, u.upload_id, u.parts, u.opts)!
    u.completed = true
}

// abort cancels the in-flight upload. Idempotent.
pub fn (mut u MultipartUploader) abort() ! {
    if u.aborted || u.completed {
        return
    }
    u.client.abort_multipart(u.key, u.upload_id, u.opts)!
    u.aborted = true
}

// upload_bytes_multipart uploads an in-memory byte buffer using multipart.
// Useful for large generated payloads (test fixtures up to 5 GiB). Uses up
// to `Client.queue_size` parallel part uploads.
pub fn (c &Client) upload_bytes_multipart(key string, data []u8, opts PutOptions) ! {
    upload_id := c.initiate_multipart(key, opts)!
    parts := c.run_parallel_parts(key, upload_id, opts, fn [data] (offset i64, want int) ![]u8 {
        end := if offset + i64(want) > i64(data.len) { i64(data.len) } else { offset + i64(want) }
        return data[offset..end]
    }, i64(data.len)) or {
        c.abort_multipart(key, upload_id, opts) or {}
        return err
    }
    c.complete_multipart(key, upload_id, parts, opts)!
}

// upload_file_multipart streams a local file part-by-part with up to
// `Client.queue_size` concurrent uploads. Peak memory is roughly
// `queue_size * part_size`.
pub fn (c &Client) upload_file_multipart(key string, local_path string, size i64, opts PutOptions) ! {
    mut f := os.open(local_path) or {
        return new_error('LocalReadFailed', 'open ${local_path}: ${err.msg()}')
    }
    defer {
        f.close()
    }
    upload_id := c.initiate_multipart(key, opts)!
    parts := c.run_parallel_parts(key, upload_id, opts, fn [mut f] (offset i64, want int) ![]u8 {
        mut buf := []u8{len: want}
        n := f.read_bytes_into(u64(offset), mut buf) or {
            return new_error('LocalReadFailed', 'read_bytes_into: ${err.msg()}')
        }
        if n <= 0 {
            return new_error('LocalReadFailed', 'unexpected EOF at offset ${offset}')
        }
        return buf[..n]
    }, size) or {
        c.abort_multipart(key, upload_id, opts) or {}
        return err
    }
    c.complete_multipart(key, upload_id, parts, opts)!
}

// PartJob is a single upload-part task pushed onto the worker queue.
struct PartJob {
    part_number int
    data        []u8
}

// PartResult carries the worker outcome back to the dispatcher.
struct PartResult {
    part_number int
    etag        string
    err         ?IError
}

// part_worker pulls jobs off `jobs` and pushes results onto `results`. Exits
// when `jobs` is closed.
fn (c &Client) part_worker(key string, upload_id string, opts PutOptions, jobs chan PartJob, results chan PartResult) {
    for {
        job := <-jobs or { return }
        etag := c.upload_part(key, upload_id, job.part_number, job.data, opts) or {
            results <- PartResult{
                part_number: job.part_number
                err:         err
            }
            continue
        }
        results <- PartResult{
            part_number: job.part_number
            etag:        etag
        }
    }
}

// run_parallel_parts dispatches `producer` chunks across a worker pool of
// `Client.queue_size` goroutines and gathers the resulting PartRefs. The
// caller is responsible for issuing `initiate_multipart` / `abort_multipart`
// / `complete_multipart` around this call. `producer(offset, want)` must
// return a freshly-allocated chunk of length ≤ `want` bytes (workers may
// hold the slice past the next call). Returns parts in completion order;
// `complete_multipart` re-sorts by part number before sending the manifest.
fn (c &Client) run_parallel_parts(key string, upload_id string, opts PutOptions, producer fn (offset i64, want int) ![]u8, total i64) ![]PartRef {
    workers := if c.queue_size < 1 { 1 } else { c.queue_size }
    part_size := if c.part_size < min_part_size { min_part_size } else { c.part_size }
    jobs := chan PartJob{cap: workers}
    results := chan PartResult{cap: workers}
    for _ in 0 .. workers {
        spawn c.part_worker(key, upload_id, opts, jobs, results)
    }

    mut parts := []PartRef{}
    mut first_err := ?IError(none)
    mut sent := 0
    mut received := 0
    mut offset := i64(0)
    mut part_number := 1

    for offset < total && first_err == none {
        if part_number > max_parts {
            first_err = new_error('TooManyParts',
                'Exceeded ${max_parts} parts — increase part_size')
            break
        }
        want := if offset + part_size > total { int(total - offset) } else { int(part_size) }
        chunk := producer(offset, want) or {
            first_err = err
            break
        }
        if chunk.len == 0 {
            break
        }
        jobs <- PartJob{
            part_number: part_number
            data:        chunk
        }
        sent++
        offset += i64(chunk.len)
        part_number++
        // Drain ready results without blocking — surfaces errors early so we
        // can stop dispatching new parts.
        for {
            select {
                r := <-results {
                    received++
                    if e := r.err {
                        if first_err == none {
                            first_err = e
                        }
                    } else {
                        parts << PartRef{
                            part_number: r.part_number
                            etag:        r.etag
                        }
                    }
                }
                else {
                    break
                }
            }
        }
    }
    jobs.close()
    for received < sent {
        r := <-results
        received++
        if e := r.err {
            if first_err == none {
                first_err = e
            }
        } else {
            parts << PartRef{
                part_number: r.part_number
                etag:        r.etag
            }
        }
    }
    results.close()
    if e := first_err {
        return e
    }
    return parts
}

// PartRef is one ETag/PartNumber pair recorded during multipart upload.
pub struct PartRef {
pub:
    part_number int
    etag        string
}

// initiate_multipart starts an upload and returns the UploadId. ACL,
// content-type, and friends from `opts` are sent here — they apply to the
// final object.
pub fn (c &Client) initiate_multipart(key string, opts PutOptions) !string {
    creds := c.creds_for(opts.bucket)
    path := build_object_path(creds, opts.bucket, key)!
    signed := sign_request(creds, SignRequest{
        method:        'POST'
        path:          path
        query:         'uploads='
        payload_hash:  empty_sha256
        extra_headers: put_object_headers(opts)
    })!
    resp := c.do_http(signed, '')!
    if resp.status_code != 200 {
        return new_http_error(resp.status_code, key, resp.body)
    }
    upload_id := extract_xml_tag(resp.body, 'UploadId')
    if upload_id == '' {
        return new_error('InvalidResponse', 'CreateMultipartUploadResult missing UploadId')
    }
    return upload_id
}

// upload_part uploads a single chunk and returns the server-side ETag. Up to
// `Client.retry` attempts with exponential backoff (200ms, 400ms, 800ms, …)
// on transient failures.
//
// The payload is always SHA-256-signed (not `UNSIGNED-PAYLOAD`) so the
// service validates byte-for-byte integrity end-to-end. Multipart uploads
// don't carry a Content-MD5 header by default; without payload signing a
// flipped bit in transit would silently produce a corrupt object that
// passes the multipart ETag check.
pub fn (c &Client) upload_part(key string, upload_id string, part_number int, data []u8, opts PutOptions) !string {
    creds := c.creds_for(opts.bucket)
    path := build_object_path(creds, opts.bucket, key)!
    query := canonical_query_string({
        'partNumber': part_number.str()
        'uploadId':   upload_id
    })
    // Hash the payload once — multi-MiB chunks should not be re-hashed per attempt.
    // The signature itself still has to be recomputed on each retry because the
    // `x-amz-date` header advances.
    payload_hash := sha256_hex(data)
    body := data.bytestr()
    max_attempts := if c.retry < 1 { 1 } else { c.retry }
    mut last_err := IError(new_error('NetworkError', 'upload_part: no attempt was made'))
    for attempt in 0 .. max_attempts {
        signed := sign_request(creds, SignRequest{
            method:       'PUT'
            path:         path
            query:        query
            payload_hash: payload_hash
        })!
        resp := c.do_http(signed, body) or {
            last_err = err
            if attempt + 1 < max_attempts {
                time.sleep((1 << attempt) * 200 * time.millisecond)
            }
            continue
        }
        if resp.status_code == 200 {
            etag := resp.header.get(.etag) or { '' }
            if etag == '' {
                return new_error('InvalidResponse', 'UploadPart returned no ETag')
            }
            return etag.trim('"')
        }
        last_err = new_http_error(resp.status_code, key, resp.body)
        if attempt + 1 < max_attempts {
            time.sleep((1 << attempt) * 200 * time.millisecond)
        }
    }
    return last_err
}

// complete_multipart finalizes a multipart upload. Parts must be in ascending
// `part_number` order; we sort defensively.
pub fn (c &Client) complete_multipart(key string, upload_id string, parts []PartRef, opts PutOptions) ! {
    mut sorted := parts.clone()
    sorted.sort(a.part_number < b.part_number)
    mut body := strings.new_builder(1024)
    body.write_string('<CompleteMultipartUpload xmlns="http://s3.amazonaws.com/doc/2006-03-01/">')
    for p in sorted {
        body.write_string('<Part><PartNumber>${p.part_number}</PartNumber><ETag>"${p.etag}"</ETag></Part>')
    }
    body.write_string('</CompleteMultipartUpload>')
    xml := body.str()
    creds := c.creds_for(opts.bucket)
    path := build_object_path(creds, opts.bucket, key)!
    q := 'uploadId=' + uri_encode_query(upload_id)
    signed := sign_request(creds, SignRequest{
        method:        'POST'
        path:          path
        query:         q
        payload_hash:  sha256_hex(xml.bytes())
        extra_headers: {
            'content-type': 'application/xml'
        }
    })!
    resp := c.do_http(signed, xml)!
    if resp.status_code != 200 {
        return new_http_error(resp.status_code, key, resp.body)
    }
    // CompleteMultipartUpload may return HTTP 200 with an <Error> body — surface those.
    if resp.body.contains('<Error>') {
        return new_http_error(200, key, resp.body)
    }
}

// abort_multipart cancels an in-flight upload. Best-effort — callers usually
// invoke it inside an error path and don't care about the result.
pub fn (c &Client) abort_multipart(key string, upload_id string, opts PutOptions) ! {
    creds := c.creds_for(opts.bucket)
    path := build_object_path(creds, opts.bucket, key)!
    q := 'uploadId=' + uri_encode_query(upload_id)
    signed := sign_request(creds, SignRequest{
        method:       'DELETE'
        path:         path
        query:        q
        payload_hash: empty_sha256
    })!
    resp := c.do_http(signed, '')!
    if resp.status_code !in [204, 200, 404] {
        return new_http_error(resp.status_code, key, resp.body)
    }
}