package agents

import (
	"context"
	"crypto/rand"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"strings"
	"sync"
	"time"

	"g.fc.run/misc/byre/internal/api"
	"g.fc.run/misc/byre/internal/claude"
	"g.fc.run/misc/byre/internal/git"
	"g.fc.run/misc/byre/internal/mcp"
	"g.fc.run/misc/byre/internal/store"
)

var (
	ErrConflict      = errors.New("conflict")
	ErrNoGitIdentity = errors.New("set your git name and email (PUT /api/me) before starting agents")
)

const agentActor = "agent"

const restartNotice = "byre restarted while you were mid-turn, so your last turn was cut off. Re-check the working tree (git status and the files you were editing) before continuing where you left off."

type Limits struct {
	Global    int
	PerVendor map[api.AgentSessionVendor]int
}

var DefaultLimits = Limits{
	Global:    8,
	PerVendor: map[api.AgentSessionVendor]int{api.Claude: 6},
}

type Manager struct {
	store        *store.Store
	limits       Limits
	worktreeRoot string

	mu         sync.Mutex
	live       map[string]*live
	closing    bool
	rateLimits map[api.AgentSessionVendor]map[string]any
	merges     map[string]string

	deliverMu sync.Mutex
	capMu     sync.Mutex
}

type live struct {
	sess        *claude.Session
	id          string
	dir         string
	model       string
	gitName     string
	gitEmail    string
	vendor      api.AgentSessionVendor
	taskID      string
	projectID   string
	stopping    bool
	interrupted bool
	asked       bool
	bounced     bool
	compacting  bool
	rateLimited map[string]any

	contextModel  string
	contextUsed   int
	contextWindow int
}

func New(st *store.Store, limits Limits, worktreeRoot string) *Manager {
	return &Manager{
		store:        st,
		limits:       limits,
		worktreeRoot: worktreeRoot,
		live:         map[string]*live{},
		rateLimits:   map[api.AgentSessionVendor]map[string]any{},
		merges:       map[string]string{},
	}
}

func (m *Manager) Create(ctx context.Context, task api.Task, creator api.User, prompt, model string, images ...string) (api.AgentSession, error) {
	return m.create(ctx, task, creator, prompt, model, creator.Login, images)
}

func (m *Manager) create(ctx context.Context, task api.Task, creator api.User, prompt, model, actor string, images []string) (api.AgentSession, error) {
	if creator.GitName == "" || creator.GitEmail == "" {
		return api.AgentSession{}, ErrNoGitIdentity
	}

	if task.Status != api.TaskStatusActive {
		return api.AgentSession{}, fmt.Errorf("%w: task %s is %s", ErrConflict, task.Id, task.Status)
	}

	existing, err := m.store.ListSessions(ctx, task.Id)
	if err != nil {
		return api.AgentSession{}, err
	}
	for _, s := range existing {
		if s.State != api.AgentSessionStateCompleted {
			return api.AgentSession{}, fmt.Errorf("%w: task %s already has session %s", ErrConflict, task.Id, s.Id)
		}
	}

	id := newUUID()
	sess := api.AgentSession{
		Id:              id,
		TaskId:          task.Id,
		Vendor:          api.Claude,
		Model:           model,
		VendorSessionId: id,
		State:           api.AgentSessionStateIdle,
		Role:            api.Implementer,
		CreatedBy:       creator.Login,
		CreatedAt:       time.Now().UTC(),
	}

	err = m.store.CreateSession(ctx, sess, task.ProjectId)
	if err != nil {
		return api.AgentSession{}, err
	}

	return m.turn(ctx, sess, task, actor, prompt, images)
}

