package graphbuilder

import (
	"encoding/json"
	"io"
	"os"
	"path/filepath"
	"sort"
	"sync"
	"time"
)

const benchmarkReportSchema = "agentteam-insight.graph.benchmark.v1"

type benchmarkPhase struct {
	Name       string `json:"name"`
	DurationMS int64  `json:"durationMs"`
}

type benchmarkBatch struct {
	BatchIndex    int    `json:"batchIndex"`
	Files         int    `json:"files"`
	Nodes         int    `json:"nodes"`
	Edges         int    `json:"edges"`
	InputBytes    int64  `json:"inputBytes"`
	OutputBytes   int64  `json:"outputBytes"`
	PrepareMS     int64  `json:"prepareMs"`
	OpenCodeMS    int64  `json:"openCodeMs"`
	ReadMS        int64  `json:"readMs"`
	DecodeMS      int64  `json:"decodeMs"`
	ApplyMS       int64  `json:"applyMs"`
	WriteMS       int64  `json:"writeMs"`
	TotalMS       int64  `json:"totalMs"`
	Success       bool   `json:"success"`
	Attempts      int    `json:"attempts"`
	Retried       bool   `json:"retried"`
	FailureReason string `json:"failureReason,omitempty"`
	started       time.Time
	openCodeAt    time.Time
}

type benchmarkSemanticCall struct {
	Stage       string `json:"stage"`
	Attempt     int    `json:"attempt"`
	InputBytes  int    `json:"inputBytes"`
	OutputBytes int    `json:"outputBytes"`
	DurationMS  int64  `json:"durationMs"`
	Success     bool   `json:"success"`
}

type benchmarkDistribution struct {
	Count int   `json:"count"`
	P50MS int64 `json:"p50Ms"`
	P95MS int64 `json:"p95Ms"`
	MaxMS int64 `json:"maxMs"`
}

type benchmarkSummary struct {
	BatchTotal    benchmarkDistribution `json:"batchTotal"`
	BatchPrepare  benchmarkDistribution `json:"batchPrepare"`
	BatchOpenCode benchmarkDistribution `json:"batchOpenCode"`
	BatchApply    benchmarkDistribution `json:"batchApply"`
}

type benchmarkReport struct {
	Schema   string                  `json:"schema"`
	Success  bool                    `json:"success"`
	TotalMS  int64                   `json:"totalMs"`
	Phases   []benchmarkPhase        `json:"phases"`
	Batches  []benchmarkBatch        `json:"batches"`
	Semantic []benchmarkSemanticCall `json:"semantic"`
	Summary  benchmarkSummary        `json:"summary"`
}

// benchmarkRecorder owns diagnostic timing for graph orchestration. It is
// inactive unless BenchmarkPath is explicitly set, so normal runs pay only
// nil checks and do not emit per-batch log noise.
type benchmarkRecorder struct {
	mu       sync.Mutex
	started  time.Time
	phases   []benchmarkPhase
	batches  map[int]benchmarkBatch
	semantic []benchmarkSemanticCall
}

func newBenchmarkRecorder(enabled bool) *benchmarkRecorder {
	if !enabled {
		return nil
	}
	return &benchmarkRecorder{started: time.Now(), batches: map[int]benchmarkBatch{}}
}

func (recorder *benchmarkRecorder) phase(name string, started time.Time) {
	if recorder == nil {
		return
	}
	recorder.mu.Lock()
	recorder.phases = append(recorder.phases, benchmarkPhase{Name: name, DurationMS: elapsedMS(started)})
	recorder.mu.Unlock()
}

func (recorder *benchmarkRecorder) startBatch(index, files, nodes, edges int) {
	if recorder == nil {
		return
	}
	recorder.mu.Lock()
	recorder.batches[index] = benchmarkBatch{BatchIndex: index, Files: files, Nodes: nodes, Edges: edges, started: time.Now()}
	recorder.mu.Unlock()
}

func (recorder *benchmarkRecorder) startOpenCode(index int) {
	recorder.updateBatch(index, func(batch *benchmarkBatch) { batch.openCodeAt = time.Now() })
}

func (recorder *benchmarkRecorder) updateBatch(index int, update func(*benchmarkBatch)) {
	if recorder == nil {
		return
	}
	recorder.mu.Lock()
	batch := recorder.batches[index]
	update(&batch)
	recorder.batches[index] = batch
	recorder.mu.Unlock()
}

func (recorder *benchmarkRecorder) semanticCall(call benchmarkSemanticCall) {
	if recorder == nil {
		return
	}
	recorder.mu.Lock()
	recorder.semantic = append(recorder.semantic, call)
	recorder.mu.Unlock()
}

func (recorder *benchmarkRecorder) write(writer io.Writer, success bool) {
	if recorder == nil || writer == nil {
		return
	}
	recorder.mu.Lock()
	report := benchmarkReport{
		Schema: benchmarkReportSchema, Success: success, TotalMS: elapsedMS(recorder.started),
		Phases:   append([]benchmarkPhase(nil), recorder.phases...),
		Semantic: append([]benchmarkSemanticCall(nil), recorder.semantic...),
	}
	for _, batch := range recorder.batches {
		if batch.TotalMS == 0 && !batch.started.IsZero() {
			batch.TotalMS = time.Since(batch.started).Milliseconds()
		}
		if batch.OpenCodeMS == 0 && !batch.openCodeAt.IsZero() {
			batch.OpenCodeMS = time.Since(batch.openCodeAt).Milliseconds()
		}
		report.Batches = append(report.Batches, batch)
	}
	recorder.mu.Unlock()
	sort.Slice(report.Batches, func(i, j int) bool { return report.Batches[i].BatchIndex < report.Batches[j].BatchIndex })
	report.Summary = summarizeBenchmark(report.Batches)
	encoder := json.NewEncoder(writer)
	encoder.SetIndent("", "  ")
	_ = encoder.Encode(report)
}

func (recorder *benchmarkRecorder) writeFile(path string, success bool) {
	if recorder == nil || path == "" {
		return
	}
	_ = os.MkdirAll(filepath.Dir(path), 0o755)
	file, err := os.Create(path)
	if err != nil {
		return
	}
	defer file.Close()
	recorder.write(file, success)
}

func summarizeBenchmark(batches []benchmarkBatch) benchmarkSummary {
	total := make([]int64, 0, len(batches))
	prepare := make([]int64, 0, len(batches))
	opencode := make([]int64, 0, len(batches))
	apply := make([]int64, 0, len(batches))
	for _, batch := range batches {
		total = append(total, batch.TotalMS)
		prepare = append(prepare, batch.PrepareMS)
		opencode = append(opencode, batch.OpenCodeMS)
		apply = append(apply, batch.ApplyMS)
	}
	return benchmarkSummary{
		BatchTotal: distribution(total), BatchPrepare: distribution(prepare),
		BatchOpenCode: distribution(opencode), BatchApply: distribution(apply),
	}
}

func distribution(values []int64) benchmarkDistribution {
	if len(values) == 0 {
		return benchmarkDistribution{}
	}
	values = append([]int64(nil), values...)
	sort.Slice(values, func(i, j int) bool { return values[i] < values[j] })
	return benchmarkDistribution{
		Count: len(values), P50MS: percentile(values, 0.50), P95MS: percentile(values, 0.95), MaxMS: values[len(values)-1],
	}
}

func percentile(sorted []int64, percentile float64) int64 {
	index := int(float64(len(sorted)-1)*percentile + 0.5)
	return sorted[index]
}

func elapsedMS(started time.Time) int64 {
	return time.Since(started).Milliseconds()
}