package service

import (
	"context"
	"fmt"
	"runtime"
	"sort"
	"sync"
	"time"

	coremodel "security-center-backend/internal/model"
	"security-center-backend/internal/modules/caasm/collector"
	caasmmodel "security-center-backend/internal/modules/caasm/model"
	"security-center-backend/internal/modules/caasm/repository"
)

type collectorResumeSetter interface {
	SetResumeCursor(cursor *collector.ResumeCursor)
}

const chunkProgressFlushInterval = 10
const chunkProgressFlushMaxAge = 30 * time.Second
const chunkPersistQueueWaitPollInterval = 50 * time.Millisecond
const chunkPersistQueueCapacity = 4

type chunkPersistRequest struct {
	ctx   context.Context
	chunk *collector.AssetChunk
}

type chunkedSyncSink struct {
	service                *CloudAccountService
	job                    *caasmmodel.SyncJob
	account                *caasmmodel.CloudAccount
	resumeSource           uint64
	windowStart            time.Time
	cursor                 collector.ResumeCursor
	mu                     sync.Mutex
	pruned                 bool
	chunkCount             int
	assetCount             int
	ipAssetCount           int
	domainCount            int
	portCount              int
	maxHeapAlloc           uint64
	queueDepthMax          int
	writerCh               chan chunkPersistRequest
	writerDone             chan struct{}
	persistCtx             context.Context
	persistCancel          context.CancelFunc
	lastProgressFlushAt    time.Time
	lastProgressFlushCount int
	writerErrMu            sync.Mutex
	writerErr              error
	closeOnce              sync.Once
}

func newChunkedSyncSink(service *CloudAccountService, job *caasmmodel.SyncJob, account *caasmmodel.CloudAccount, executionCtx context.Context, cursor collector.ResumeCursor, resumeSource uint64, windowStart time.Time) *chunkedSyncSink {
	persistCtx, persistCancel := newChunkPersistContext(executionCtx)
	sink := &chunkedSyncSink{
		service:       service,
		job:           job,
		account:       account,
		cursor:        normalizeResumeCursor(cursor, windowStart),
		resumeSource:  resumeSource,
		windowStart:   windowStart,
		writerCh:      make(chan chunkPersistRequest, chunkPersistQueueCapacity),
		writerDone:    make(chan struct{}),
		persistCtx:    persistCtx,
		persistCancel: persistCancel,
	}
	go sink.runPersistWriter()
	return sink
}

func (s *chunkedSyncSink) PersistChunk(ctx context.Context, chunk *collector.AssetChunk) error {
	if chunk == nil {
		return nil
	}
	persistCtx := s.persistContext()
	if err := s.service.blockIfKillSwitchEnabled(persistCtx, s.job, s.account.ID, "chunk_persist"); err != nil {
		return err
	}
	if err := s.persistStopErr(); err != nil {
		return err
	}
	if err := s.getWriterErr(); err != nil {
		return err
	}

	request := chunkPersistRequest{ctx: persistCtx, chunk: chunk}
	// Prefer handing already-collected chunks to the bounded writer when capacity
	// is immediately available. If the queue is full, keep waiting on the durable
	// persist context instead of the collection call-site ctx so a fully collected
	// chunk is not discarded only because the caller context ended.
	select {
	case s.writerCh <- request:
		s.recordQueueDepth()
		if err := s.getWriterErr(); err != nil {
			return err
		}
		return nil
	default:
	}

	ticker := time.NewTicker(chunkPersistQueueWaitPollInterval)
	defer ticker.Stop()

	for {
		if err := s.persistStopErr(); err != nil {
			return err
		}
		select {
		case s.writerCh <- request:
			s.recordQueueDepth()
			if err := s.getWriterErr(); err != nil {
				return err
			}
			return nil
		case <-persistCtx.Done():
			return persistCtx.Err()
		case <-s.writerDone:
			if err := s.getWriterErr(); err != nil {
				return err
			}
			return context.Canceled
		case <-ticker.C:
		}
	}
}

func (s *chunkedSyncSink) Wait() error {
	s.closeOnce.Do(func() {
		close(s.writerCh)
	})
	<-s.writerDone
	if s.persistCancel != nil {
		s.persistCancel()
	}
	return s.getWriterErr()
}

func (s *chunkedSyncSink) runPersistWriter() {
	defer close(s.writerDone)
	for request := range s.writerCh {
		if request.chunk == nil {
			continue
		}
		if s.getWriterErr() != nil {
			continue
		}
		if err := s.persistChunk(request.ctx, request.chunk); err != nil {
			s.setWriterErr(err)
		}
	}
}