func (m *Manager) Send(ctx context.Context, sessionID, actor, text, model string, images ...string) (api.AgentSession, error) {
	sess, err := m.store.GetSession(ctx, sessionID)
	if err != nil {
		return api.AgentSession{}, err
	}

	if sess.State != api.AgentSessionStateIdle {
		return api.AgentSession{}, fmt.Errorf("%w: session %s is %s", ErrConflict, sess.Id, sess.State)
	}

	if model != "" && model != sess.Model {
		err = m.store.SetSessionModel(ctx, sess.Id, model, actor)
		if err != nil {
			return api.AgentSession{}, err
		}
		sess.Model = model
	}

	task, err := m.store.GetTask(ctx, sess.TaskId)
	if err != nil {
		return api.AgentSession{}, err
	}

	return m.turn(ctx, sess, task, actor, text, images)
}

func (m *Manager) Interrupt(ctx context.Context, sessionID string) (api.AgentSession, error) {
	sess, err := m.store.GetSession(ctx, sessionID)
	if err != nil {
		return api.AgentSession{}, err
	}

	if sess.State == api.AgentSessionStateQueued {
		return m.cancelQueued(ctx, sess)
	}

	m.mu.Lock()
	l := m.live[sessionID]
	m.mu.Unlock()

	if sess.State != api.AgentSessionStateRunning || l == nil {
		return api.AgentSession{}, fmt.Errorf("%w: session %s has no turn running", ErrConflict, sess.Id)
	}

	m.mu.Lock()
	l.interrupted = true
	m.mu.Unlock()

	ictx, cancel := context.WithTimeout(ctx, 10*time.Second)
	defer cancel()

	err = l.sess.Interrupt(ictx)
	if err != nil {
		return api.AgentSession{}, err
	}

	return sess, nil
}

func (m *Manager) StopTask(ctx context.Context, taskID, actor string) error {
	sessions, err := m.store.ListSessions(ctx, taskID)
	if err != nil {
		return err
	}

	for _, sess := range sessions {
		if sess.State == api.AgentSessionStateCompleted {
			continue
		}
		err = m.retire(ctx, sess, actor)
		if err != nil {
			return err
		}
	}

	return m.drain(ctx)
}

func (m *Manager) retire(ctx context.Context, sess api.AgentSession, actor string) error {
	m.mu.Lock()
	l := m.live[sess.Id]
	if l != nil {
		l.stopping = true
	}
	m.mu.Unlock()

	if l != nil {
		l.sess.Kill()
	}

	if sess.State == api.AgentSessionStateQueued {
		err := m.store.DeleteQueuedTurn(ctx, sess.Id)
		if err != nil {
			return err
		}
	}

	return m.store.SetSessionState(ctx, sess.Id, api.AgentSessionStateCompleted, actor, nil)
}

func (m *Manager) Recover(ctx context.Context) error {
	sessions, err := m.store.SessionsInState(ctx, api.AgentSessionStateRunning)
	if err != nil {
		return err
	}

	for _, sess := range sessions {
		task, err := m.store.GetTask(ctx, sess.TaskId)
		if err != nil {
			return err
		}

		l, err := m.ensureLive(ctx, sess, task)
		if err != nil {
			slog.Error("resume session", "session", sess.Id, "err", err)
			err = m.store.SetSessionState(ctx, sess.Id, api.AgentSessionStateIdle, "byre", map[string]any{"error": err.Error()})
			if err != nil {
				return err
			}
			continue
		}

		err = m.message(ctx, l, "byre", restartNotice, nil)
		if err != nil {
			return err
		}
		slog.Info("resumed session", "session", sess.Id, "task", task.Id)
	}

	return m.drain(ctx)
}

func (m *Manager) Close() {
	m.mu.Lock()
	m.closing = true
	all := []*live{}
	for _, l := range m.live {
		all = append(all, l)
	}
	m.mu.Unlock()

	for _, l := range all {
		l.sess.Kill()
	}
}

func (m *Manager) LiveCount() int {
	m.mu.Lock()
	defer m.mu.Unlock()
	return len(m.live)
}

type Usage struct {
	Running    int
	Queued     int
	RateLimits map[string]map[string]any
}

