package agent

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"strings"
	"sync"
	"time"

	"ollie/backend"
	"ollie/toolsrv"
)

const maxTransientRetries = 3

// retryBaseDelay is the base delay for rate-limit retries (5<<attempt seconds).
// Tests can override this to speed up retry tests.
var retryBaseDelay = 5 * time.Second

// streamDropBaseDelay is the base delay for stream-drop retries (2<<attempt seconds).
var streamDropBaseDelay = 2 * time.Second

// Consecutive tool-error thresholds. At the soft limit the model is nudged
// to try a different approach; at the hard limit the loop aborts.
const (
	consecutiveErrorSoftLimit = 5
	consecutiveErrorHardLimit = 10
)

type toolExecutor func(ctx context.Context, name string, args json.RawMessage) (string, []backend.ContentBlock, error)

// cachedResult stores a tool result with file metadata for staleness detection.
type cachedResult struct {
	Result  string
	ModTime int64 // unix nanoseconds of the file at cache time; 0 if no file
	Size    int64 // file size at cache time; -1 if no file
}

// extractFilePath attempts to pull a "path" field from JSON tool args.
func extractFilePath(args json.RawMessage) string {
	var m map[string]json.RawMessage
	if json.Unmarshal(args, &m) != nil {
		return ""
	}
	raw, ok := m["path"]
	if !ok {
		return ""
	}
	var p string
	if json.Unmarshal(raw, &p) != nil {
		return ""
	}
	return p
}

// cacheValid checks if a cached result is still valid by stat-ing the file.
func cacheValid(entry cachedResult, args json.RawMessage) bool {
	if entry.Size < 0 {
		// No file was associated — always valid (non-file tool).
		return true
	}
	path := extractFilePath(args)
	if path == "" {
		return true
	}
	info, err := os.Stat(path)
	if err != nil {
		// File gone or unreadable — cache is stale.
		return false
	}
	return info.ModTime().UnixNano() == entry.ModTime && info.Size() == entry.Size
}

// fileStat returns (modtime nanoseconds, size) for a path, or (0, -1) if unavailable.
func fileStat(path string) (int64, int64) {
	if path == "" {
		return 0, -1
	}
	info, err := os.Stat(path)
	if err != nil {
		return 0, -1
	}
	return info.ModTime().UnixNano(), info.Size()
}

func toolOutputFormat(rt *Runtime, name string) string {
	if ti, ok := rt.ToolMeta[name]; ok {
		return ti.OutputFormat
	}
	return ""
}

func toolMemoryTier(rt *Runtime, name string) MemoryTier {
	if ti, ok := rt.ToolMeta[name]; ok {
		switch ti.Tier {
		case "cold":
			return MemoryCold
		case "warm":
			return MemoryWarm
		}
	}
	return MemoryHot
}

// TurnCtx holds per-turn closures and state that vary between turns within
// the same agent. Everything else comes from *Runtime. Embeds context.Context
// so the turn's cancellation is carried with the turn state rather than as a
// separate parameter.
type TurnCtx struct {
	context.Context
	Output            EventHandler
	PopInject         func() string             // returns and clears pending inject, or ""
	AutoCompact       func(ctx context.Context) // called after each tool round; may compact in-place
	Save              func()                    // called after each state.update(); persists mid-turn progress
	IncrToolCallCount func() int64
	ResultCache       *sync.Map
}

