package csvpipe

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/csv"
	"encoding/hex"
	"errors"
	"fmt"
	"hash"
	"io"
	"os"
	"path/filepath"
)

// Row is one CSV row in canonical column order.
type Row []string

// ObjectManifest records one emitted CSV source object.
type ObjectManifest struct {
	ObjectRef            string `json:"object_ref"`
	TableRef             string `json:"table_ref"`
	SourceURI            string `json:"source_uri"`
	HeaderPresent        bool   `json:"header_present"`
	RowCount             int64  `json:"row_count"`
	SizeBytes            int64  `json:"size_bytes"`
	SHA256               string `json:"sha256"`
	FirstKey             string `json:"first_key"`
	LastKey              string `json:"last_key"`
	OwnerRef             string `json:"owner_ref"`
	OwnerEpoch           int64  `json:"owner_epoch"`
	GenerationNodeRef    string `json:"generation_node_ref"`
	CSVObjectTargetBytes int64  `json:"csv_object_target_bytes"`
	StorageMode          string `json:"storage_mode"`
	UploadPartCount      int    `json:"upload_part_count"`
}

// WriteOptions configures CSV source-object writing for one table.
type WriteOptions struct {
	Context              context.Context
	OutputDir            string
	Sink                 ObjectSink
	TableRef             string
	ObjectRefPrefix      string
	Header               []string
	CSVObjectTargetBytes int64
	OwnerRef             string
	OwnerEpoch           int64
	GenerationNodeRef    string
	ObjectSeqStart       int
	ObjectObserver       func(ObjectManifest)
}

// ObjectWrite carries a complete CSV object to its configured source location.
type ObjectWrite struct {
	TableRef  string
	ObjectRef string
	FileName  string
	Data      []byte
	SizeBytes int64
	SHA256    string
}

// ObjectLocation identifies where TiDB can read an emitted CSV object.
type ObjectLocation struct {
	SourceURI string
}

// ObjectWriteResult records bytes and digest produced by an object stream.
type ObjectWriteResult struct {
	SizeBytes   int64
	SHA256      string
	StorageMode string
	PartCount   int
}

// ObjectSink stores one bounded CSV object and returns the importable source URI.
type ObjectSink interface {
	WriteObject(context.Context, ObjectWrite) (ObjectLocation, error)
}

// ObjectStreamSink stores one CSV object incrementally as rows are encoded.
type ObjectStreamSink interface {
	BeginObject(context.Context, ObjectWrite) (ObjectStream, error)
}

// ObjectStream accepts one object's bytes and closes with the final source URI.
type ObjectStream interface {
	io.Writer
	Close() (ObjectLocation, ObjectWriteResult, error)
	Abort() error
}

// LocalFilesystemSink stores CSV objects under one root directory.
type LocalFilesystemSink struct {
	RootDir string
}

// WriteObject writes one CSV object while preserving table directories.
func (s LocalFilesystemSink) WriteObject(_ context.Context, object ObjectWrite) (ObjectLocation, error) {
	if s.RootDir == "" {
		return ObjectLocation{}, fmt.Errorf("output dir is empty")
	}
	tableDir := filepath.Join(s.RootDir, object.TableRef)
	if err := os.MkdirAll(tableDir, 0o755); err != nil {
		return ObjectLocation{}, err
	}
	path := filepath.Join(tableDir, object.FileName)
	sourceURI, err := filepath.Abs(path)
	if err != nil {
		return ObjectLocation{}, err
	}
	if err := os.WriteFile(sourceURI, object.Data, 0o644); err != nil {
		return ObjectLocation{}, err
	}
	return ObjectLocation{SourceURI: sourceURI}, nil
}