func (s *chunkedSyncSink) persistChunk(ctx context.Context, chunk *collector.AssetChunk) error {
	var (
		progressValues coremodel.JSONB
		progressDue    bool
		pruneEventDue  bool
		checkpointDue  bool
		checkpointVals coremodel.JSONB
		eventMetadata  map[string]interface{}
		eventDue       bool
		shouldPrune    bool
	)

	if err := s.persistStopErr(); err != nil {
		return err
	}
	if err := s.service.blockIfKillSwitchEnabled(ctx, s.job, s.account.ID, "chunk_persist"); err != nil {
		return err
	}
	payload := &caasmmodel.IngestPayload{
		SyncJobID:      s.job.ID,
		CloudAccountID: s.account.ID,
		Assets:         chunk.Result.Assets,
		IPAssets:       chunk.Result.IPAssets,
		DomainAssets:   chunk.Result.DomainAssets,
		Ports:          chunk.Result.Ports,
	}
	s.service.filterSyncPayloadByCollectionScope(ctx, payload)

	s.mu.Lock()
	if !s.pruned {
		s.pruned = true
		shouldPrune = true
	}
	s.mu.Unlock()

	if shouldPrune {
		if err := s.service.prunePersistedDataByCollectionScope(ctx, s.account.ID); err != nil {
			s.mu.Lock()
			s.pruned = false
			s.mu.Unlock()
			return err
		}
		pruneEventDue = true
	}

	if err := s.service.persistSyncPayload(ctx, s.job, payload); err != nil {
		return err
	}

	s.mu.Lock()
	s.chunkCount++
	s.assetCount += len(payload.Assets)
	s.ipAssetCount += len(payload.IPAssets)
	s.domainCount += len(payload.DomainAssets)
	s.portCount += len(payload.Ports)
	s.observeMemory()
	s.advanceCheckpoint(chunk.Checkpoint)
	checkpointDue = chunk.Checkpoint.Eligible
	if checkpointDue {
		checkpointVals = s.checkpointMetadata()
	}
	progressDue = s.shouldFlushProgress(chunk)
	if progressDue {
		progressValues = s.progressMetadata()
		progressValues["last_chunk_scope"] = chunk.Scope
		progressValues["last_chunk_target"] = chunk.Target
		progressValues["last_chunk_sequence"] = chunk.Sequence
		s.lastProgressFlushCount = s.chunkCount
		s.lastProgressFlushAt = time.Now()
	}
	eventDue = s.shouldRecordChunkEvent(chunk)
	if eventDue {
		eventMetadata = map[string]interface{}{
			"scope":               chunk.Scope,
			"target":              chunk.Target,
			"sequence":            chunk.Sequence,
			"asset_count":         len(payload.Assets),
			"ip_asset_count":      len(payload.IPAssets),
			"domain_asset_count":  len(payload.DomainAssets),
			"port_count":          len(payload.Ports),
			"chunk_persist_count": s.chunkCount,
			"resume_source_job":   s.resumeSource,
		}
		if registry, ok := chunk.Metadata["coverage_registry"]; ok {
			eventMetadata["coverage_registry"] = registry
		}
		if coverage, ok := chunk.Metadata["coverage"]; ok {
			eventMetadata["coverage"] = coverage
		}
	}
	s.mu.Unlock()

	if checkpointDue && len(checkpointVals) > 0 {
		if err := s.service.updateSyncJobMetadata(ctx, s.job, checkpointVals); err != nil {
			return fmt.Errorf("persist checkpoint metadata for chunk %s/%s#%d: %w", chunk.Scope, chunk.Target, chunk.Sequence, err)
		}
	}
	if progressDue {
		s.service.updateSyncJobMetadata(ctx, s.job, progressValues)
	}
	if pruneEventDue {
		s.service.recordJobEvent(ctx, s.job.ID, "info", "Chunk sync prune completed", map[string]interface{}{
			"cloud_account_id": s.account.ID,
		})
	}
	if eventDue {
		s.service.recordJobEvent(ctx, s.job.ID, "info", "Chunk persist finished", eventMetadata)
	}
	return nil
}

func (s *chunkedSyncSink) getWriterErr() error {
	s.writerErrMu.Lock()
	defer s.writerErrMu.Unlock()
	return s.writerErr
}

func (s *chunkedSyncSink) setWriterErr(err error) {
	if err == nil {
		return
	}
	s.writerErrMu.Lock()
	defer s.writerErrMu.Unlock()
	if s.writerErr == nil {
		s.writerErr = err
	}
}

