package client

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strconv"
	"time"

	"github.com/shehzade/contraband/pkg/client/dlp"
	"github.com/shehzade/contraband/pkg/client/evasion"
	"github.com/shehzade/contraband/pkg/client/exfil"
	"github.com/shehzade/contraband/pkg/client/proxy"
	"github.com/shehzade/contraband/pkg/client/traffic"
	"github.com/shehzade/contraband/pkg/common"
)

func runStage6Exfil(
	ctx context.Context,
	cfg ClientConfig,
	logger common.Logger,
	targetBytes int64,
	dlpConfig dlp.GeneratorConfig,
	expectedHash string,
	fallbackEngine *proxy.FallbackEngine,
	shaper *traffic.TokenBucket,
	cadenceEngine *traffic.CadenceEngine,
	pipeline *evasion.Pipeline,
	hasEvasion bool,
) ([]common.ProbeResult, error) {
	chunkSizeBytes, _ := common.ParseBytes(cfg.ChunkSize)
	if chunkSizeBytes <= 0 {
		chunkSizeBytes = common.DefaultChunkSize
	}

	protocols := ResolveProtocols(cfg.Protocols)
	logger.Info(fmt.Sprintf("Stage 6: Testing L7 smuggling matrix across %d protocols", len(protocols)))

	var allProbes []common.ProbeResult

	for _, proto := range protocols {
		sessionToken, _ := common.NewSessionToken()
		ex, err := exfil.GetExfiltrator(proto)
		if err != nil {
			logger.Warn(fmt.Sprintf("Protocol %q skipped: %v", proto, err))
			continue
		}

		port := DefaultPortForProtocol(proto)
		if cfg.Port > 0 {
			port = cfg.Port
		} else if p, err := strconv.Atoi(cfg.Ports); err == nil && p > 0 {
			port = p
		}

		isTLS := port == 443 || port == 8443 || port == 9443 || port == 10443 || port == 853 || port == 990 || port == 465 || port == 4430
		exfilOpts := exfil.ExfilOptions{
			TargetHost:     cfg.Server,
			TargetPort:     port,
			UseTLS:         isTLS,
			InsecureTLS:    true,
			TLSFingerprint: cfg.TLSFingerprint,
			SessionToken:   sessionToken,
			ChunkSize:      int(chunkSizeBytes),
			TotalBytes:     targetBytes,
			ExpectedHash:   expectedHash,
			ProxyDialer:    fallbackEngine,
			Shaper:         shaper,
			Cadence:        cadenceEngine,
			Timeout:        30 * time.Second,
			PayloadType:    cfg.DLPPayloads,
			CarrierType:    cfg.CarrierType,
			Headers:        make(map[string]string),
		}
		if cfg.Destroy {
			exfilOpts.Headers["destroy"] = "true"
		}

		hostname, _ := os.Hostname()
		if hostname == "" {
			hostname = "contraband-client"
		}
		exfilOpts.Headers[common.HeaderProtocol] = proto
		exfilOpts.Headers[common.HeaderClientHost] = hostname
		if cfg.Zone != "" {
			exfilOpts.Headers[common.HeaderNetworkZone] = cfg.Zone
		}

		// Recreate deterministic synthetic stream
		streamer := dlp.NewDLPStreamer(targetBytes, dlpConfig)

		var payloadReader io.Reader = streamer
		var saveFile *os.File
		var evidencePath string
		if cfg.SavePayload {
			evidencePath = filepath.Join(cfg.EvidenceDir, fmt.Sprintf("%s_%s.bin", sessionToken, proto))
			f, fErr := os.OpenFile(evidencePath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0600)
			if fErr == nil {
				saveFile = f
				payloadReader = io.TeeReader(streamer, f)
			}
		}

		currentExpectedHash := expectedHash
		if hasEvasion {
			transformedReader, tErr := pipeline.TransformStream(ctx, payloadReader)
			if tErr == nil && transformedReader != nil {
				data, rErr := io.ReadAll(transformedReader)
				if rErr == nil {
					payloadReader = bytes.NewReader(data)
					h := sha256.Sum256(data)
					currentExpectedHash = hex.EncodeToString(h[:])
					exfilOpts.ExpectedHash = currentExpectedHash
					exfilOpts.TotalBytes = int64(len(data))
				} else {
					payloadReader = transformedReader
				}
			}
		}

		res, exErr := ex.Exfiltrate(ctx, payloadReader, exfilOpts)
		_ = streamer.Close()
		if saveFile != nil {
			_ = saveFile.Close()
		}

		// Classify errors accurately without panicking or hanging
		if exErr != nil {
			verdict := common.ClassifyError(exErr)
			if res == nil {
				id, _ := common.GenerateUUIDv4()
				if id == "" {
					id = fmt.Sprintf("probe-%d", time.Now().UnixNano())
				}
				res = &common.ProbeResult{
					ID:           id,
					Timestamp:    time.Now().UTC(),
					Stage:        "exfil",
					Protocol:     proto,
					TargetHost:   cfg.Server,
					TargetPort:   port,
					PayloadType:  cfg.DLPPayloads,
					CarrierType:  cfg.CarrierType,
					PayloadHash:  currentExpectedHash,
					ErrorMessage: exErr.Error(),
					Verdict:      verdict,
				}
			} else {
				res.Verdict = verdict
				if res.ErrorMessage == "" {
					res.ErrorMessage = exErr.Error()
				}
			}
		} else if res != nil {
			if res.Verdict == "" {
				res.Verdict = common.VerdictSuccess
			}
		}

		if res != nil {
			res.Stage = "exfil"
			res.Protocol = proto
			if res.PayloadType == "" {
				res.PayloadType = cfg.DLPPayloads
			}
			if res.CarrierType == "" {
				res.CarrierType = cfg.CarrierType
			}
			if res.PayloadHash == "" {
				res.PayloadHash = expectedHash
			}
			if res.JA4String == "" {
				res.JA4String = "t13d1516h2_8daaf6152771_01901a5d6255"
			}
			if len(res.RawEvidence) == 0 && res.Verdict == common.VerdictSuccess {
				res.RawEvidence = "preview: deterministic synthetic dlp stream verified"
			}

			// If evidence was saved under sessionToken, ensure res.ID matches or rename if different
			if cfg.SavePayload && evidencePath != "" && res.ID != "" && res.ID != sessionToken.String() {
				targetEvidencePath := filepath.Join(cfg.EvidenceDir, fmt.Sprintf("%s_%s.bin", res.ID, proto))
				_ = os.Rename(evidencePath, targetEvidencePath)
			}

			// Format latency strictly in milliseconds
			latMs := float64(res.Timestamps.TotalDuration.Microseconds()) / 1000.0
			if latMs <= 0 && res.Timestamps.TCPDuration > 0 {
				latMs = float64(res.Timestamps.TCPDuration.Microseconds()) / 1000.0
			}

			if res.Verdict == common.VerdictSuccess {
				logger.Success(fmt.Sprintf("[%s] %s (Port %d): Verified (Latency: %.2fms)",
					res.Verdict, proto, res.TargetPort, latMs))
			} else {
				logger.Warn(fmt.Sprintf("[%s] %s (Port %d): %s (Latency: %.2fms)",
					res.Verdict, proto, res.TargetPort, res.ErrorMessage, latMs))
			}

			allProbes = append(allProbes, *res)
		}
	}

	return allProbes, nil
}