func run(rt *Runtime, ctx TurnCtx, h *History) error {
	var step int
	var consecutiveErrors int // rounds where every tool call returned an error
	var lastErrorSig string   // signature of previous round's error(s) for repeat detection
	var repeatErrorCount int  // consecutive rounds with identical error signature
	// resultCache stores outputs of read-safe tool calls keyed by name+args.
	// Only tools classified as parallel-safe (immutable reads) are cached.
	// sync.Map is required because execOne may run in concurrent goroutines.
	// Uses the cross-turn cache from TurnCtx if available.
	resultCache := ctx.ResultCache
	if resultCache == nil {
		resultCache = &sync.Map{}
	}

	for {
		emit(ctx, Event{Role: "state", Content: "thinking"})
		// Proactive context gate: strip cold material before calling the backend.
		if budget := contextBudget(ctx, rt.Backend); budget > 0 && h.estimateTokens() > budget {
			h.stripCold(ctx, rt.Backend)
		}

		history := h.history()
		if preamble := rt.PreambleString(); preamble != "" {
			history = append([]backend.Message{{Role: "system", Content: preamble}}, history...)
		}

		// Stream the assistant's response, retrying on rate limits, transient
		// backend errors (5xx, network), and mid-stream drops. One stable ID is
		// shared by all chunks and the completed assistant message.
		responseID := NewResponseID()
		var content strings.Builder
		var reasoning strings.Builder
		var toolCalls []backend.ToolCall
		var stopReason string

		for attempt := range maxTransientRetries + 1 {
			content.Reset()
			reasoning.Reset()
			toolCalls = nil

			ch, err := rt.Backend.ChatStream(ctx, history, rt.Tools, rt.GenParams)
			if err != nil {
				if ctx.Err() != nil {
					return ctx.Err()
				}
				wait, retryable := transientWait(err, attempt)
				if !retryable || attempt >= maxTransientRetries {
					return fmt.Errorf("step %d: %w", step, err)
				}
				var rlErr *backend.RateLimitError
				if errors.As(err, &rlErr) {
					emit(ctx, Event{Role: "limitretry"})
				}
				if err := retryCountdown(ctx, wait); err != nil {
					return fmt.Errorf("step %d: %w", step, err)
				}
				continue
			}

			var done bool
			var hadReasoning bool
			for ev := range ch {
				if ev.Reasoning != "" {
					if !hadReasoning {
						emit(ctx, Event{Role: "reasoning", Content: "<think>\n"})
						hadReasoning = true
					}
					reasoning.WriteString(ev.Reasoning)
					emit(ctx, Event{Role: "reasoning", Content: ev.Reasoning})
				}
				if ev.Content != "" {
					if hadReasoning {
						emit(ctx, Event{Role: "reasoning", Content: "\n</think>\n"})
						hadReasoning = false
					}
					content.WriteString(ev.Content)
					emit(ctx, Event{Role: "assistant", Content: ev.Content, ResponseID: responseID})
				}
				toolCalls = append(toolCalls, ev.ToolCalls...)
				if ev.Done {
					if hadReasoning {
						emit(ctx, Event{Role: "reasoning", Content: "\n</think>\n"})
						hadReasoning = false
					}
					stopReason = ev.StopReason
					done = true
					if ev.Usage.InputTokens > 0 || ev.Usage.OutputTokens > 0 {
						emit(ctx, Event{
							Role:    "usage",
							Content: fmt.Sprintf("%d %d 0 %g %d %d", ev.Usage.InputTokens, ev.Usage.OutputTokens, ev.Usage.CostUSD, ev.Usage.CachedInputTokens, ev.Usage.CacheCreationTokens),
						})
					} else {
						// Backend didn't report usage; estimate from content.
						inChars := 0
						for _, m := range history {
							inChars += len(m.Content)
							for _, tc := range m.ToolCalls {
								inChars += len(tc.Name) + len(tc.Arguments)
							}
						}
						emit(ctx, Event{
							Role:    "usage",
							Content: fmt.Sprintf("%d %d 1 0", inChars/4, content.Len()/4),
						})
					}
					break
				}
			}

			if done {
				break
			}

			if ctx.Err() != nil {
				// User pause: record partial state and return.
				if hadReasoning {
					emit(ctx, Event{Role: "reasoning", Content: "\n</think>\n"})
				}
				msg := backend.Message{ID: responseID, Role: "assistant", Content: content.String(), Reasoning: reasoning.String(), ToolCalls: toolCalls}
				var results []toolResult
				for _, tc := range toolCalls {
					results = append(results, toolResult{
						ToolCallID: tc.ID,
						Name:       tc.Name,
						Content:    `{"status":"cancelled","error":"stream interrupted"}`,
						IsError:    true,
					})
				}
				h.update(msg, results)
				ctx.Save()
				return ctx.Err()
			}

			// Pure stream drop — retry if attempts remain.
			if attempt >= maxTransientRetries {
				return fmt.Errorf("step %d: stream dropped (no more retries)", step)
			}
			if hadReasoning {
				emit(ctx, Event{Role: "reasoning", Content: "\n</think>\n"})
			}
			wait := streamDropBaseDelay << attempt
			if err := retryCountdown(ctx, wait); err != nil {
				return fmt.Errorf("step %d: %w", step, err)
			}
		}

		switch stopReason {
		case "stop", "tool_calls", "length", "error", "":
			// normal
		default:
			return fmt.Errorf("step %d: %s", step, stopReason)
		}

		// Parse text-based tool calls: models that don't support the function
		// calling API emit tool invocations as plain text (e.g.
		// "file_read: args=[...]" (text-based tool call syntax)). Parse these into proper ToolCall structs.
		if len(toolCalls) == 0 && len(rt.Tools) > 0 {
			if parsed := parseTextToolCalls(content.String(), rt.Tools); len(parsed) > 0 {
				toolCalls = parsed
			}
		}

		// Execute tool calls, running consecutive parallel-read-safe tools concurrently.
		msg := backend.Message{ID: responseID, Role: "assistant", Content: content.String(), Reasoning: reasoning.String(), ToolCalls: toolCalls}
		results := make([]toolResult, 0, len(toolCalls))
		interrupted := false

		cancelledResult := func(call backend.ToolCall) toolResult {
			return toolResult{
				ToolCallID: call.ID,
				Name:       call.Name,
				Content:    `{"status":"cancelled","error":"interrupted"}`,
				IsError:    true,
			}
		}

		// execOne runs a single tool call end-to-end and reports whether the
		// context was cancelled during execution.
		execOne := func(call backend.ToolCall) (toolResult, bool) {
			if call.Name == "" {
				return toolResult{ToolCallID: call.ID, Name: call.Name, Content: "error: empty tool name", IsError: true}, false
			}
			if ctx.Err() != nil {
				cr := cancelledResult(call)
				emit(ctx, Event{Role: "tool", Name: call.Name, Content: cr.Content, OutputFormat: toolOutputFormat(rt, call.Name)})
				return cr, true
			}
			emit(ctx, Event{Role: "call", Name: call.Name, Content: string(call.Arguments)})
			readSafe := rt.ToolMeta[call.Name].ReadOnly
			if readSafe {
				key := call.Name + "\x00" + string(call.Arguments)
				if v, ok := resultCache.Load(key); ok {
					entry := v.(cachedResult)
					if cacheValid(entry, call.Arguments) {
						emit(ctx, Event{Role: "tool", Name: call.Name, Content: entry.Result, OutputFormat: toolOutputFormat(rt, call.Name)})
						return toolResult{ToolCallID: call.ID, Name: call.Name, Content: entry.Result}, false
					}
					// Stale — remove and re-execute.
					resultCache.Delete(key)
				}
			}
			var result string
			var resultBlocks []backend.ContentBlock
			var isErr bool
			streamed := false
			if rt.Exec != nil {
				streamBytes := 0
				streamCtx := toolsrv.WithOutputStream(ctx, func(data string) {
					if streamBytes >= defaultToolResultMaxBytes {
						return // already at ceiling, drop further chunks
					}
					if !streamed {
						emit(ctx, Event{Role: "tool", Name: call.Name, OutputFormat: toolOutputFormat(rt, call.Name)})
						streamed = true
					}
					streamBytes += len(data)
					if streamBytes > defaultToolResultMaxBytes {
						// Emit only the portion within the ceiling.
						excess := streamBytes - defaultToolResultMaxBytes
						data = data[:len(data)-excess]
					}
					emit(ctx, Event{Role: "tool", Name: call.Name, Content: data, OutputFormat: toolOutputFormat(rt, call.Name)})
				})
				out, blocks, err := rt.Exec(streamCtx, call.Name, call.Arguments)
				if err != nil {
					isErr = true
					if ctx.Err() != nil {
						result = "error: tool execution interrupted by user"
						if ctx.PopInject != nil {
							if injected := ctx.PopInject(); injected != "" {
								result += "\n\n<system-user-interruption>\n" + injected + "\n</system-user-interruption>"
							}
						}
						emit(ctx, Event{Role: "tool", Name: call.Name, Content: result, OutputFormat: toolOutputFormat(rt, call.Name)})
						return toolResult{ToolCallID: call.ID, Name: call.Name, Content: result, IsError: true}, true
					}
					var rlErr *toolsrv.RateLimitedError
					if errors.As(err, &rlErr) {
						result = fmt.Sprintf("error: shell is rate-limited — blocked for %v. Do not call shell() again until the block expires. Use other tools or wait.", rlErr.Remaining)
					} else {
						result = fmt.Sprintf("error: %v", err)
					}
				} else {
					result = out
					resultBlocks = blocks
				}
			} else {
				result = "error: no tool executor configured"
				isErr = true
			}
			// Accumulate suffix text (user-interruptions, truncation hints)
			// that must be emitted after streaming completes.
			var suffix string
			if ctx.PopInject != nil {
				if injected := ctx.PopInject(); injected != "" {
					suffix += "\n\n<system-user-interruption>\n" + injected + "\n</system-user-interruption>"
				}
			}
			result += suffix
			// Safety ceiling: unconditionally cap all tool results at 128KB.
			if len(result) > defaultToolResultMaxBytes {
				orig := len(result)
				result = strings.ToValidUTF8(result[:defaultToolResultMaxBytes], "")
				suffix = fmt.Sprintf("\n\n[HARD LIMIT: %s output truncated — %d of %d bytes shown. This is a safety ceiling, not a semantic boundary.]",
					call.Name, defaultToolResultMaxBytes, orig)
				result += suffix
			}
			if readSafe && !isErr {
				path := extractFilePath(call.Arguments)
				mtime, size := fileStat(path)
				resultCache.Store(call.Name+"\x00"+string(call.Arguments), cachedResult{
					Result:  result,
					ModTime: mtime,
					Size:    size,
				})
			}
			if !streamed {
				emit(ctx, Event{Role: "tool", Name: call.Name, Content: result, OutputFormat: toolOutputFormat(rt, call.Name)})
			} else {
				// Emit suffixes that were appended after streaming completed
				// so they appear at the end of the chat output.
				if suffix != "" {
					emit(ctx, Event{Role: "tool", Name: call.Name, Content: suffix, OutputFormat: toolOutputFormat(rt, call.Name)})
				}
				// Signal the log handler to close the source fence and emit
				// [[[end]]]. Without this, streamingRole stays set and the
				// fence is only closed on the next role transition — or never,
				// if the agent turn ends or is interrupted.
				emit(ctx, Event{Role: "tool", Name: call.Name, OutputFormat: toolOutputFormat(rt, call.Name)})
			}
			tier := MemoryHot
			if !isErr {
				tier = toolMemoryTier(rt, call.Name)
			}
			return toolResult{ToolCallID: call.ID, Name: call.Name, Content: result, ContentBlocks: resultBlocks, IsError: isErr, Tier: tier}, false
		}

		isParallelSafe := func(name string) bool {
			return name != "" && rt.ToolMeta[name].ReadOnly
		}

		for i := 0; i < len(toolCalls) && !interrupted; {
			if ctx.Err() != nil {
				for _, remaining := range toolCalls[i:] {
					cr := cancelledResult(remaining)
					emit(ctx, Event{Role: "tool", Name: remaining.Name, Content: cr.Content, OutputFormat: toolOutputFormat(rt, remaining.Name)})
					results = append(results, cr)
				}
				interrupted = true
				break
			}

			// Collect a run of consecutive parallel-safe calls.
			j := i + 1
			if isParallelSafe(toolCalls[i].Name) {
				for j < len(toolCalls) && isParallelSafe(toolCalls[j].Name) {
					j++
				}
			}
			batch := toolCalls[i:j]

			fillCancelled := func(from int) {
				for _, remaining := range toolCalls[from:] {
					cr := cancelledResult(remaining)
					emit(ctx, Event{Role: "tool", Name: remaining.Name, Content: cr.Content, OutputFormat: toolOutputFormat(rt, remaining.Name)})
					results = append(results, cr)
				}
			}

			if len(batch) == 1 {
				tr, wasInt := execOne(batch[0])
				results = append(results, tr)
				if wasInt {
					fillCancelled(j)
					interrupted = true
				}
			} else {
				// Fan out the batch concurrently; preserve submission order in results.
				// Deduplicate identical calls so only one executes per unique key.
				type inflightResult struct {
					tr     toolResult
					wasInt bool
				}
				inflight := make(map[string]int) // key -> index of first occurrence
				batchResults := make([]toolResult, len(batch))
				batchInt := make([]bool, len(batch))
				var wg sync.WaitGroup
				uniqueResults := make([]inflightResult, len(batch))
				for k, call := range batch {
					key := call.Name + "\x00" + string(call.Arguments)
					if first, dup := inflight[key]; dup {
						// Will copy result from first occurrence after wg.Wait.
						batchResults[k] = toolResult{} // placeholder
						_ = first                      // used below
						continue
					}
					inflight[key] = k
					wg.Add(1)
					go func(k int, call backend.ToolCall) {
						defer wg.Done()
						uniqueResults[k].tr, uniqueResults[k].wasInt = execOne(call)
					}(k, call)
				}
				wg.Wait()
				// Fill results: unique calls get their own result, duplicates copy from first.
				for k, call := range batch {
					key := call.Name + "\x00" + string(call.Arguments)
					first := inflight[key]
					if k == first {
						batchResults[k] = uniqueResults[k].tr
						batchInt[k] = uniqueResults[k].wasInt
					} else {
						batchResults[k] = toolResult{
							ToolCallID: call.ID,
							Name:       call.Name,
							Content:    uniqueResults[first].tr.Content,
							IsError:    uniqueResults[first].tr.IsError,
						}
						batchInt[k] = uniqueResults[first].wasInt
						emit(ctx, Event{Role: "call", Name: call.Name, Content: string(call.Arguments)})
						emit(ctx, Event{Role: "tool", Name: call.Name, Content: batchResults[k].Content, OutputFormat: toolOutputFormat(rt, call.Name)})
					}
				}
				// Add all batch results — model needs a result for every tool call.
				results = append(results, batchResults...)
				for _, wasInt := range batchInt {
					if wasInt {
						fillCancelled(j)
						interrupted = true
						break
					}
				}
			}
			i = j
		}

		h.update(msg, results)
		ctx.Save()

		if interrupted {
			return ctx.Err()
		}

		if len(toolCalls) == 0 {
			break
		}

		// Track consecutive rounds where every tool call errored.
		// Rate-limited errors don't count — they're a temporary block,
		// not a real failure. Reset consecutiveErrors so the hard abort
		// doesn't fire.
		allErrors := len(results) > 0
		allRateLimited := allErrors
		for _, r := range results {
			if !r.IsError {
				allErrors = false
				allRateLimited = false
				break
			}
			if !strings.HasPrefix(r.Content, "error: shell is rate-limited") {
				allRateLimited = false
			}
		}
		if allRateLimited {
			consecutiveErrors = 0
		} else if allErrors {
			consecutiveErrors++
		} else {
			consecutiveErrors = 0
		}
		if consecutiveErrors >= consecutiveErrorHardLimit {
			emit(ctx, Event{Role: "error", Content: fmt.Sprintf("%d consecutive tool errors — aborting", consecutiveErrors)})
			return fmt.Errorf("step %d: %d consecutive tool errors", step, consecutiveErrors)
		}
		if consecutiveErrors == consecutiveErrorSoftLimit {
			// Nudge the model to try a different approach and keep the plan current.
			h.update(backend.Message{
				Role:    "user",
				Content: "<system-consecutive-errors>\nYour last several tool calls all failed. Try a different approach, or ask the user for help. Keep your plan current.\n</system-consecutive-errors>",
			}, nil)
		}

		// Repeated identical error detection: if the same error occurs 3 times
		// in a row, the model is stuck on a format/argument issue. Inject the
		// error content as a corrective nudge immediately rather than waiting
		// for the generic soft limit.
		if allErrors && len(results) > 0 {
			var errSig strings.Builder
			for _, r := range results {
				if r.IsError {
					errSig.WriteString(r.Content)
					errSig.WriteByte('\n')
				}
			}
			sig := errSig.String()
			if sig == lastErrorSig {
				repeatErrorCount++
			} else {
				repeatErrorCount = 1
				lastErrorSig = sig
			}
			if repeatErrorCount == 3 {
				h.update(backend.Message{
					Role:    "user",
					Content: "<system-repeated-error>\nYou have made the same malformed tool call 3 times in a row. Read the error message carefully and fix the arguments. The error was:\n" + strings.TrimSpace(results[0].Content) + "\n</system-repeated-error>",
				}, nil)
			}
		} else {
			repeatErrorCount = 0
			lastErrorSig = ""
		}

		// Reset step counter when an action-oriented tool is called successfully.
		if rt.ToolMeta != nil {
			for _, r := range results {
				if !r.IsError && rt.ToolMeta[r.Name].ResetsCounter {
					step = 0
					break
				}
			}
		}

		// Step budget warnings: nudge the agent at 50% and when exhausted.
		if rt.MaxSteps > 0 {
			halfBudget := rt.MaxSteps / 2
			if step == halfBudget {
				h.update(backend.Message{
					Role:    "user",
					Content: fmt.Sprintf("<system-step-budget-warning>\nYou have used %d/%d research steps without taking action. Consider making progress — write code, edit files, or run commands. Action tools reset this counter.\n</system-step-budget-warning>", step, rt.MaxSteps),
				}, nil)
			} else if step >= rt.MaxSteps-1 {
				emit(ctx, Event{Role: "maxsteps", Content: fmt.Sprintf("%d", step+1)})
				h.update(backend.Message{
					Role:    "user",
					Content: fmt.Sprintf("<system-step-budget-stop>\nStep budget exhausted (%d/%d steps used). Stop calling tools. Summarize what you have done and what remains, then stop.\n</system-step-budget-stop>", step+1, rt.MaxSteps),
				}, nil)
				break
			}
		}

		if ctx.AutoCompact != nil {
			ctx.AutoCompact(ctx)
		}

		step++
	}

	return nil
}