// BeginObject opens one local CSV object for incremental writing.
func (s LocalFilesystemSink) BeginObject(_ context.Context, object ObjectWrite) (ObjectStream, error) {
	if s.RootDir == "" {
		return nil, fmt.Errorf("output dir is empty")
	}
	tableDir := filepath.Join(s.RootDir, object.TableRef)
	if err := os.MkdirAll(tableDir, 0o755); err != nil {
		return nil, err
	}
	path := filepath.Join(tableDir, object.FileName)
	sourceURI, err := filepath.Abs(path)
	if err != nil {
		return nil, err
	}
	file, err := os.Create(sourceURI)
	if err != nil {
		return nil, err
	}
	return &localObjectStream{file: file, path: sourceURI, hasher: sha256.New()}, nil
}

type localObjectStream struct {
	file   *os.File
	path   string
	hasher hash.Hash
	size   int64
	closed bool
}

func (s *localObjectStream) Write(data []byte) (int, error) {
	if s.closed {
		return 0, fmt.Errorf("CSV object stream is closed")
	}
	n, err := s.file.Write(data)
	if n > 0 {
		_, _ = s.hasher.Write(data[:n])
		s.size += int64(n)
	}
	return n, err
}

func (s *localObjectStream) Close() (ObjectLocation, ObjectWriteResult, error) {
	if s.closed {
		return ObjectLocation{}, ObjectWriteResult{}, fmt.Errorf("CSV object stream is closed")
	}
	s.closed = true
	closeErr := s.file.Close()
	result := ObjectWriteResult{SizeBytes: s.size, SHA256: hex.EncodeToString(s.hasher.Sum(nil)), StorageMode: "local_file", PartCount: 1}
	if closeErr != nil {
		return ObjectLocation{}, result, closeErr
	}
	return ObjectLocation{SourceURI: s.path}, result, nil
}

func (s *localObjectStream) Abort() error {
	if !s.closed {
		s.closed = true
		_ = s.file.Close()
	}
	return os.Remove(s.path)
}

// RowIterator streams CSV rows without requiring a full table in memory.
type RowIterator interface {
	Next() (Row, bool, error)
}

// SliceIterator adapts an in-memory row slice to RowIterator.
type SliceIterator struct {
	Rows  []Row
	index int
}

// Next returns the next row from the slice iterator.
func (i *SliceIterator) Next() (Row, bool, error) {
	if i.index >= len(i.Rows) {
		return nil, false, nil
	}
	row := i.Rows[i.index]
	i.index++
	return row, true, nil
}

// WriteObjects writes a bounded slice of rows into CSV source objects.
func WriteObjects(opts WriteOptions, rows []Row) ([]ObjectManifest, error) {
	return WriteIterator(opts, &SliceIterator{Rows: rows})
}

// WriteIterator streams rows into CSV source objects without splitting rows.
func WriteIterator(opts WriteOptions, rows RowIterator) (objects []ObjectManifest, returnErr error) {
	if opts.Context == nil {
		opts.Context = context.Background()
	}
	sink := opts.Sink
	if sink == nil {
		sink = LocalFilesystemSink{RootDir: opts.OutputDir}
	}
	streamSink := streamSinkFor(sink)
	if streamSink == nil {
		return nil, fmt.Errorf("CSV object stream sink is empty")
	}
	if opts.TableRef == "" {
		return nil, fmt.Errorf("table ref is empty")
	}
	if len(opts.Header) == 0 {
		return nil, fmt.Errorf("CSV header is empty")
	}
	if opts.CSVObjectTargetBytes <= 0 {
		return nil, fmt.Errorf("CSV object target bytes must be positive")
	}
	header, err := EncodeRecord(Row(opts.Header))
	if err != nil {
		return nil, err
	}
	if int64(len(header)) > opts.CSVObjectTargetBytes {
		return nil, fmt.Errorf("CSV header for table %s exceeds object target", opts.TableRef)
	}

	objects = make([]ObjectManifest, 0)
	current := newObjectBuilder(opts.Context, opts, streamSink, len(objects), header)
	defer func() {
		returnErr = errors.Join(returnErr, current.abortOpen())
	}()
	for {
		row, ok, err := rows.Next()
		if err != nil {
			return nil, err
		}
		if !ok {
			break
		}
		encoded, err := EncodeRecord(row)
		if err != nil {
			return nil, err
		}
		if int64(len(header)+len(encoded)) > opts.CSVObjectTargetBytes {
			return nil, fmt.Errorf("single CSV row for table %s exceeds object target", opts.TableRef)
		}
		// The target is a rolling object size, not a hard split point: rows
		// stay intact and objects are never padded to meet the target.
		if current.rowCount > 0 && current.size()+int64(len(encoded)) > opts.CSVObjectTargetBytes {
			current.ownsProducingStream = false
			manifest, err := current.flush()
			if err != nil {
				return nil, err
			}
			objects = append(objects, manifest)
			current = newObjectBuilder(opts.Context, opts, streamSink, len(objects), header)
		}
		if err := current.add(row, encoded); err != nil {
			return nil, err
		}
	}
	if current.rowCount > 0 {
		current.ownsProducingStream = false
		manifest, err := current.flush()
		if err != nil {
			return nil, err
		}
		objects = append(objects, manifest)
	}
	return objects, nil
}

