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"`
}

type AgentReady struct {
	Machine          string `json:"machine"`
	ReadyConnections int    `json:"ready_connections"`
}

type coordinatorState struct {
	mu             sync.Mutex
	start          StartInfo
	startDelay     time.Duration
	expectedAgents int
	ready          map[string]AgentReady
	ledgers        map[string][]AcceptedMessage
	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, expectedAgents: expectedAgents,
		ready: make(map[string]AgentReady), ledgers: make(map[string][]AcceptedMessage),
		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/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) 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.ReadyConnections <= 0 {
		if ready.Machine != "a" && ready.Machine != "b" {
			http.Error(w, "invalid machine", http.StatusBadRequest)
		}
		return
	}
	s.mu.Lock()
	s.ready[ready.Machine] = ready
	if len(s.ready) == s.expectedAgents && s.start.StartAt.IsZero() {
		s.start.StartAt = time.Now().UTC().Add(s.startDelay)
		log.Printf("all agents ready; 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, "ledger_agents": len(s.ledgers), "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"`
		Summaries map[string]AgentSummary      `json:"summaries"`
	}{s.start, s.ready, s.ledgers, 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, ready AgentReady) (StartInfo, error) {
	var start StartInfo
	if _, err := c.request(ctx, http.MethodPost, "/v1/ready", ready, &start); err != nil {
		return start, err
	}
	for start.StartAt.IsZero() {
		if err := waitContext(ctx, time.Second); err != nil {
			return start, err
		}
		status, err := c.request(ctx, http.MethodGet, "/v1/start", nil, &start)
		if err != nil {
			return start, err
		}
		if status != http.StatusOK {
			start = StartInfo{}
		}
	}
	return start, nil
}

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) 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
	}
}