package main

import (
	"encoding/json"
	"errors"
	"net/http"
	"net/url"
	"os"
	"strings"
	"sync"
	"time"
)

type talkWorker struct {
	Sandbox   string `json:"sandbox,omitempty"`
	ID        string `json:"id"`
	Name      string `json:"name"`
	Workspace string `json:"workspace"`
	LastSeen  int64  `json:"last_seen"`
}

type talkTask struct {
	Sandbox        string `json:"sandbox,omitempty"`
	ID             string `json:"id"`
	WorkerID       string `json:"worker_id"`
	Title          string `json:"title"`
	Instructions   string `json:"instructions"`
	Status         string `json:"status"`
	Note           string `json:"note"`
	Created        int64  `json:"created"`
	Updated        int64  `json:"updated"`
	Source         string `json:"source"`
	ConversationID string `json:"conversation_id,omitempty"`
	ResumeOf       string `json:"resume_of,omitempty"`
}

type talkTaskFile struct {
	Workers       []talkWorker       `json:"workers"`
	Tasks         []talkTask         `json:"tasks"`
	Conversations []talkConversation `json:"conversations"`
	Messages      []talkChatMessage  `json:"messages"`
}

type talkPair struct {
	Name    string
	Expires time.Time
}

type talkTaskStore struct {
	mu    sync.Mutex
	path  string
	data  talkTaskFile
	pairs map[string]talkPair // hashes only, short-lived and never persisted
	wake  chan struct{}
}

func newTalkTaskStore(path string) (*talkTaskStore, error) {
	t := &talkTaskStore{path: path, data: talkTaskFile{Workers: []talkWorker{}, Tasks: []talkTask{}}, pairs: map[string]talkPair{}, wake: make(chan struct{})}
	if path != "" {
		b, err := os.ReadFile(path)
		if err != nil && !os.IsNotExist(err) {
			return nil, err
		}
		if err == nil {
			if err := json.Unmarshal(b, &t.data); err != nil {
				return nil, err
			}
		}
	}
	if t.data.Conversations == nil {
		t.data.Conversations = []talkConversation{}
	}
	if t.data.Messages == nil {
		t.data.Messages = []talkChatMessage{}
	}
	next := t.copyLocked()
	for _, worker := range next.Workers {
		ensureWorkerChat(&next, worker)
	}
	if len(next.Conversations) != len(t.data.Conversations) {
		if err := t.commitLocked(next); err != nil {
			return nil, err
		}
	}
	return t, nil
}

// Commit a copy before publishing it, so failed disk writes never acknowledge lost work.
func (t *talkTaskStore) commitLocked(next talkTaskFile) error {
	if t.path != "" {
		b, err := json.MarshalIndent(next, "", "  ")
		if err != nil {
			return err
		}
		if err := os.WriteFile(t.path+".tmp", b, 0600); err != nil {
			return err
		}
		if err := os.Rename(t.path+".tmp", t.path); err != nil {
			return err
		}
	}
	t.data = next
	close(t.wake)
	t.wake = make(chan struct{})
	return nil
}

func (t *talkTaskStore) copyLocked() talkTaskFile {
	return talkTaskFile{Workers: append([]talkWorker{}, t.data.Workers...), Tasks: append([]talkTask{}, t.data.Tasks...), Conversations: append([]talkConversation{}, t.data.Conversations...), Messages: append([]talkChatMessage{}, t.data.Messages...)}
}

func talkDecode(w http.ResponseWriter, r *http.Request, v any) bool {
	d := json.NewDecoder(http.MaxBytesReader(w, r.Body, 40<<10))
	d.DisallowUnknownFields()
	if d.Decode(v) != nil {
		http.Error(w, "invalid JSON request", http.StatusBadRequest)
		return false
	}
	return true
}