func (s *chunkedSyncSink) recordQueueDepth() {
	s.mu.Lock()
	defer s.mu.Unlock()
	if depth := len(s.writerCh); depth > s.queueDepthMax {
		s.queueDepthMax = depth
	}
}

func newChunkPersistContext(executionCtx context.Context) (context.Context, context.CancelFunc) {
	if executionCtx == nil {
		return context.WithCancel(context.Background())
	}
	return context.WithCancel(executionCtx)
}

func (s *chunkedSyncSink) persistContext() context.Context {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.persistCtx
}

func (s *chunkedSyncSink) persistStopErr() error {
	if s == nil {
		return nil
	}
	if s.persistCtx != nil {
		if err := s.persistCtx.Err(); err != nil {
			return err
		}
	}
	if s.service != nil && s.job != nil && s.job.ID != 0 && s.service.wasJobStoppedExternally(s.job.ID) {
		return context.Canceled
	}
	return nil
}

func (s *chunkedSyncSink) progressMetadata() coremodel.JSONB {
	chunkTotals := coremodel.JSONB{
		"assets":        s.assetCount,
		"ip_assets":     s.ipAssetCount,
		"domain_assets": s.domainCount,
		"ports":         s.portCount,
	}

	values := coremodel.JSONB{
		"chunk_persist_count":           s.chunkCount,
		"chunk_totals":                  chunkTotals,
		"chunk_max_heap_alloc":          s.maxHeapAlloc,
		"chunk_persist_queue_capacity":  chunkPersistQueueCapacity,
		"chunk_persist_queue_depth_max": s.queueDepthMax,
		"full_run_completed":            false,
		"stale_mark_eligible":           false,
	}
	return values
}

func (s *chunkedSyncSink) checkpointMetadata() coremodel.JSONB {
	values := coremodel.JSONB{
		"resume_cursor":         resumeCursorJSONB(s.cursor),
		"resume_window_started": s.windowStart.Format(time.RFC3339Nano),
	}
	if s.resumeSource != 0 {
		values["resume_source_job_id"] = s.resumeSource
	}
	return values
}

func (s *chunkedSyncSink) shouldFlushProgress(chunk *collector.AssetChunk) bool {
	if chunk == nil {
		return false
	}
	if s.chunkCount <= 1 {
		return true
	}
	if s.chunkCount%chunkProgressFlushInterval == 0 {
		return true
	}
	if s.lastProgressFlushCount >= s.chunkCount {
		return false
	}
	if s.lastProgressFlushAt.IsZero() {
		return false
	}
	return time.Since(s.lastProgressFlushAt) >= chunkProgressFlushMaxAge
}

func (s *chunkedSyncSink) shouldRecordChunkEvent(chunk *collector.AssetChunk) bool {
	return s.shouldFlushProgress(chunk)
}

func (s *chunkedSyncSink) observeMemory() {
	var stats runtime.MemStats
	runtime.ReadMemStats(&stats)
	if stats.HeapAlloc > s.maxHeapAlloc {
		s.maxHeapAlloc = stats.HeapAlloc
	}
}

func (s *chunkedSyncSink) advanceCheckpoint(checkpoint collector.ChunkCheckpoint) {
	if !checkpoint.Eligible {
		return
	}
	if checkpoint.GlobalService != "" {
		s.cursor.CompletedGlobal = appendUniqueString(s.cursor.CompletedGlobal, checkpoint.GlobalService)
	}
	if checkpoint.Region != "" {
		s.cursor.CompletedRegions = appendUniqueString(s.cursor.CompletedRegions, checkpoint.Region)
	}
	if checkpoint.Route53ZoneID != "" {
		if s.cursor.Route53ZoneProgress == nil {
			s.cursor.Route53ZoneProgress = make(map[string]collector.Route53ZoneProgress)
		}
		if checkpoint.Route53ZoneComplete {
			s.cursor.Route53CompletedZones = appendUniqueString(s.cursor.Route53CompletedZones, checkpoint.Route53ZoneID)
			delete(s.cursor.Route53ZoneProgress, checkpoint.Route53ZoneID)
		} else if checkpoint.Route53NextRecordName != "" && checkpoint.Route53NextRecordType != "" {
			s.cursor.Route53ZoneProgress[checkpoint.Route53ZoneID] = collector.Route53ZoneProgress{
				ZoneID:               checkpoint.Route53ZoneID,
				ZoneName:             checkpoint.Route53Zone,
				NextRecordName:       checkpoint.Route53NextRecordName,
				NextRecordType:       checkpoint.Route53NextRecordType,
				NextRecordIdentifier: checkpoint.Route53NextRecordIdentifier,
			}
		}
	}
	sort.Strings(s.cursor.CompletedGlobal)
	sort.Strings(s.cursor.CompletedRegions)
	sort.Strings(s.cursor.Route53CompletedZones)
}

