package engine

import (
	"bytes"
	"context"
	"crypto/tls"
	"fmt"
	"io"
	"net/http"
	"sync"
	"time"

	"github.com/r3g1tpr0cs/burterm/internal/payload"
)

// AttackMode задаёт способ комбинирования генераторов и точек вставки,
// по аналогии с режимами Burp Intruder.
type AttackMode int

const (
	// ModeSniper — один генератор, payload'ы по очереди проходят через
	// каждую точку вставки (остальные точки остаются пустыми/базовыми).
	ModeSniper AttackMode = iota
	// ModeBattering — один генератор, payload вставляется во все точки
	// одновременно (одно и то же значение везде).
	ModeBattering
	// ModePitchfork — по генератору на каждую точку вставки, но идут они
	// в ногу по индексу (payloads[i] из каждого генератора вместе
	// формируют один запрос), а не декартовым произведением. Останавливается
	// на самом коротком генераторе — как в Burp: если один список короче
	// остальных, лишние значения из более длинных просто не используются.
	ModePitchfork
	// ModeClusterBomb — разные генераторы на разные позиции,
	// декартово произведение всех комбинаций.
	ModeClusterBomb
)

// Result — результат одного запроса атаки.
type Result struct {
	Payload []string

	// URL и Method — куда именно ушёл этот конкретный запрос атаки.
	// Заполняются в send() из уже собранного *http.Request — нужны,
	// в частности, чтобы результаты Intruder-атаки можно было сохранить
	// в project.Store (там для каждой записи нужны host/method/url,
	// а Result сам по себе их не содержал до этого момента).
	URL    string
	Method string

	StatusCode int
	// Length — фактическое количество прочитанных байт тела ответа,
	// а не значение заголовка Content-Length (оно часто равно -1 при
	// chunked encoding и вводило бы в заблуждение как база для фильтра
	// по длине).
	Length int64
	// Body — тело ответа, обрезанное до maxBodyCapture байт. Полное
	// тело не хранится намеренно: при тысячах запросов в атаке это
	// быстро съело бы память; для grep-фильтра обрезанного фрагмента
	// обычно достаточно, чтобы поймать характерную строку в ответе.
	Body string
	// Headers — заголовки ответа (используются, в частности, для
	// цветовой маркировки по Content-Type в TUI).
	Headers http.Header

	Duration time.Duration
	Err      error
}

// maxBodyCapture — сколько байт тела ответа сохраняется в Result.Body.
const maxBodyCapture = 4096

// Config описывает параметры выполнения атаки.
type Config struct {
	Concurrency        int
	RateLimit          time.Duration // пауза между запросами на воркер
	Timeout            time.Duration
	InsecureSkipVerify bool

	// ForceScheme принудительно проставляет схему ("http" или "https")
	// для каждого запроса атаки, перекрывая дефолт из BuildRequest
	// (который всегда ставит https, если в запросе нет абсолютного URI).
	// Пустая строка — оставить как есть (дефолтное поведение BuildRequest).
	ForceScheme string
}

// Engine выполняет HTTP-атаку с заданной конкурентностью.
type Engine struct {
	client *http.Client
	cfg    Config
}

// New создаёт Engine с настроенным http.Client (таймаут, TLS).
//
// MaxIdleConnsPerHost/MaxIdleConns заданы явно — у голого http.Transport
// дефолт MaxIdleConnsPerHost всего 2 (значение из net/http.DefaultTransport),
// а Intruder-атака почти всегда бьёт в один и тот же хост множеством
// конкурентных воркеров. С дефолтом это означает, что при Concurrency
// больше двух подавляющее большинство запросов не переиспользует
// соединение из пула — TCP- (а для HTTPS ещё и TLS-) хендшейк на каждый
// запрос вместо одного на воркер. Ставим пул размером с Concurrency —
// ровно столько одновременных соединений к одному хосту нам и нужно.
func New(cfg Config) *Engine {
	transport := &http.Transport{
		TLSClientConfig:     &tls.Config{InsecureSkipVerify: cfg.InsecureSkipVerify},
		MaxIdleConnsPerHost: cfg.Concurrency,
		MaxIdleConns:        cfg.Concurrency * 2, // с запасом на случай редиректов на другой хост
		IdleConnTimeout:     90 * time.Second,
	}
	return &Engine{
		client: &http.Client{Timeout: cfg.Timeout, Transport: transport},
		cfg:    cfg,
	}
}