func (m *Manager) Usage(ctx context.Context) (Usage, error) {
	u := Usage{RateLimits: map[string]map[string]any{}}

	counts, err := m.store.CountRunning(ctx)
	if err != nil {
		return u, err
	}
	for _, n := range counts {
		u.Running += n
	}

	queued, err := m.store.QueuedTurns(ctx)
	if err != nil {
		return u, err
	}
	u.Queued = len(queued)

	m.mu.Lock()
	defer m.mu.Unlock()
	for vendor, info := range m.rateLimits {
		u.RateLimits[string(vendor)] = info
	}
	return u, nil
}

func (m *Manager) turn(ctx context.Context, sess api.AgentSession, task api.Task, actor, text string, images []string) (api.AgentSession, error) {
	m.capMu.Lock()
	defer m.capMu.Unlock()

	room, err := m.hasRoom(ctx, sess.Vendor)
	if err != nil {
		return api.AgentSession{}, err
	}

	if room {
		return m.start(ctx, sess, task, actor, text, images, api.AgentSessionStateIdle)
	}

	err = transition(ctx, m.store, sess.Id, api.AgentSessionStateIdle, api.AgentSessionStateQueued, actor)
	if err != nil {
		return api.AgentSession{}, err
	}

	err = m.store.EnqueueTurn(ctx, store.QueuedTurn{
		SessionID: sess.Id,
		Vendor:    sess.Vendor,
		Actor:     actor,
		Text:      text,
		Images:    images,
		CreatedAt: time.Now().UTC(),
	})
	if err != nil {
		return api.AgentSession{}, err
	}

	return m.store.GetSession(ctx, sess.Id)
}

func (m *Manager) start(ctx context.Context, sess api.AgentSession, task api.Task, actor, text string, images []string, from api.AgentSessionState) (api.AgentSession, error) {
	err := transition(ctx, m.store, sess.Id, from, api.AgentSessionStateRunning, actor)
	if err != nil {
		return api.AgentSession{}, err
	}

	l, err := m.ensureLive(ctx, sess, task)
	if err == nil {
		err = m.syncModel(ctx, l, sess.Model)
	}
	if err != nil {
		serr := m.store.SetSessionState(ctx, sess.Id, api.AgentSessionStateIdle, actor, map[string]any{"error": err.Error()})
		if serr != nil {
			slog.Error("record failed start", "session", sess.Id, "err", serr)
		}
		return api.AgentSession{}, err
	}

	err = m.message(ctx, l, actor, text, images)
	if err != nil {
		return api.AgentSession{}, err
	}

	return m.store.GetSession(ctx, sess.Id)
}

func (m *Manager) syncModel(ctx context.Context, l *live, model string) error {
	if model == l.model {
		return nil
	}

	cctx, cancel := context.WithTimeout(ctx, 10*time.Second)
	defer cancel()

	_, err := l.sess.Control(cctx, map[string]any{"subtype": "set_model", "model": model})
	if err != nil {
		return fmt.Errorf("set model %s: %w", model, err)
	}
	l.model = model
	return nil
}

func (m *Manager) hasRoom(ctx context.Context, vendor api.AgentSessionVendor) (bool, error) {
	counts, err := m.store.CountRunning(ctx)
	if err != nil {
		return false, err
	}

	total := 0
	for _, n := range counts {
		total += n
	}

	return total < m.limits.Global && counts[vendor] < m.limits.PerVendor[vendor], nil
}

func (m *Manager) drain(ctx context.Context) error {
	m.capMu.Lock()
	defer m.capMu.Unlock()

	queued, err := m.store.QueuedTurns(ctx)
	if err != nil {
		return err
	}

	for _, q := range queued {
		room, err := m.hasRoom(ctx, q.Vendor)
		if err != nil {
			return err
		}
		if !room {
			continue
		}

		sess, err := m.store.GetSession(ctx, q.SessionID)
		if err != nil {
			return err
		}
		task, err := m.store.GetTask(ctx, sess.TaskId)
		if err != nil {
			return err
		}

		err = m.store.DeleteQueuedTurn(ctx, q.SessionID)
		if err != nil {
			return err
		}

		_, err = m.start(ctx, sess, task, q.Actor, q.Text, q.Images, api.AgentSessionStateQueued)
		if err != nil {
			slog.Error("start queued turn", "session", q.SessionID, "err", err)
		}
	}

	return nil
}