func loadResumeCursor(ctx context.Context, jobs repository.SyncJobRepository, accountID uint64, currentJobID uint64, jobType string) (collector.ResumeCursor, uint64) {
	if jobs == nil || accountID == 0 {
		return collector.ResumeCursor{}, 0
	}

	item, err := jobs.FindLatestResumableJobForAccount(ctx, accountID, currentJobID, jobType)
	if err != nil {
		return collector.ResumeCursor{}, 0
	}
	if item == nil {
		return collector.ResumeCursor{}, 0
	}
	cursor := resumeCursorFromMetadata(item.Metadata)
	if !resumeCursorHasProgress(cursor) {
		return collector.ResumeCursor{}, 0
	}
	return cursor, item.ID
}

func normalizeResumeCursor(cursor collector.ResumeCursor, windowStart time.Time) collector.ResumeCursor {
	if cursor.WindowStartedAt == "" {
		cursor.WindowStartedAt = windowStart.Format(time.RFC3339Nano)
	}
	cursor.CompletedGlobal = dedupeStrings(cursor.CompletedGlobal)
	cursor.CompletedRegions = dedupeStrings(cursor.CompletedRegions)
	cursor.Route53CompletedZones = dedupeStrings(cursor.Route53CompletedZones)
	cursor.Route53ZoneProgress = dedupeRoute53ZoneProgress(cursor.Route53ZoneProgress)
	return cursor
}

func resumeCursorHasProgress(cursor collector.ResumeCursor) bool {
	return cursor.WindowStartedAt != "" ||
		len(cursor.CompletedGlobal) > 0 ||
		len(cursor.CompletedRegions) > 0 ||
		len(cursor.Route53CompletedZones) > 0 ||
		len(cursor.Route53ZoneProgress) > 0
}

func resumeCursorJSONB(cursor collector.ResumeCursor) coremodel.JSONB {
	return coremodel.JSONB{
		"window_started_at":       cursor.WindowStartedAt,
		"completed_global":        append([]string(nil), cursor.CompletedGlobal...),
		"completed_regions":       append([]string(nil), cursor.CompletedRegions...),
		"route53_completed_zones": append([]string(nil), cursor.Route53CompletedZones...),
		"route53_zone_progress":   route53ZoneProgressJSONB(cursor.Route53ZoneProgress),
	}
}

func resumeCursorFromMetadata(metadata coremodel.JSONB) collector.ResumeCursor {
	if metadata == nil {
		return collector.ResumeCursor{}
	}
	raw, ok := metadata["resume_cursor"]
	if !ok {
		return collector.ResumeCursor{}
	}
	if cursor, ok := raw.(collector.ResumeCursor); ok {
		return cursor
	}
	values, ok := raw.(map[string]interface{})
	if !ok {
		if jsonbValues, ok := raw.(coremodel.JSONB); ok {
			values = map[string]interface{}(jsonbValues)
		} else {
			return collector.ResumeCursor{}
		}
	}
	cursor := collector.ResumeCursor{}
	if startedAt, ok := values["window_started_at"].(string); ok {
		cursor.WindowStartedAt = startedAt
	}
	cursor.CompletedGlobal = interfaceSliceToStrings(values["completed_global"])
	cursor.CompletedRegions = interfaceSliceToStrings(values["completed_regions"])
	cursor.Route53CompletedZones = interfaceSliceToStrings(values["route53_completed_zones"])
	cursor.Route53ZoneProgress = route53ZoneProgressFromValue(values["route53_zone_progress"])
	return cursor
}

func resumeWindowStart(cursor collector.ResumeCursor, fallback time.Time) time.Time {
	if cursor.WindowStartedAt == "" {
		return fallback
	}
	parsed, err := time.Parse(time.RFC3339Nano, cursor.WindowStartedAt)
	if err != nil {
		return fallback
	}
	return parsed
}

func appendUniqueString(values []string, next string) []string {
	for _, value := range values {
		if value == next {
			return values
		}
	}
	return append(values, next)
}