// EstimateTotal вычисляет ожидаемое общее число запросов атаки заранее,
// без её запуска — используется в TUI для индикации прогресса (N/Total).
// Опирается только на payload.Generator.Len() каждого генератора, не
// потребляя их состояние (Len() не двигает Next()/Reset()), поэтому
// безопасно вызывать перед реальным Run с теми же генераторами.
func EstimateTotal(mode AttackMode, points []InsertionPoint, gens []payload.Generator) int {
	switch mode {
	case ModeSniper:
		if len(gens) == 0 {
			return 0
		}
		// RunSniper проходит генератор заново (Reset) на каждую точку —
		// то есть суммарно len(points) полных проходов одного генератора.
		return len(points) * gens[0].Len()

	case ModeBattering:
		if len(gens) == 0 {
			return 0
		}
		return gens[0].Len()

	case ModePitchfork:
		if len(gens) == 0 {
			return 0
		}
		min := -1
		for _, g := range gens {
			l := g.Len()
			if min == -1 || l < min {
				min = l
			}
		}
		if min < 0 {
			return 0
		}
		return min

	case ModeClusterBomb:
		if len(gens) == 0 {
			return 0
		}
		total := 1
		for _, g := range gens {
			total *= g.Len()
		}
		return total

	default:
		return 0
	}
}

// Run — единая точка входа для запуска атаки в любом из режимов.
// Сигнатура унифицирована под все три режима через gens []payload.Generator:
//
//   - ModeSniper: используется только gens[0], остальные позиции на каждом
//     запросе берутся из base (если base == nil, они остаются пустыми
//     строками). Если позиций несколько, base обязателен и должен быть
//     той же длины, что и points.
//   - ModeBattering: используется только gens[0], его значение вставляется
//     сразу во все точки.
//   - ModePitchfork: используются все gens, len(gens) должен равняться
//     len(points) — по генератору на каждую точку, но идут они в ногу
//     по индексу, не декартовым произведением.
//   - ModeClusterBomb: используются все gens, len(gens) должен равняться
//     len(points) — по генератору на каждую точку.
//
// Так TUI/CLI могут работать с одной сигнатурой, не зная деталей
// каждого режима.
func (e *Engine) Run(ctx context.Context, mode AttackMode, clean string, points []InsertionPoint, gens []payload.Generator, base []string) (<-chan Result, error) {
	if len(gens) == 0 {
		return nil, fmt.Errorf("нужен хотя бы один генератор payload'ов")
	}

	switch mode {
	case ModeSniper:
		b := base
		if b == nil {
			b = make([]string, len(points))
		}
		if len(b) != len(points) {
			return nil, fmt.Errorf("base должен быть длины %d (по числу точек вставки), получено %d", len(points), len(b))
		}
		return e.RunSniper(ctx, clean, points, gens[0], b), nil

	case ModeBattering:
		return e.RunBattering(ctx, clean, points, gens[0]), nil

	case ModePitchfork:
		return e.RunPitchfork(ctx, clean, points, gens)

	case ModeClusterBomb:
		return e.RunClusterBomb(ctx, clean, points, gens)

	default:
		return nil, fmt.Errorf("неизвестный режим атаки: %v", mode)
	}
}

// readBodyCapped дочитывает r до конца (это важно само по себе — см.
// комментарий в send про keep-alive), считая полную длину total, но
// сохраняя в captured не больше cap байт. Ошибку чтения намеренно не
// прокидывает наверх как фатальную для результата атаки: то, что тело
// оборвалось на середине, не должно скрывать уже полученные статус-код
// и частично прочитанные данные.
func readBodyCapped(r io.Reader, cap int) (captured []byte, total int64, err error) {
	buf := make([]byte, 32*1024)
	var b bytes.Buffer
	for {
		n, rerr := r.Read(buf)
		if n > 0 {
			total += int64(n)
			if b.Len() < cap {
				remaining := cap - b.Len()
				if n < remaining {
					remaining = n
				}
				b.Write(buf[:remaining])
			}
		}
		if rerr == io.EOF {
			break
		}
		if rerr != nil {
			return b.Bytes(), total, rerr
		}
	}
	return b.Bytes(), total, nil
}