func (m *Manager) cancelQueued(ctx context.Context, sess api.AgentSession) (api.AgentSession, error) {
	m.capMu.Lock()
	defer m.capMu.Unlock()

	err := m.store.DeleteQueuedTurn(ctx, sess.Id)
	if errors.Is(err, store.ErrNotFound) {
		return api.AgentSession{}, fmt.Errorf("%w: session %s has no queued turn", ErrConflict, sess.Id)
	}
	if err != nil {
		return api.AgentSession{}, err
	}

	err = transition(ctx, m.store, sess.Id, api.AgentSessionStateQueued, api.AgentSessionStateIdle, "byre")
	if err != nil {
		return api.AgentSession{}, err
	}

	return m.store.GetSession(ctx, sess.Id)
}

func transition(ctx context.Context, st *store.Store, id string, from, to api.AgentSessionState, actor string) error {
	err := st.TransitionSession(ctx, id, from, to, actor)
	if errors.Is(err, store.ErrStateChanged) {
		return fmt.Errorf("%w: %w", ErrConflict, err)
	}
	return err
}

func (m *Manager) message(ctx context.Context, l *live, actor, text string, images []string) error {
	imgs := []claude.Image{}
	for _, id := range images {
		b, data, err := m.store.ReadBlob(ctx, id)
		if err != nil {
			return err
		}
		imgs = append(imgs, claude.Image{MediaType: b.ContentType, Data: data})
	}

	head, err := git.Head(l.dir)
	if err != nil {
		return err
	}
	payload := map[string]any{"text": text, "head": head}
	if len(images) > 0 {
		payload["images"] = images
	}
	err = m.emit(ctx, l, actor, "user.message", payload)
	if err != nil {
		return err
	}

	m.mu.Lock()
	l.interrupted = false
	l.asked = false
	l.bounced = false
	l.compacting = actor == "byre" && text == handoffPrompt
	l.rateLimited = nil
	m.mu.Unlock()

	return l.sess.Send(text, imgs...)
}

func (m *Manager) ensureLive(ctx context.Context, sess api.AgentSession, task api.Task) (*live, error) {
	m.mu.Lock()
	defer m.mu.Unlock()

	if m.closing {
		return nil, errors.New("shutting down")
	}

	l := m.live[sess.Id]
	if l != nil {
		return l, nil
	}

	creator, err := m.store.GetUser(ctx, sess.CreatedBy)
	if err != nil {
		return nil, err
	}

	l = &live{
		id:        sess.Id,
		dir:       task.WorktreePath,
		model:     sess.Model,
		gitName:   creator.GitName,
		gitEmail:  creator.GitEmail,
		vendor:    sess.Vendor,
		taskID:    task.Id,
		projectID: task.ProjectId,
	}

	opts := claude.Options{
		Cwd:        task.WorktreePath,
		Model:      sess.Model,
		GitName:    creator.GitName,
		GitEmail:   creator.GitEmail,
		Guidance:   guidance,
		MCPServers: []*mcp.Server{{Name: "byre", Version: "0.0.0", Tools: []mcp.Tool{m.askTool(l), m.createTaskTool(l)}}},
		Hooks: map[string][]claude.HookMatcher{
			"UserPromptSubmit": {{Func: remindGuidance}},
			"Stop":             {{Func: m.bounceProseQuestion(l)}},
		},
	}
	started, err := m.store.SessionStarted(ctx, sess.Id)
	if err != nil {
		return nil, err
	}
	if started {
		opts.Resume = sess.VendorSessionId
	} else {
		opts.SessionID = sess.VendorSessionId
	}

	sctx, cancel := context.WithTimeout(ctx, 60*time.Second)
	defer cancel()

	l.sess, err = claude.Start(sctx, opts)
	if err != nil {
		return nil, fmt.Errorf("start claude: %w", err)
	}

	m.live[sess.Id] = l
	go m.pump(l)
	return l, nil
}

