package main

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"sync"
	"time"
)

type StartInfo struct {
	RunID     string    `json:"run_id"`
	Scenario  string    `json:"scenario"`
	StartAt   time.Time `json:"start_at"`
	Committed bool      `json:"committed"`
}

type AgentReady struct {
	Machine             string `json:"machine"`
	ReadyConnections    int    `json:"ready_connections"`
	ExpectedConnections int    `json:"expected_connections"`
	Armed               bool   `json:"armed"`
}

type coordinatorState struct {
	mu             sync.Mutex
	start          StartInfo
	startDelay     time.Duration
	commitDelay    time.Duration
	expectedAgents int
	ready          map[string]AgentReady
	armed          map[string]bool
	ledgers        map[string][]AcceptedMessage
	observed       map[string][]AcceptedMessage
	history        map[string][]string
	summaries      map[string]AgentSummary
	complete       chan struct{}
	completeOnce   sync.Once
}

func runCoordinator(ctx context.Context, listen, runID, scenario, output string, expectedAgents int, startDelay time.Duration) error {
	if !runIDPattern.MatchString(runID) {
		return fmt.Errorf("run-id must match %s", runIDPattern.String())
	}
	if !validScenario(scenario) {
		return fmt.Errorf("unsupported scenario %q", scenario)
	}
	if expectedAgents <= 0 {
		expectedAgents = 2
	}
	state := &coordinatorState{
		start: StartInfo{RunID: runID, Scenario: scenario}, startDelay: startDelay, commitDelay: 5 * time.Second,
		expectedAgents: expectedAgents,
		ready:          make(map[string]AgentReady), ledgers: make(map[string][]AcceptedMessage),
		observed: make(map[string][]AcceptedMessage), history: make(map[string][]string),
		armed: make(map[string]bool), summaries: make(map[string]AgentSummary), complete: make(chan struct{}),
	}
	mux := http.NewServeMux()
	mux.HandleFunc("/v1/ready", state.handleReady)
	mux.HandleFunc("/v1/start", state.handleStart)
	mux.HandleFunc("/v1/ledger", state.handleLedger)
	mux.HandleFunc("/v1/observed", state.handleObserved)
	mux.HandleFunc("/v1/history", state.handleHistory)
	mux.HandleFunc("/v1/summary", state.handleSummary)
	mux.HandleFunc("/v1/status", state.handleStatus)
	server := &http.Server{Addr: listen, Handler: mux, ReadHeaderTimeout: 5 * time.Second}
	errCh := make(chan error, 1)
	go func() {
		log.Printf("coordinator listening=%s run_id=%s scenario=%s expected_agents=%d", listen, runID, scenario, expectedAgents)
		if err := server.ListenAndServe(); err != nil && err != http.ErrServerClosed {
			errCh <- err
		}
	}()
	select {
	case <-ctx.Done():
	case err := <-errCh:
		return err
	case <-state.complete:
		log.Printf("all agent summaries received")
	}
	shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	_ = server.Shutdown(shutdownCtx)
	if output != "" {
		if err := state.writeOutput(output); err != nil {
			return err
		}
	}
	return nil
}

func (s *coordinatorState) handleHistory(w http.ResponseWriter, r *http.Request) {
	machine := strings.ToLower(strings.TrimSpace(r.URL.Query().Get("machine")))
	switch r.Method {
	case http.MethodPost:
		if machine != "a" && machine != "b" {
			http.Error(w, "invalid machine", http.StatusBadRequest)
			return
		}
		var messageIDs []string
		if !decodeLimitedJSON(w, r, &messageIDs) {
			return
		}
		s.mu.Lock()
		if s.history == nil {
			s.history = make(map[string][]string)
		}
		s.history[machine] = messageIDs
		count := len(s.history)
		s.mu.Unlock()
		writeJSON(w, http.StatusOK, map[string]int{"submitted_agents": count})
	case http.MethodGet:
		s.mu.Lock()
		defer s.mu.Unlock()
		if len(s.history) < s.expectedAgents {
			writeJSON(w, http.StatusAccepted, map[string]int{"submitted_agents": len(s.history), "expected_agents": s.expectedAgents})
			return
		}
		unique := make(map[string]bool)
		for _, messageIDs := range s.history {
			for _, mid := range messageIDs {
				if mid != "" {
					unique[mid] = true
				}
			}
		}
		merged := make([]string, 0, len(unique))
		for mid := range unique {
			merged = append(merged, mid)
		}
		sort.Strings(merged)
		writeJSON(w, http.StatusOK, merged)
	default:
		http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
	}
}