// RunSniper гоняет атаку в режиме Sniper: один генератор, одна точка
// вставки за проход. Результаты идут в возвращаемый канал.
func (e *Engine) RunSniper(ctx context.Context, clean string, points []InsertionPoint, gen payload.Generator, base []string) <-chan Result {
	out := make(chan Result)
	type job struct {
		pointIdx int
		value    string
	}
	jobs := make(chan job)

	var wg sync.WaitGroup
	for i := 0; i < e.cfg.Concurrency; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for j := range jobs {
				payloads := append([]string(nil), base...)
				payloads[j.pointIdx] = j.value
				e.send(ctx, clean, points, payloads, out)
			}
		}()
	}

	go func() {
		defer close(jobs)
		for pointIdx := range points {
			gen.Reset()
			for {
				v, ok := gen.Next()
				if !ok {
					break
				}
				select {
				case jobs <- job{pointIdx: pointIdx, value: v}:
				case <-ctx.Done():
					return
				}
			}
		}
	}()

	go func() {
		wg.Wait()
		close(out)
	}()

	return out
}

// send выполняет один запрос с заданным набором payload'ов и пишет
// результат в out. Обе потенциально блокирующие операции — отправка в out
// и пауза rate limit'а — слушают ctx.Done(), поэтому отмена контекста
// гарантированно освобождает воркер, даже если читатель канала (TUI)
// уже перестал вычитывать результаты этой атаки.
func (e *Engine) send(ctx context.Context, clean string, points []InsertionPoint, payloads []string, out chan<- Result) {
	start := time.Now()
	req, err := BuildRequest(clean, points, payloads)
	if err != nil {
		select {
		case out <- Result{Payload: payloads, Err: err}:
		case <-ctx.Done():
		}
		return
	}
	if e.cfg.ForceScheme != "" {
		req.URL.Scheme = e.cfg.ForceScheme
	}
	req = req.WithContext(ctx)

	resp, err := e.client.Do(req)
	res := Result{Payload: payloads, Duration: time.Since(start), Err: err, URL: req.URL.String(), Method: req.Method}
	if err == nil {
		res.StatusCode = resp.StatusCode
		res.Headers = resp.Header
		body, total, _ := readBodyCapped(resp.Body, maxBodyCapture)
		resp.Body.Close()
		res.Body = string(body)
		res.Length = total
	}

	select {
	case out <- res:
	case <-ctx.Done():
		return
	}

	if e.cfg.RateLimit > 0 {
		select {
		case <-time.After(e.cfg.RateLimit):
		case <-ctx.Done():
		}
	}
}

// RunBattering гоняет атаку в режиме Battering Ram: один генератор,
// каждое его значение вставляется одновременно во все точки вставки
// (в отличие от Sniper, где точки перебираются по очереди, здесь на
// каждый запрос приходится ровно одно значение генератора, но во
// все §маркеры§ сразу — удобно, когда один и тот же токен/ID
// повторяется в нескольких местах запроса).
func (e *Engine) RunBattering(ctx context.Context, clean string, points []InsertionPoint, gen payload.Generator) <-chan Result {
	out := make(chan Result)
	jobs := make(chan string)

	var wg sync.WaitGroup
	for i := 0; i < e.cfg.Concurrency; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for v := range jobs {
				payloads := make([]string, len(points))
				for i := range payloads {
					payloads[i] = v
				}
				e.send(ctx, clean, points, payloads, out)
			}
		}()
	}

	go func() {
		defer close(jobs)
		gen.Reset()
		for {
			v, ok := gen.Next()
			if !ok {
				return
			}
			select {
			case jobs <- v:
			case <-ctx.Done():
				return
			}
		}
	}()

	go func() {
		wg.Wait()
		close(out)
	}()

	return out
}

