package parquet

import (
	"context"
	"fmt"
	"os"
	"path/filepath"
	"time"

	"github.com/parquet-go/parquet-go"
	"github.com/pingcap/log"
	"github.com/tidbcloud/observability/cmd/diagnosis-compute/pkg/config"
	"github.com/tidbcloud/observability/cmd/diagnosis-compute/pkg/metrics"
	"github.com/tidbcloud/observability/cmd/diagnosis-compute/pkg/types"
	"github.com/tidbcloud/observability/o11ydiagnosis/pkg/storage"
	"go.uber.org/zap"
)

type rollingWriter struct {
	component string
	seq       int
	localPath string
	s3Key     string
	file      *os.File
	writer    *parquet.GenericWriter[LogRow]
	batch     []LogRow
}

// Manager writes Parquet parts per (hour, component, cluster) with size-based rolling.
type Manager struct {
	info        types.ClusterInfo
	hourBegin   int64
	baseDir     string
	maxBytes    int64
	writeBatch  int
	writers     map[string]*rollingWriter
	closedParts []*rollingWriter
}

func NewManager(info types.ClusterInfo, hourBegin int64) (*Manager, error) {
	baseDir := filepath.Join(
		config.LocalMergedFileDir,
		info.StorageID,
		"parquet",
		time.Unix(hourBegin, 0).UTC().Format("2006010215"),
	)
	if err := os.MkdirAll(baseDir, 0o755); err != nil {
		return nil, err
	}
	return &Manager{
		info:       info,
		hourBegin:  hourBegin,
		baseDir:    baseDir,
		maxBytes:   config.ParquetMaxFileBytes,
		writeBatch: config.ParquetWriteBatchSize,
		writers:    map[string]*rollingWriter{},
	}, nil
}

func (m *Manager) Write(component string, row LogRow) error {
	key := fmt.Sprintf("%s|%s", component, m.info.ClusterID)
	w, ok := m.writers[key]
	if !ok {
		var err error
		w, err = m.openWriter(component, 1)
		if err != nil {
			return err
		}
		m.writers[key] = w
	}

	w.batch = append(w.batch, row)
	if len(w.batch) >= m.writeBatch {
		if err := m.flushBatch(w); err != nil {
			return err
		}
	}
	size, err := w.file.Stat()
	if err != nil {
		return err
	}
	if size.Size() < m.maxBytes {
		return nil
	}
	if err := m.closeWriter(w); err != nil {
		return err
	}
	m.closedParts = append(m.closedParts, w)
	next, err := m.openWriter(component, w.seq+1)
	if err != nil {
		return err
	}
	m.writers[key] = next
	return nil
}

func (m *Manager) flushBatch(w *rollingWriter) error {
	if len(w.batch) == 0 {
		return nil
	}
	if _, err := w.writer.Write(w.batch); err != nil {
		return err
	}
	w.batch = w.batch[:0]
	return nil
}

func (m *Manager) openWriter(component string, seq int) (*rollingWriter, error) {
	t := time.Unix(m.hourBegin, 0).UTC()
	hour := t.Format("15")
	filename := fmt.Sprintf("%s_%s_%02d.parquet", hour, m.info.ClusterID, seq)
	// Local temp dir is per storage/hour; isolate by component to avoid name clashes.
	localDir := filepath.Join(m.baseDir, component)
	if err := os.MkdirAll(localDir, 0o755); err != nil {
		return nil, err
	}
	localPath := filepath.Join(localDir, filename)
	file, err := os.Create(localPath)
	if err != nil {
		return nil, err
	}
	writer := parquet.NewGenericWriter[LogRow](
		file,
		parquet.Compression(&parquet.Zstd),
		parquet.MaxRowsPerRowGroup(config.ParquetMaxRowsPerRowGroup),
	)
	return &rollingWriter{
		component: component,
		seq:       seq,
		localPath: localPath,
		s3Key:     S3ObjectKey(m.info, m.hourBegin, component, seq),
		file:      file,
		writer:    writer,
		batch:     make([]LogRow, 0, m.writeBatch),
	}, nil
}

func (m *Manager) closeWriter(w *rollingWriter) error {
	if err := m.flushBatch(w); err != nil {
		return err
	}
	if w.writer != nil {
		if err := w.writer.Close(); err != nil {
			return err
		}
		w.writer = nil
	}
	if w.file != nil {
		if err := w.file.Close(); err != nil {
			return err
		}
		w.file = nil
	}
	return nil
}

// UploadClosedParts uploads parquet parts that were rolled due to size limits.
// Call this periodically during hourly processing to release disk and memory
// instead of waiting until the end of the hour.
func (m *Manager) UploadClosedParts(ctx context.Context, store storage.Storage) error {
	if len(m.closedParts) == 0 {
		return nil
	}
	bucket := config.BucketName()
	for _, w := range m.closedParts {
		if err := m.uploadPart(ctx, store, bucket, w); err != nil {
			return err
		}
	}
	m.closedParts = nil
	return nil
}

func (m *Manager) FlushAndUpload(ctx context.Context, store storage.Storage) error {
	if err := m.UploadClosedParts(ctx, store); err != nil {
		return err
	}
	bucket := config.BucketName()
	parts := make([]*rollingWriter, 0, len(m.writers))
	for _, w := range m.writers {
		parts = append(parts, w)
		if err := m.closeWriter(w); err != nil {
			return err
		}
	}
	m.writers = map[string]*rollingWriter{}
	for _, w := range parts {
		if err := m.uploadPart(ctx, store, bucket, w); err != nil {
			return err
		}
	}
	parent := filepath.Dir(m.baseDir)
	if parent != "" && parent != "." {
		_ = os.RemoveAll(parent)
	}
	return nil
}

func (m *Manager) uploadPart(ctx context.Context, store storage.Storage, bucket string, w *rollingWriter) error {
	info, err := os.Stat(w.localPath)
	if err != nil {
		if os.IsNotExist(err) {
			return nil
		}
		return err
	}
	if info.Size() == 0 {
		return os.Remove(w.localPath)
	}
	start := time.Now()
	f, err := os.Open(w.localPath)
	if err != nil {
		return err
	}
	defer f.Close()
	if err := store.SaveObject(ctx, bucket, w.s3Key, f); err != nil {
		return err
	}
	metrics.RecordParquetUpload(m.info, w.component, info.Size(), time.Since(start).Seconds())
	log.Info("parquet uploaded",
		zap.String("cluster", m.info.StorageID),
		zap.String("key", w.s3Key),
		zap.Int64("size", info.Size()))
	return os.Remove(w.localPath)
}