func (s *Server) talkOwnerAPI(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Cache-Control", "no-store")
	if !s.authed(r) {
		http.Error(w, "login required", http.StatusUnauthorized)
		return
	}
	if r.Method != http.MethodGet {
		// Cookies identify the owner; a foreign origin must never enqueue instructions.
		if origin := r.Header.Get("Origin"); origin != "" {
			u, err := url.Parse(origin)
			if err != nil || !strings.EqualFold(u.Host, r.Host) || (u.Scheme != "https" && u.Scheme != "http") {
				http.Error(w, "foreign origin", http.StatusForbidden)
				return
			}
		}
		if r.Header.Get("Sec-Fetch-Site") == "cross-site" || !strings.HasPrefix(r.Header.Get("Content-Type"), "application/json") {
			http.Error(w, "same-site JSON request required", http.StatusForbidden)
			return
		}
	}
	t := s.talkTasks
	if t == nil {
		http.Error(w, "task store unavailable", http.StatusServiceUnavailable)
		return
	}
	if strings.HasPrefix(r.URL.Path, "/api/talk/chat") {
		s.talkChatOwnerAPI(w, r)
		return
	}
	t.mu.Lock()
	defer t.mu.Unlock()
	switch {
	case strings.HasPrefix(r.URL.Path, "/api/talk/workers/") && strings.HasSuffix(r.URL.Path, "/cancel-pending") && r.Method == http.MethodPost:
		id := strings.TrimSuffix(strings.TrimPrefix(r.URL.Path, "/api/talk/workers/"), "/cancel-pending")
		found := false
		for _, worker := range t.data.Workers {
			found = found || worker.ID == id
		}
		if !found {
			http.NotFound(w, r)
			return
		}
		for _, task := range t.data.Tasks {
			if task.WorkerID == id && task.Status == "running" {
				http.Error(w, "receiver has running work; stop it locally before cancelling pending work", http.StatusConflict)
				return
			}
		}
		next, count := t.copyLocked(), 0
		for i, task := range next.Tasks {
			if task.WorkerID == id && task.Source == "owner" && (task.Status == "queued" || task.Status == "waiting" || task.Status == "blocked") {
				next.Tasks[i].Status, next.Tasks[i].Updated = "cancelled", time.Now().Unix()
				count++
			}
		}
		if count > 0 && t.commitLocked(next) != nil {
			http.Error(w, "could not cancel pending work", http.StatusInternalServerError)
			return
		}
		writeJSON(w, map[string]any{"cancelled": count})
	case strings.HasPrefix(r.URL.Path, "/api/talk/tasks/") && strings.HasSuffix(r.URL.Path, "/resume") && r.Method == http.MethodPost:
		id := strings.TrimSuffix(strings.TrimPrefix(r.URL.Path, "/api/talk/tasks/"), "/resume")
		for _, task := range t.data.Tasks {
			if task.ResumeOf == id {
				writeJSON(w, task) // A repeated click/lost response must not create another run.
				return
			}
		}
		for _, task := range t.data.Tasks {
			if task.ID != id {
				continue
			}
			if task.Source != "owner" || (task.Status != "waiting" && task.Status != "blocked") {
				http.Error(w, "only waiting or blocked owner work can be continued", http.StatusConflict)
				return
			}
			enabled := false
			for _, key := range s.keys.list() {
				if key.ID == task.WorkerID && !key.Disabled && key.Scope == "talk-worker" {
					enabled = true
				}
			}
			if !enabled || len(t.data.Tasks) >= 2000 {
				http.Error(w, "receiver disabled or task history full", http.StatusConflict)
				return
			}
			next := t.copyLocked()
			task.ResumeOf, task.ID = task.ID, randHex(12)
			task.Status, task.Created, task.Updated = "queued", time.Now().Unix(), time.Now().Unix()
			next.Tasks = append(next.Tasks, task)
			if t.commitLocked(next) != nil {
				http.Error(w, "could not save continuation", http.StatusInternalServerError)
				return
			}
			writeJSON(w, task)
			return
		}
		http.NotFound(w, r)
	case r.URL.Path == "/api/talk/state" && r.Method == http.MethodGet:
		workers := []map[string]any{}
		for _, worker := range t.data.Workers {
			enabled := false
			for _, key := range s.keys.list() {
				if key.ID == worker.ID && !key.Disabled && key.Scope == "talk-worker" {
					enabled = true
					break
				}
			}
			workers = append(workers, map[string]any{"id": worker.ID, "name": worker.Name, "workspace": worker.Workspace, "sandbox": worker.Sandbox, "last_seen": worker.LastSeen, "enabled": enabled, "online": enabled && time.Now().Unix()-worker.LastSeen < 75})
		}
		writeJSON(w, map[string]any{"workers": workers, "tasks": t.data.Tasks})
	case strings.HasPrefix(r.URL.Path, "/api/talk/workers/") && strings.HasSuffix(r.URL.Path, "/sandbox") && r.Method == http.MethodPost:
		id := strings.TrimSuffix(strings.TrimPrefix(r.URL.Path, "/api/talk/workers/"), "/sandbox")
		var req struct {
			Sandbox string `json:"sandbox"`
		}
		if !talkDecode(w, r, &req) {
			return
		}
		if req.Sandbox != "" && req.Sandbox != "read-only" && req.Sandbox != "workspace-write" {
			http.Error(w, "unsupported sandbox mode", http.StatusBadRequest)
			return
		}
		next := t.copyLocked()
		for i := range next.Workers {
			if next.Workers[i].ID != id {
				continue
			}
			next.Workers[i].Sandbox = req.Sandbox
			if t.commitLocked(next) != nil {
				http.Error(w, "could not save sandbox mode", http.StatusInternalServerError)
				return
			}
			writeJSON(w, next.Workers[i])
			return
		}
		http.NotFound(w, r)
	case r.URL.Path == "/api/talk/pair" && r.Method == http.MethodPost:
		var req struct{ Name string }
		if !talkDecode(w, r, &req) {
			return
		}
		req.Name = strings.TrimSpace(req.Name)
		if req.Name == "" || len([]rune(req.Name)) > 60 {
			http.Error(w, "machine name required (max 60 characters)", http.StatusBadRequest)
			return
		}
		for hash, pair := range t.pairs {
			if time.Now().After(pair.Expires) || pair.Name == req.Name {
				delete(t.pairs, hash)
			}
		}
		if len(t.pairs) >= 100 || len(t.data.Workers) >= 200 {
			http.Error(w, "too many pairings", http.StatusConflict)
			return
		}
		code := randHex(12)
		expires := time.Now().Add(10 * time.Minute)
		t.pairs[hashKey(code)] = talkPair{Name: req.Name, Expires: expires}
		writeJSON(w, map[string]any{"code": code, "expires": expires.Unix()})
	case r.URL.Path == "/api/talk/tasks" && r.Method == http.MethodPost:
		var req struct {
			WorkerID     string `json:"worker_id"`
			Title        string `json:"title"`
			Instructions string `json:"instructions"`
		}
		if !talkDecode(w, r, &req) {
			return
		}
		req.Title, req.Instructions = strings.TrimSpace(req.Title), strings.TrimSpace(req.Instructions)
		if req.Title == "" || len([]rune(req.Title)) > 160 || req.Instructions == "" || len([]rune(req.Instructions)) > maxTalkText {
			http.Error(w, "title (1-160) and instructions (1-8000 characters) required", http.StatusBadRequest)
			return
		}
		found := false
		for _, worker := range t.data.Workers {
			if worker.ID == req.WorkerID {
				found = true
			}
		}
		enabled := false
		for _, key := range s.keys.list() {
			if key.ID == req.WorkerID && !key.Disabled && key.Scope == "talk-worker" {
				enabled = true
			}
		}
		if !found || !enabled {
			http.Error(w, "pair an enabled receiver first", http.StatusConflict)
			return
		}
		if len(t.data.Tasks) >= 2000 {
			http.Error(w, "task history full", http.StatusConflict)
			return
		}
		now := time.Now().Unix()
		task := talkTask{ID: randHex(12), WorkerID: req.WorkerID, Title: req.Title, Instructions: req.Instructions, Status: "queued", Created: now, Updated: now, Source: "owner"}
		next := t.copyLocked()
		next.Tasks = append(next.Tasks, task)
		if t.commitLocked(next) != nil {
			http.Error(w, "could not save task", http.StatusInternalServerError)
			return
		}
		writeJSON(w, task)
	default:
		http.NotFound(w, r)
	}
}