// RunPitchfork гоняет атаку в режиме Pitchfork: у каждой точки вставки
// свой генератор (len(gens) должен равняться len(points)), но в отличие
// от Cluster Bomb они не перемножаются, а идут в ногу — i-й запрос берёт
// i-е значение из каждого генератора одновременно. Останавливается на
// самом коротком генераторе: если один список payload'ов короче
// остальных, лишние значения из более длинных просто не используются —
// то же поведение, что и в Burp Pitchfork.
func (e *Engine) RunPitchfork(ctx context.Context, clean string, points []InsertionPoint, gens []payload.Generator) (<-chan Result, error) {
	if len(gens) != len(points) {
		return nil, fmt.Errorf("ожидалось %d генераторов (по числу точек вставки), получено %d", len(points), len(gens))
	}

	values := make([][]string, len(gens))
	minLen := -1
	for i, g := range gens {
		values[i] = drainGenerator(g)
		if len(values[i]) == 0 {
			return nil, fmt.Errorf("генератор для точки %q не выдал ни одного payload'а", points[i].Name)
		}
		if minLen == -1 || len(values[i]) < minLen {
			minLen = len(values[i])
		}
	}

	out := make(chan Result)
	jobs := make(chan int) // общий индекс i в пределах minLen

	var wg sync.WaitGroup
	for i := 0; i < e.cfg.Concurrency; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for idx := range jobs {
				payloads := make([]string, len(points))
				for pos := range points {
					payloads[pos] = values[pos][idx]
				}
				e.send(ctx, clean, points, payloads, out)
			}
		}()
	}

	go func() {
		defer close(jobs)
		for i := 0; i < minLen; i++ {
			select {
			case jobs <- i:
			case <-ctx.Done():
				return
			}
		}
	}()

	go func() {
		wg.Wait()
		close(out)
	}()

	return out, nil
}

// RunClusterBomb гоняет атаку в режиме Cluster Bomb: каждой точке вставки
// сопоставлен свой генератор (len(gens) должен равняться len(points)),
// и перебираются все комбинации значений (декартово произведение).
//
// Значения каждого генератора вычитываются один раз в память
// (drainGenerator), после чего перебор комбинаций идёт по индексам —
// это ограничивает память суммой длин словарей, а не числом комбинаций,
// которое при N позициях и M вариантах растёт как M^N.
func (e *Engine) RunClusterBomb(ctx context.Context, clean string, points []InsertionPoint, gens []payload.Generator) (<-chan Result, error) {
	if len(gens) != len(points) {
		return nil, fmt.Errorf("ожидалось %d генераторов (по числу точек вставки), получено %d", len(points), len(gens))
	}

	values := make([][]string, len(gens))
	for i, g := range gens {
		values[i] = drainGenerator(g)
		if len(values[i]) == 0 {
			return nil, fmt.Errorf("генератор для точки %q не выдал ни одного payload'а", points[i].Name)
		}
	}

	out := make(chan Result)
	jobs := make(chan []int)

	var wg sync.WaitGroup
	for i := 0; i < e.cfg.Concurrency; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for idx := range jobs {
				payloads := make([]string, len(idx))
				for pos, vi := range idx {
					payloads[pos] = values[pos][vi]
				}
				e.send(ctx, clean, points, payloads, out)
			}
		}()
	}

	go func() {
		defer close(jobs)
		odometer(values, func(idx []int) bool {
			cp := append([]int(nil), idx...)
			select {
			case jobs <- cp:
				return true
			case <-ctx.Done():
				return false
			}
		})
	}()

	go func() {
		wg.Wait()
		close(out)
	}()

	return out, nil
}

// drainGenerator вычитывает генератор до конца в слайс значений.
func drainGenerator(g payload.Generator) []string {
	g.Reset()
	var out []string
	for {
		v, ok := g.Next()
		if !ok {
			break
		}
		out = append(out, v)
	}
	return out
}

// odometer перебирает все комбинации индексов по values (декартово
// произведение), вызывая emit для каждой комбинации. Работает как счётчик
// с разными основаниями по разрядам: последний разряд крутится быстрее
// всего, при переполнении перенос идёт влево — отсюда и название.
// Останавливается досрочно, если emit вернул false.
func odometer(values [][]string, emit func(idx []int) bool) {
	n := len(values)
	if n == 0 {
		return
	}
	idx := make([]int, n)

	for {
		if !emit(idx) {
			return
		}

		pos := n - 1
		for pos >= 0 {
			idx[pos]++
			if idx[pos] < len(values[pos]) {
				break
			}
			idx[pos] = 0
			pos--
		}
		if pos < 0 {
			return // все разряды переполнились — перебор завершён
		}
	}
}