func dedupeStrings(values []string) []string {
	if len(values) <= 1 {
		return values
	}
	unique := make([]string, 0, len(values))
	seen := make(map[string]struct{}, len(values))
	for _, value := range values {
		if value == "" {
			continue
		}
		if _, ok := seen[value]; ok {
			continue
		}
		seen[value] = struct{}{}
		unique = append(unique, value)
	}
	sort.Strings(unique)
	return unique
}

func interfaceSliceToStrings(value interface{}) []string {
	switch typed := value.(type) {
	case []string:
		return dedupeStrings(typed)
	case []interface{}:
		values := make([]string, 0, len(typed))
		for _, entry := range typed {
			if item, ok := entry.(string); ok {
				values = append(values, item)
			}
		}
		return dedupeStrings(values)
	default:
		return nil
	}
}

func dedupeRoute53ZoneProgress(progress map[string]collector.Route53ZoneProgress) map[string]collector.Route53ZoneProgress {
	if len(progress) == 0 {
		return nil
	}
	normalized := make(map[string]collector.Route53ZoneProgress, len(progress))
	for zoneID, cursor := range progress {
		id := zoneID
		if cursor.ZoneID != "" {
			id = cursor.ZoneID
		}
		if id == "" || cursor.NextRecordName == "" || cursor.NextRecordType == "" {
			continue
		}
		cursor.ZoneID = id
		normalized[id] = cursor
	}
	if len(normalized) == 0 {
		return nil
	}
	return normalized
}

func route53ZoneProgressJSONB(progress map[string]collector.Route53ZoneProgress) coremodel.JSONB {
	if len(progress) == 0 {
		return coremodel.JSONB{}
	}
	values := coremodel.JSONB{}
	for zoneID, cursor := range dedupeRoute53ZoneProgress(progress) {
		values[zoneID] = coremodel.JSONB{
			"zone_id":                cursor.ZoneID,
			"zone_name":              cursor.ZoneName,
			"next_record_name":       cursor.NextRecordName,
			"next_record_type":       cursor.NextRecordType,
			"next_record_identifier": cursor.NextRecordIdentifier,
		}
	}
	return values
}

func route53ZoneProgressFromValue(value interface{}) map[string]collector.Route53ZoneProgress {
	if value == nil {
		return nil
	}
	rawMap, ok := value.(map[string]interface{})
	if !ok {
		if jsonbMap, ok := value.(coremodel.JSONB); ok {
			rawMap = map[string]interface{}(jsonbMap)
		} else {
			return nil
		}
	}
	progress := make(map[string]collector.Route53ZoneProgress, len(rawMap))
	for zoneID, rawCursor := range rawMap {
		cursorMap, ok := rawCursor.(map[string]interface{})
		if !ok {
			if jsonbCursor, ok := rawCursor.(coremodel.JSONB); ok {
				cursorMap = map[string]interface{}(jsonbCursor)
			} else {
				continue
			}
		}
		cursor := collector.Route53ZoneProgress{
			ZoneID:               stringValue(cursorMap["zone_id"]),
			ZoneName:             stringValue(cursorMap["zone_name"]),
			NextRecordName:       stringValue(cursorMap["next_record_name"]),
			NextRecordType:       stringValue(cursorMap["next_record_type"]),
			NextRecordIdentifier: stringValue(cursorMap["next_record_identifier"]),
		}
		if cursor.ZoneID == "" {
			cursor.ZoneID = zoneID
		}
		progress[cursor.ZoneID] = cursor
	}
	return dedupeRoute53ZoneProgress(progress)
}

func stringValue(value interface{}) string {
	if str, ok := value.(string); ok {
		return str
	}
	return ""
}

func chunkSummaryMetadata(summary *collector.CollectionSummary, sink *chunkedSyncSink, fullRunCompleted bool) coremodel.JSONB {
	values := sink.progressMetadata()
	for key, value := range sink.checkpointMetadata() {
		values[key] = value
	}
	values["collection_completed"] = true
	values["persist_completed"] = true
	values["full_run_completed"] = fullRunCompleted
	values["stale_mark_eligible"] = fullRunCompleted
	if summary != nil {
		values["chunk_count"] = summary.ChunkCount
		values["partial_failure_count"] = len(summary.PartialFailures)
	}
	return values
}

func chunkCollectionMessage(summary *collector.CollectionSummary, sink *chunkedSyncSink) string {
	if summary == nil {
		return fmt.Sprintf("Chunked collection finished with %d persisted chunks", sink.chunkCount)
	}
	return fmt.Sprintf("Chunked collection finished with %d collector chunks and %d persisted chunks", summary.ChunkCount, sink.chunkCount)
}