func (s *coordinatorState) handleObserved(w http.ResponseWriter, r *http.Request) {
	machine := strings.ToLower(strings.TrimSpace(r.URL.Query().Get("machine")))
	switch r.Method {
	case http.MethodPost:
		if machine != "a" && machine != "b" {
			http.Error(w, "invalid machine", http.StatusBadRequest)
			return
		}
		var messages []AcceptedMessage
		if !decodeLimitedJSON(w, r, &messages) {
			return
		}
		s.mu.Lock()
		if s.observed == nil {
			s.observed = make(map[string][]AcceptedMessage)
		}
		s.observed[machine] = messages
		count := len(s.observed)
		s.mu.Unlock()
		writeJSON(w, http.StatusOK, map[string]int{"submitted_agents": count})
	case http.MethodGet:
		s.mu.Lock()
		defer s.mu.Unlock()
		if len(s.observed) < s.expectedAgents {
			writeJSON(w, http.StatusAccepted, map[string]int{"submitted_agents": len(s.observed), "expected_agents": s.expectedAgents})
			return
		}
		var merged []AcceptedMessage
		for _, messages := range s.observed {
			merged = append(merged, messages...)
		}
		merged = mergeMessageLedgers(nil, merged)
		writeJSON(w, http.StatusOK, merged)
	default:
		http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
	}
}

func (s *coordinatorState) handleReady(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
		return
	}
	var ready AgentReady
	if !decodeLimitedJSON(w, r, &ready) || (ready.Machine != "a" && ready.Machine != "b") ||
		ready.ExpectedConnections <= 0 || ready.ReadyConnections < 0 || ready.ReadyConnections > ready.ExpectedConnections {
		if ready.Machine != "a" && ready.Machine != "b" {
			http.Error(w, "invalid machine", http.StatusBadRequest)
		} else {
			http.Error(w, "invalid connection counts", http.StatusBadRequest)
		}
		return
	}
	s.mu.Lock()
	s.ready[ready.Machine] = ready
	now := time.Now().UTC()
	allStable := len(s.ready) == s.expectedAgents
	for _, agent := range s.ready {
		if agent.ReadyConnections != agent.ExpectedConnections {
			allStable = false
			break
		}
	}
	if !allStable && (!s.start.Committed || now.Before(s.start.StartAt)) {
		log.Printf("agent readiness dropped before start; resetting start barrier")
		s.start.StartAt = time.Time{}
		s.start.Committed = false
		clear(s.armed)
	}
	if allStable && s.start.StartAt.IsZero() {
		s.start.StartAt = now.Add(s.startDelay)
		s.start.Committed = false
		clear(s.armed)
		log.Printf("all agents ready; validation_at=%s", s.start.StartAt.Format(time.RFC3339Nano))
	}
	if allStable && !s.start.Committed && ready.Armed && !now.Before(s.start.StartAt) {
		s.armed[ready.Machine] = true
	}
	if allStable && !s.start.Committed && len(s.armed) == s.expectedAgents {
		delay := s.commitDelay
		if delay <= 0 {
			delay = 5 * time.Second
		}
		s.start.StartAt = now.Add(delay)
		s.start.Committed = true
		log.Printf("all agents validated; committed start_at=%s", s.start.StartAt.Format(time.RFC3339Nano))
	}
	start := s.start
	s.mu.Unlock()
	writeJSON(w, http.StatusOK, start)
}

func (s *coordinatorState) handleStart(w http.ResponseWriter, r *http.Request) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.start.StartAt.IsZero() {
		writeJSON(w, http.StatusAccepted, map[string]any{"ready_agents": len(s.ready), "expected_agents": s.expectedAgents})
		return
	}
	writeJSON(w, http.StatusOK, s.start)
}