func (m *Manager) pump(l *live) {
	ctx := context.Background()

	for ev := range l.sess.Events() {
		err := m.handle(ctx, l, ev)
		if err != nil {
			slog.Error("record agent event", "session", l.id, "type", ev.Type, "err", err)
		}
	}

	m.mu.Lock()
	delete(m.live, l.id)
	quiet := m.closing || l.stopping
	m.mu.Unlock()

	if quiet {
		return
	}

	sess, err := m.store.GetSession(ctx, l.id)
	if err != nil {
		slog.Error("load exited session", "session", l.id, "err", err)
		return
	}

	if sess.State != api.AgentSessionStateRunning {
		return
	}

	err = m.store.SetSessionState(ctx, l.id, api.AgentSessionStateIdle, agentActor, map[string]any{"error": "agent process exited mid-turn"})
	if err != nil {
		slog.Error("record session exit", "session", l.id, "err", err)
	}
	m.fileError(l, "The agent process exited mid-turn. Send a message to resume it.")

	err = m.drain(ctx)
	if err != nil {
		slog.Error("start queued turns", "err", err)
	}
}

func (m *Manager) handle(ctx context.Context, l *live, ev claude.Event) error {
	switch ev.Type {
	case "assistant":
		blocks, err := ev.Blocks()
		if err != nil {
			return err
		}
		for _, b := range blocks {
			err = m.handleAssistantBlock(ctx, l, b)
			if err != nil {
				return err
			}
		}
		usage, err := ev.ContextUsage()
		if err != nil {
			return err
		}
		if usage.Tokens == 0 {
			return nil
		}
		l.contextModel = usage.Model
		return m.updateContext(ctx, l, usage.Tokens, l.contextWindow)

	case "user":
		blocks, err := ev.Blocks()
		if err != nil {
			return err
		}
		for _, b := range blocks {
			if b.Type != "tool_result" {
				continue
			}
			content, err := decode(b.Content)
			if err != nil {
				return err
			}
			err = m.emit(ctx, l, agentActor, "agent.tool_result", map[string]any{
				"tool_use_id": b.ToolUseID,
				"is_error":    b.IsError,
				"content":     content,
			})
			if err != nil {
				return err
			}
		}
		return nil

	case "result":
		res, err := ev.Result()
		if err != nil {
			return err
		}

		m.mu.Lock()
		interrupted := l.interrupted
		rateLimited := l.rateLimited
		compacting := l.compacting
		m.mu.Unlock()

		finished := !res.IsError && !interrupted && strings.TrimSpace(res.Result) != ""
		handoff := compacting && finished

		payload := store.ToPayload(res)
		sha, err := git.CommitAll(l.dir, commitMessage(res), l.gitName, l.gitEmail)
		if err != nil {
			slog.Error("commit turn", "session", l.id, "err", err)
		}
		if sha != "" {
			payload["commit"] = sha
		}
		switch {
		case handoff:
			payload["handoff"] = true
		case finished && !compacting:
			payload["inbox_item_id"] = m.fileSummary(l, res.Result)
		}
		err = m.emit(ctx, l, agentActor, "turn.ended", payload)
		if err != nil {
			return err
		}

		window := res.ModelUsage[l.contextModel].ContextWindow
		if window > 0 {
			err = m.updateContext(ctx, l, l.contextUsed, window)
			if err != nil {
				return err
			}
		}

		failure := fmt.Sprintf("The turn ended with an error (%s): %s", res.Subtype, res.Result)
		switch {
		case !res.IsError || interrupted:
		case rateLimited != nil:
			m.fileRateLimited(l, rateLimited, failure)
		default:
			m.fileError(l, failure)
		}

		err = m.store.SetSessionState(ctx, l.id, api.AgentSessionStateIdle, agentActor, nil)
		if err != nil {
			return err
		}

		err = m.drain(ctx)
		if err != nil {
			slog.Error("start queued turns", "err", err)
		}
		if handoff {
			go m.continueFromHandoff(l.id, res.Result)
			return nil
		}
		mergeBy, merging := m.takeMerge(l.taskID)
		if merging && finished {
			go m.finishMerge(l, mergeBy)
			return nil
		}
		err = m.Deliver(ctx, l.id)
		if err != nil {
			slog.Error("deliver inbox responses", "session", l.id, "err", err)
		}
		return nil

	case "stream_event":
		text, err := ev.TextDelta()
		if err != nil || text == "" {
			return err
		}
		m.store.Broadcast(api.Event{
			Kind:           "agent.delta",
			ProjectId:      &l.projectID,
			TaskId:         &l.taskID,
			AgentSessionId: &l.id,
			Actor:          agentActor,
			Payload:        map[string]any{"text": text},
		})
		return nil

	case "rate_limit_event":
		info := struct {
			RateLimitInfo map[string]any `json:"rate_limit_info"`
		}{}
		err := json.Unmarshal(ev.Raw, &info)
		if err != nil {
			return err
		}

		m.mu.Lock()
		m.rateLimits[l.vendor] = info.RateLimitInfo
		if info.RateLimitInfo["status"] == "rejected" {
			l.rateLimited = info.RateLimitInfo
		}
		m.mu.Unlock()

	case "system":
		if ev.Subtype == "init" {
			return nil
		}
	}

	raw := map[string]any{}
	err := json.Unmarshal(ev.Raw, &raw)
	if err != nil {
		return err
	}
	return m.emit(ctx, l, agentActor, "agent.vendor", raw)
}