func emit(ctx TurnCtx, msg Event) {
	if ctx.Output != nil {
		ctx.Output(msg)
	}
}

// contextBudget returns the token threshold (50% of context length) above which
// cold material should be proactively stripped. Returns 0 if unknown.
func contextBudget(ctx context.Context, b backend.Backend) int {
	if b == nil {
		return 0
	}
	ctxLen := b.ContextLength(ctx)
	if ctxLen <= 0 {
		return 0
	}
	return ctxLen / 2
}

// defaultToolResultMaxBytes caps tool result content sent back to the model.
const defaultToolResultMaxBytes = 131072



// MemoryTier classifies how long a tool result stays in the hot message list.
type MemoryTier int

const (
	MemoryHot  MemoryTier = iota // stays verbatim in messages (default)
	MemoryWarm                   // summarized on next compaction pass
	MemoryCold                   // immediately summarized at update time
)

// toolResult holds the output of a single tool call.
type toolResult struct {
	ToolCallID    string // may be empty (Ollama does not always set this)
	Name          string
	Content       string
	ContentBlocks []backend.ContentBlock // image blocks extracted from tool output
	IsError       bool
	Tier          MemoryTier // default MemoryHot; tools may self-classify
}
// retryable. Rate limits use longer waits; transient/network errors use shorter.
func transientWait(err error, attempt int) (time.Duration, bool) {
	var rlErr *backend.RateLimitError
	if errors.As(err, &rlErr) {
		wait := rlErr.RetryAfter
		if wait == 0 {
			wait = retryBaseDelay << attempt
		}
		return wait, true
	}
	var tErr *backend.TransientError
	if errors.As(err, &tErr) {
		return time.Duration(2<<attempt) * time.Second, true
	}
	return 0, false
}

func retryCountdown(ctx TurnCtx, wait time.Duration) error {
	deadline := time.Now().Add(wait)
	for {
		remaining := time.Until(deadline)
		if remaining <= 0 {
			return nil
		}
		totalSecs := int(remaining.Seconds()) + 1
		h := totalSecs / 3600
		m := (totalSecs % 3600) / 60
		s := totalSecs % 60
		emit(ctx, Event{Role: "retry", Content: fmt.Sprintf("%02d:%02d:%02d", h, m, s)})
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-time.After(min(remaining, time.Second)):
		}
	}
}