func streamSinkFor(sink ObjectSink) ObjectStreamSink {
	if streamSink, ok := sink.(ObjectStreamSink); ok {
		return streamSink
	}
	if sink == nil {
		return nil
	}
	// This compatibility wrapper buffers a whole object for legacy sinks; formal
	// S3 publication must provide ObjectStreamSink to stay bounded-memory.
	return bufferedObjectSink{sink: sink}
}

type bufferedObjectSink struct {
	sink ObjectSink
}

func (s bufferedObjectSink) BeginObject(_ context.Context, object ObjectWrite) (ObjectStream, error) {
	return &bufferedObjectStream{sink: s.sink, object: object, hasher: sha256.New()}, nil
}

type bufferedObjectStream struct {
	sink   ObjectSink
	object ObjectWrite
	buffer bytes.Buffer
	hasher hash.Hash
	closed bool
}

func (s *bufferedObjectStream) Write(data []byte) (int, error) {
	if s.closed {
		return 0, fmt.Errorf("CSV object stream is closed")
	}
	n, err := s.buffer.Write(data)
	if n > 0 {
		_, _ = s.hasher.Write(data[:n])
	}
	return n, err
}

func (s *bufferedObjectStream) Close() (ObjectLocation, ObjectWriteResult, error) {
	if s.closed {
		return ObjectLocation{}, ObjectWriteResult{}, fmt.Errorf("CSV object stream is closed")
	}
	s.closed = true
	data := s.buffer.Bytes()
	s.object.Data = data
	s.object.SizeBytes = int64(len(data))
	s.object.SHA256 = hex.EncodeToString(s.hasher.Sum(nil))
	location, err := s.sink.WriteObject(context.Background(), s.object)
	return location, ObjectWriteResult{SizeBytes: s.object.SizeBytes, SHA256: s.object.SHA256, StorageMode: "buffered_object_sink", PartCount: 1}, err
}

func (s *bufferedObjectStream) Abort() error {
	s.closed = true
	s.buffer.Reset()
	return nil
}

type objectBuilder struct {
	ctx     context.Context
	opts    WriteOptions
	sink    ObjectStreamSink
	ordinal int
	header  []byte
	stream  ObjectStream
	// ownsProducingStream is cleared before flush transfers cleanup to Close.
	ownsProducingStream bool
	sizeBytes           int64
	rowCount            int64
	firstKey            string
	lastKey             string
}

func newObjectBuilder(ctx context.Context, opts WriteOptions, sink ObjectStreamSink, ordinal int, header []byte) *objectBuilder {
	return &objectBuilder{ctx: ctx, opts: opts, sink: sink, ordinal: ordinal, header: header, sizeBytes: int64(len(header))}
}

func (b *objectBuilder) size() int64 {
	return b.sizeBytes
}