func commitMessage(res claude.Result) string {
	line, _, _ := strings.Cut(strings.TrimSpace(res.Result), "\n")
	sentence, _, found := strings.Cut(line, ". ")
	if found {
		sentence += "."
	}
	sentence = strings.TrimSpace(strings.ReplaceAll(sentence, "**", ""))
	if sentence == "" {
		return fmt.Sprintf("Work from a turn that ended without a summary (%s)", res.Subtype)
	}
	runes := []rune(sentence)
	if len(runes) > 100 {
		return string(runes[:99]) + "…"
	}
	return sentence
}

func (m *Manager) updateContext(ctx context.Context, l *live, used, window int) error {
	if used == l.contextUsed && window == l.contextWindow {
		return nil
	}
	l.contextUsed = used
	l.contextWindow = window
	return m.emit(ctx, l, agentActor, "session.context", map[string]any{"model": l.contextModel, "used": used, "window": window})
}

func (m *Manager) handleAssistantBlock(ctx context.Context, l *live, b claude.Block) error {
	switch b.Type {
	case "text":
		return m.emit(ctx, l, agentActor, "agent.message", map[string]any{"text": b.Text})

	case "thinking":
		if b.Thinking == "" {
			return nil
		}
		return m.emit(ctx, l, agentActor, "agent.thinking", map[string]any{"text": b.Thinking})

	case "tool_use":
		input, err := decode(b.Input)
		if err != nil {
			return err
		}
		return m.emit(ctx, l, agentActor, "agent.tool_call", map[string]any{
			"id":    b.ID,
			"name":  b.Name,
			"input": input,
		})
	}

	return m.emit(ctx, l, agentActor, "agent.vendor", store.ToPayload(b))
}

func (m *Manager) emit(ctx context.Context, l *live, actor, kind string, payload map[string]any) error {
	return m.store.AppendEvent(ctx, api.Event{
		Kind:           kind,
		ProjectId:      &l.projectID,
		TaskId:         &l.taskID,
		AgentSessionId: &l.id,
		Actor:          actor,
		Payload:        payload,
	})
}

func decode(raw json.RawMessage) (any, error) {
	if len(raw) == 0 {
		return nil, nil
	}
	var v any
	err := json.Unmarshal(raw, &v)
	return v, err
}

func newUUID() string {
	b := make([]byte, 16)
	rand.Read(b)
	b[6] = b[6]&0x0f | 0x40
	b[8] = b[8]&0x3f | 0x80
	return fmt.Sprintf("%x-%x-%x-%x-%x", b[0:4], b[4:6], b[6:8], b[8:10], b[10:])
}