// Pair codes are high-entropy one-use credentials, issued only behind the owner login.
func (s *Server) talkPairAPI(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Cache-Control", "no-store")
	if r.Method != http.MethodPost {
		http.Error(w, "POST required", http.StatusMethodNotAllowed)
		return
	}
	var req struct{ Code, Workspace string }
	if !talkDecode(w, r, &req) {
		return
	}
	if len(req.Code) != 24 || req.Workspace == "" || len(req.Workspace) > 1024 {
		http.Error(w, "valid pairing code and workspace required", http.StatusBadRequest)
		return
	}
	t := s.talkTasks
	if t == nil {
		http.Error(w, "pairing unavailable", http.StatusServiceUnavailable)
		return
	}
	t.mu.Lock()
	defer t.mu.Unlock()
	hash := hashKey(req.Code)
	pair, ok := t.pairs[hash]
	if !ok || time.Now().After(pair.Expires) {
		delete(t.pairs, hash)
		http.Error(w, "invalid or expired pairing code", http.StatusUnauthorized)
		return
	}
	delete(t.pairs, hash)
	key, k, err := s.keys.createWorker(pair.Name)
	if err != nil {
		http.Error(w, "could not save credential; create another pairing code", http.StatusInternalServerError)
		return
	}
	next := t.copyLocked()
	worker := talkWorker{ID: k.ID, Name: k.Name, Workspace: req.Workspace}
	next.Workers = append(next.Workers, worker)
	ensureWorkerChat(&next, worker)
	if t.commitLocked(next) != nil {
		s.keys.setDisabled(k.ID, true)
		http.Error(w, "could not save receiver; create another pairing code", http.StatusInternalServerError)
		return
	}
	writeJSON(w, map[string]any{"key": key, "worker": worker})
}