func (s *coordinatorState) handleLedger(w http.ResponseWriter, r *http.Request) {
	machine := strings.ToLower(strings.TrimSpace(r.URL.Query().Get("machine")))
	switch r.Method {
	case http.MethodPost:
		if machine != "a" && machine != "b" {
			http.Error(w, "invalid machine", http.StatusBadRequest)
			return
		}
		var messages []AcceptedMessage
		if !decodeLimitedJSON(w, r, &messages) {
			return
		}
		s.mu.Lock()
		s.ledgers[machine] = messages
		count := len(s.ledgers)
		s.mu.Unlock()
		writeJSON(w, http.StatusOK, map[string]int{"submitted_agents": count})
	case http.MethodGet:
		s.mu.Lock()
		defer s.mu.Unlock()
		if len(s.ledgers) < s.expectedAgents {
			writeJSON(w, http.StatusAccepted, map[string]int{"submitted_agents": len(s.ledgers), "expected_agents": s.expectedAgents})
			return
		}
		var merged []AcceptedMessage
		for _, messages := range s.ledgers {
			merged = append(merged, messages...)
		}
		sort.Slice(merged, func(i, j int) bool { return merged[i].ClientMid < merged[j].ClientMid })
		writeJSON(w, http.StatusOK, merged)
	default:
		http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
	}
}

func (s *coordinatorState) handleSummary(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
		return
	}
	var summary AgentSummary
	if !decodeLimitedJSON(w, r, &summary) {
		return
	}
	if summary.Machine != "a" && summary.Machine != "b" {
		http.Error(w, "invalid machine", http.StatusBadRequest)
		return
	}
	s.mu.Lock()
	s.summaries[summary.Machine] = summary
	complete := len(s.summaries) == s.expectedAgents
	s.mu.Unlock()
	writeJSON(w, http.StatusOK, map[string]any{"complete": complete})
	if complete {
		s.completeOnce.Do(func() {
			go func() {
				time.Sleep(500 * time.Millisecond)
				close(s.complete)
			}()
		})
	}
}

func (s *coordinatorState) handleStatus(w http.ResponseWriter, _ *http.Request) {
	s.mu.Lock()
	defer s.mu.Unlock()
	writeJSON(w, http.StatusOK, map[string]any{
		"start": s.start, "ready": s.ready, "armed": s.armed,
		"ledger_agents": len(s.ledgers), "observed_agents": len(s.observed),
		"history_agents": len(s.history), "summary_agents": len(s.summaries),
	})
}

func (s *coordinatorState) writeOutput(path string) error {
	s.mu.Lock()
	data := struct {
		Start     StartInfo                    `json:"start"`
		Ready     map[string]AgentReady        `json:"ready"`
		Ledgers   map[string][]AcceptedMessage `json:"ledgers"`
		Observed  map[string][]AcceptedMessage `json:"observed"`
		History   map[string][]string          `json:"history"`
		Summaries map[string]AgentSummary      `json:"summaries"`
	}{
		Start: s.start, Ready: s.ready, Ledgers: s.ledgers, Observed: s.observed,
		History: s.history, Summaries: s.summaries,
	}
	s.mu.Unlock()
	raw, err := json.MarshalIndent(data, "", "  ")
	if err != nil {
		return err
	}
	if err = os.MkdirAll(filepath.Dir(path), 0o750); err != nil && filepath.Dir(path) != "." {
		return err
	}
	return os.WriteFile(path, append(raw, '\n'), 0o600)
}

func decodeLimitedJSON(w http.ResponseWriter, r *http.Request, output any) bool {
	dec := json.NewDecoder(io.LimitReader(r.Body, 32<<20))
	if err := dec.Decode(output); err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return false
	}
	return true
}

func writeJSON(w http.ResponseWriter, status int, value any) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	_ = json.NewEncoder(w).Encode(value)
}

type coordinatorClient struct {
	base   string
	client *http.Client
}

func newCoordinatorClient(base string) *coordinatorClient {
	return &coordinatorClient{base: strings.TrimRight(base, "/"), client: &http.Client{Timeout: 35 * time.Second}}
}