func (b *objectBuilder) add(row Row, encoded []byte) error {
	if b.stream == nil {
		stream, err := b.sink.BeginObject(b.ctx, b.objectWrite(nil, ""))
		if err != nil {
			return err
		}
		b.stream = stream
		b.ownsProducingStream = true
		if err := writeAll(b.stream, b.header); err != nil {
			return errors.Join(err, b.abortOpen())
		}
	}
	if len(row) > 0 {
		if b.rowCount == 0 {
			b.firstKey = row[0]
		}
		b.lastKey = row[0]
	}
	if err := writeAll(b.stream, encoded); err != nil {
		return errors.Join(err, b.abortOpen())
	}
	b.sizeBytes += int64(len(encoded))
	b.rowCount++
	return nil
}

func (b *objectBuilder) abortOpen() error {
	if b == nil || b.stream == nil || !b.ownsProducingStream {
		return nil
	}
	stream := b.stream
	b.stream = nil
	b.ownsProducingStream = false
	return stream.Abort()
}

func (b *objectBuilder) flush() (ObjectManifest, error) {
	if b.stream == nil {
		return ObjectManifest{}, fmt.Errorf("CSV object has no rows")
	}
	location, result, err := b.stream.Close()
	if err != nil {
		return ObjectManifest{}, err
	}
	if result.SizeBytes != b.sizeBytes {
		return ObjectManifest{}, fmt.Errorf("CSV object %s size drifted: %d != %d", b.objectRef(), result.SizeBytes, b.sizeBytes)
	}
	if result.SHA256 == "" {
		return ObjectManifest{}, fmt.Errorf("CSV object %s has empty checksum", b.objectRef())
	}
	if location.SourceURI == "" {
		return ObjectManifest{}, fmt.Errorf("CSV object %s has empty source URI", b.objectRef())
	}
	manifest := ObjectManifest{
		ObjectRef:            b.objectRef(),
		TableRef:             b.opts.TableRef,
		SourceURI:            location.SourceURI,
		HeaderPresent:        true,
		RowCount:             b.rowCount,
		SizeBytes:            result.SizeBytes,
		SHA256:               result.SHA256,
		FirstKey:             b.firstKey,
		LastKey:              b.lastKey,
		OwnerRef:             b.opts.OwnerRef,
		OwnerEpoch:           b.opts.OwnerEpoch,
		GenerationNodeRef:    b.opts.GenerationNodeRef,
		CSVObjectTargetBytes: b.opts.CSVObjectTargetBytes,
		StorageMode:          result.StorageMode,
		UploadPartCount:      result.PartCount,
	}
	if b.opts.ObjectObserver != nil {
		b.opts.ObjectObserver(manifest)
	}
	return manifest, nil
}

func (b *objectBuilder) objectWrite(data []byte, sha string) ObjectWrite {
	objectRef := b.objectRef()
	return ObjectWrite{
		TableRef:  b.opts.TableRef,
		ObjectRef: objectRef,
		FileName:  objectRef + ".csv",
		Data:      data,
		SizeBytes: int64(len(data)),
		SHA256:    sha,
	}
}

func (b *objectBuilder) objectRef() string {
	prefix := b.opts.ObjectRefPrefix
	if prefix == "" {
		prefix = b.opts.TableRef
	}
	return fmt.Sprintf("%s-%06d", prefix, b.opts.ObjectSeqStart+b.ordinal)
}

func writeAll(writer io.Writer, data []byte) error {
	for len(data) > 0 {
		n, err := writer.Write(data)
		if err != nil {
			return err
		}
		if n <= 0 {
			return io.ErrShortWrite
		}
		data = data[n:]
	}
	return nil
}

// EncodeRecord encodes one CSV data record using the same canonical encoder as
// WriteIterator. The returned bytes include the record terminator.
func EncodeRecord(record Row) ([]byte, error) {
	var buf bytes.Buffer
	writer := csv.NewWriter(&buf)
	if err := writer.Write(record); err != nil {
		return nil, err
	}
	writer.Flush()
	if err := writer.Error(); err != nil {
		return nil, err
	}
	return buf.Bytes(), nil
}

// EncodedRecordSize returns the byte length of EncodeRecord's canonical output.
func EncodedRecordSize(record Row) (uint64, error) {
	encoded, err := EncodeRecord(record)
	if err != nil {
		return 0, err
	}
	return uint64(len(encoded)), nil
}