func (s *Server) talkWorkerAPI(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Cache-Control", "no-store")
	k, ok := s.keys.authenticateScope(strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer "), "talk-worker")
	if !ok {
		http.Error(w, "invalid or missing receiver credential", http.StatusUnauthorized)
		return
	}
	done, why := s.keys.admit(k.ID)
	if done == nil {
		http.Error(w, why, http.StatusTooManyRequests)
		return
	}
	defer done()
	s.keys.record(k.ID, 0, 0, false)
	if r.Method != http.MethodPost {
		http.Error(w, "POST required", http.StatusMethodNotAllowed)
		return
	}
	t := s.talkTasks
	if t == nil {
		http.Error(w, "task store unavailable", http.StatusServiceUnavailable)
		return
	}
	if r.URL.Path == "/api/tasks/current" {
		t.mu.Lock()
		defer t.mu.Unlock()
		var current *talkTask
		for _, task := range t.data.Tasks {
			if task.WorkerID == k.ID && task.Status == "running" {
				copy := task
				current = &copy
				break
			}
		}
		writeJSON(w, map[string]any{"task": current})
		return
	}
	if r.URL.Path == "/api/tasks/claim" {
		var req struct{ Wait int }
		if !talkDecode(w, r, &req) {
			return
		}
		deadline := time.Now().Add(time.Duration(min(max(req.Wait, 0), 25)) * time.Second)
		for {
			task, wake, err := t.claim(k.ID)
			if err != nil {
				http.Error(w, "could not claim task", http.StatusInternalServerError)
				return
			}
			if task != nil || time.Until(deadline) <= 0 {
				writeJSON(w, map[string]any{"task": task})
				return
			}
			select {
			case <-wake:
			case <-time.After(time.Until(deadline)):
			case <-r.Context().Done():
				return
			}
			if _, ok := s.keys.authenticateScope(strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer "), "talk-worker"); !ok {
				http.Error(w, "receiver revoked", http.StatusUnauthorized)
				return
			}
		}
	}
	var req struct{ Status, Note string }
	if !talkDecode(w, r, &req) {
		return
	}
	if (req.Status != "running" && req.Status != "done" && req.Status != "waiting" && req.Status != "blocked" && req.Status != "failed") || len([]rune(req.Note)) > maxTalkText {
		http.Error(w, "invalid status or note too long", http.StatusBadRequest)
		return
	}
	t.mu.Lock()
	defer t.mu.Unlock()
	id := strings.TrimPrefix(r.URL.Path, "/api/tasks/")
	next := t.copyLocked()
	for i, task := range next.Tasks {
		if task.ID != id || task.WorkerID != k.ID {
			continue
		}
		if task.Status != "running" {
			if task.Status == req.Status && task.Note == req.Note { // a lost HTTP response can be retried safely
				writeJSON(w, task)
				return
			}
			http.Error(w, "task is not running", http.StatusConflict)
			return
		}
		next.Tasks[i].Status, next.Tasks[i].Note, next.Tasks[i].Updated = req.Status, req.Note, time.Now().Unix()
		if req.Status != "running" && task.ConversationID != "" {
			appendChatMessage(&next, talkChatMessage{ConversationID: task.ConversationID, FromID: k.ID, Text: req.Note, TaskID: task.ID, Status: req.Status, Kind: "result"})
		}
		for j := range next.Workers {
			if next.Workers[j].ID == k.ID {
				next.Workers[j].LastSeen = time.Now().Unix()
			}
		}
		if t.commitLocked(next) != nil {
			http.Error(w, "could not save task status", http.StatusInternalServerError)
			return
		}
		writeJSON(w, next.Tasks[i])
		return
	}
	http.NotFound(w, r)
}

func (t *talkTaskStore) claim(workerID string) (*talkTask, <-chan struct{}, error) {
	t.mu.Lock()
	defer t.mu.Unlock()
	found, busy := false, false
	sandbox := ""
	for i := range t.data.Workers {
		if t.data.Workers[i].ID == workerID {
			sandbox = t.data.Workers[i].Sandbox
			t.data.Workers[i].LastSeen = time.Now().Unix()
			found = true
		}
	}
	if !found {
		return nil, t.wake, errors.New("receiver not registered")
	}
	for _, task := range t.data.Tasks {
		if task.WorkerID == workerID && task.Status == "running" {
			busy = true
		}
	}
	if !busy {
		next := t.copyLocked()
		for i, task := range next.Tasks {
			if task.WorkerID == workerID && task.Status == "queued" {
				// Only the authenticated owner's saved policy determines execution rights.
				next.Tasks[i].Sandbox = sandbox
				next.Tasks[i].Status, next.Tasks[i].Updated = "running", time.Now().Unix()
				if err := t.commitLocked(next); err != nil {
					return nil, t.wake, err
				}
				task = next.Tasks[i]
				return &task, t.wake, nil
			}
		}
	}
	return nil, t.wake, nil
}