func (c *coordinatorClient) ready(ctx context.Context, current func() AgentReady) (StartInfo, error) {
	var start StartInfo
	armed := false
	for {
		state := current()
		state.Armed = armed
		if _, err := c.request(ctx, http.MethodPost, "/v1/ready", state, &start); err != nil {
			return start, err
		}
		if start.StartAt.IsZero() {
			armed = false
		} else {
			remaining := time.Until(start.StartAt)
			if start.Committed && remaining <= 0 {
				return start, nil
			}
			if !start.Committed && remaining <= 0 {
				armed = true
				continue
			}
			if remaining < time.Second {
				if err := waitContext(ctx, remaining); err != nil {
					return start, err
				}
				continue
			}
		}
		if err := waitContext(ctx, time.Second); err != nil {
			return start, err
		}
	}
}

func (c *coordinatorClient) submitLedger(ctx context.Context, machine string, messages []AcceptedMessage) error {
	_, err := c.request(ctx, http.MethodPost, "/v1/ledger?machine="+machine, messages, nil)
	return err
}

func (c *coordinatorClient) ledger(ctx context.Context) ([]AcceptedMessage, error) {
	for {
		var messages []AcceptedMessage
		status, err := c.request(ctx, http.MethodGet, "/v1/ledger", nil, &messages)
		if err != nil {
			return nil, err
		}
		if status == http.StatusOK {
			return messages, nil
		}
		if err = waitContext(ctx, time.Second); err != nil {
			return nil, err
		}
	}
}

func (c *coordinatorClient) submitObserved(ctx context.Context, machine string, messages []AcceptedMessage) error {
	_, err := c.request(ctx, http.MethodPost, "/v1/observed?machine="+machine, messages, nil)
	return err
}

func (c *coordinatorClient) observed(ctx context.Context) ([]AcceptedMessage, error) {
	for {
		var messages []AcceptedMessage
		status, err := c.request(ctx, http.MethodGet, "/v1/observed", nil, &messages)
		if err != nil {
			return nil, err
		}
		if status == http.StatusOK {
			return messages, nil
		}
		if err = waitContext(ctx, time.Second); err != nil {
			return nil, err
		}
	}
}

func (c *coordinatorClient) submitHistory(ctx context.Context, machine string, messageIDs []string) error {
	_, err := c.request(ctx, http.MethodPost, "/v1/history?machine="+machine, messageIDs, nil)
	return err
}

func (c *coordinatorClient) history(ctx context.Context) ([]string, error) {
	for {
		var messageIDs []string
		status, err := c.request(ctx, http.MethodGet, "/v1/history", nil, &messageIDs)
		if err != nil {
			return nil, err
		}
		if status == http.StatusOK {
			return messageIDs, nil
		}
		if err = waitContext(ctx, time.Second); err != nil {
			return nil, err
		}
	}
}

func (c *coordinatorClient) submitSummary(ctx context.Context, summary AgentSummary) error {
	_, err := c.request(ctx, http.MethodPost, "/v1/summary", summary, nil)
	return err
}

func (c *coordinatorClient) request(ctx context.Context, method, path string, input, output any) (int, error) {
	var body io.Reader
	if input != nil {
		raw, err := json.Marshal(input)
		if err != nil {
			return 0, err
		}
		body = strings.NewReader(string(raw))
	}
	req, err := http.NewRequestWithContext(ctx, method, c.base+path, body)
	if err != nil {
		return 0, err
	}
	if input != nil {
		req.Header.Set("Content-Type", "application/json")
	}
	resp, err := c.client.Do(req)
	if err != nil {
		return 0, err
	}
	defer resp.Body.Close()
	if resp.StatusCode >= 400 {
		raw, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
		return resp.StatusCode, fmt.Errorf("coordinator HTTP %d: %s", resp.StatusCode, strings.TrimSpace(string(raw)))
	}
	if output != nil && resp.StatusCode == http.StatusOK {
		if err = json.NewDecoder(resp.Body).Decode(output); err != nil {
			return resp.StatusCode, err
		}
	}
	return resp.StatusCode, nil
}

func waitContext(ctx context.Context, delay time.Duration) error {
	timer := time.NewTimer(delay)
	defer timer.Stop()
	select {
	case <-ctx.Done():
		return ctx.Err()
	case <-timer.C:
		return nil
	}
}