package agent

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"strings"

	"ollie/backend"
	"ollie/toolsrv"
)

// BuildRuntime constructs a Runtime from a pre-configured Dispatcher and
// optional agent config. cwd sets the working directory reported in the
// system prompt; if empty, the process working directory is used.
// env provides additional environment variables injected into prompt resolution
// subprocesses (e.g. OLLIE_SESSION_ID=xxx).
// The caller is responsible for registering all servers on d before calling this.
func BuildRuntime(cfg *AgentConfig, srv toolsrv.Runner, cwd string, env []string, baseLayers ...string) *Runtime {
	var messages []string

	var allToolInfos []toolsrv.ToolInfo
	var allTools []backend.Tool

	if srv != nil && (cfg == nil || cfg.ToolsEnabled()) {
		var listErr error
		allToolInfos, listErr = srv.ListTools()
		if listErr != nil {
			messages = append(messages, fmt.Sprintf("list tools: %v", listErr))
		}
		// Only loaded tools become backend callable tools.
		allTools = toolInfosToBackend(allToolInfos)
	}

	hooks := Hooks{}
	var preamble string
	var genParams backend.GenerationParams
	var maxSteps int
	if cfg != nil {
		for k, v := range cfg.Hooks {
			hooks[k] = []string(v)
		}
		if resolved, err := resolvePrompt(cfg.Prompt, cwd, env); err != nil {
			fmt.Fprintf(os.Stderr, "resolve prompt: %v\n", err)
		} else {
			preamble = resolved
		}
		genParams = backend.GenerationParams{
			MaxTokens:           cfg.MaxTokens,
			MaxCompletionTokens: cfg.MaxCompletionTokens,
			Temperature:         cfg.Temperature,
			TopP:                cfg.TopP,
			TopK:                cfg.TopK,
			MinP:                cfg.MinP,
			TopA:                cfg.TopA,
			FrequencyPenalty:    cfg.FrequencyPenalty,
			PresencePenalty:     cfg.PresencePenalty,
			RepetitionPenalty:   cfg.RepetitionPenalty,
			ThinkingBudget:      cfg.Reasoning,
			ReasoningEffort:     cfg.ReasoningEffort,
			IncludeReasoning:    cfg.IncludeReasoning,
			ResponseFormat:      cfg.ResponseFormat,
			Stop:                cfg.Stop,
			Verbosity:           cfg.Verbosity,
		}
		maxSteps = cfg.MaxSteps
		if len(cfg.AllowTools) > 0 {
			if rs, ok := srv.(interface{ SetAllowTools([]string) }); ok {
				rs.SetAllowTools(cfg.AllowTools)
			}
		}
	}

	exec := func(ctx context.Context, name string, args json.RawMessage) (string, []backend.ContentBlock, error) {
		if srv == nil {
			return "", nil, fmt.Errorf("no tool server available")
		}
		raw, err := srv.CallTool(ctx, name, args)
		if err != nil {
			return "", nil, err
		}
		text, blocks, isErr := extractToolResult(raw)
		if isErr {
			return "", nil, fmt.Errorf("%s", text)
		}
		return text, blocks, nil
	}

	var classify toolClassifier
	if pc, ok := srv.(toolsrv.ParallelClassifier); ok {
		classify = pc.IsParallelRead
	}

	// Capture output formats from the same tool listing used to build the
	// backend tool definitions. Remote runners expose ToolInfo through
	// ListTools, but do not necessarily implement OutputFormatClassifier.
	outputFormats := make(map[string]string, len(allToolInfos))
	for _, ti := range allToolInfos {
		outputFormats[ti.Name] = ti.OutputFormat
	}
	outputFormat := func(name string) string {
		return outputFormats[name]
	}
	if oc, ok := srv.(toolsrv.OutputFormatClassifier); ok {
		outputFormat = oc.OutputFormat
	}

	var tierFn func(string, json.RawMessage) ResultTier
	if tc, ok := srv.(toolsrv.TierClassifier); ok {
		tierFn = func(name string, args json.RawMessage) ResultTier {
			switch tc.ResultTierArgs(name, args) {
			case "cold":
				return TierCold
			case "warm":
				return TierWarm
			default:
				return TierHot
			}
		}
	}

	var backendName, modelName, compactionModel string
	if cfg != nil {
		backendName = cfg.Backend
		modelName = cfg.Model
		compactionModel = cfg.CompactionModel
	}

	// Prepend base layers (system prompt, operational model, environment)
	// before the agent-specific preamble.
	if len(baseLayers) > 0 {
		var prefix strings.Builder
		for _, layer := range baseLayers {
			if layer != "" {
				prefix.WriteString(layer)
				prefix.WriteByte('\n')
			}
		}
		if prefix.Len() > 0 {
			preamble = prefix.String() + preamble
		}
	}

	// Append compact tool surface listing (name + description) plus
	// full documentation for loaded tools (schema, args, examples).
	var toolListing strings.Builder
	for _, ti := range allToolInfos {
		if ti.Description != "" && ti.Server == "" {
			fmt.Fprintf(&toolListing, "- **%s** — %s\n", ti.Name, ti.Description)
		}
	}
	if toolListing.Len() > 0 {
		preamble += "\n# Available Tools\n\n" + toolListing.String()
	}
	// Inject full tool prompts so models that ignore native tool schemas
	// still see usage docs, arg definitions, and examples in-context.
	var toolDocs strings.Builder
	for _, ti := range allToolInfos {
		if ti.Prompt != "" && ti.Server == "" {
			toolDocs.WriteString(ti.Prompt)
			toolDocs.WriteString("\n\n")
			if len(ti.InputSchema) > 0 && string(ti.InputSchema) != "{}" {
				fmt.Fprintf(&toolDocs, "**Schema**:\n```json\n%s\n```\n\n", ti.InputSchema)
			}
		}
	}
	if toolDocs.Len() > 0 {
		preamble += "\n# Tool Documentation\n\n" + toolDocs.String()
	}

	// Append 9P discovery instructions for loading additional tools.
	preamble += "\n## Tool discovery\n\nUse the 9P filesystem to discover and load tools dynamically:\n- **List available tools** — `ollie-9p read session/{id}/agent/{aid}/tools`\n- **Load a tool** — `echo \"toolname\" | ollie-9p write session/{id}/agent/{aid}/tools`\n- **See loaded tools** — `ollie-9p read session/{id}/agent/{aid}/tools.loaded`\n"

	return &Runtime{
		ToolServer:      srv,
		Tools:           allTools,
		Exec:            exec,
		ClassifyTool:    classify,
		OutputFormat:    outputFormat,
		ClassifyTier:    tierFn,
		Hooks:           hooks,
		Preamble:        preamble,
		GenParams:       genParams,
		MaxSteps:        maxSteps,
		CfgBackend:      backendName,
		CfgModel:        modelName,
		CompactionModel: compactionModel,
		Messages:        messages,
	}
}

func toolInfosToBackend(infos []toolsrv.ToolInfo) []backend.Tool {
	out := make([]backend.Tool, len(infos))
	for i, t := range infos {
		out[i] = backend.Tool{
			Name:        t.Name,
			Description: t.Description,
			Parameters:  t.InputSchema,
		}
	}
	return out
}

func extractToolResult(raw json.RawMessage) (text string, contentBlocks []backend.ContentBlock, isError bool) {
	var result struct {
		IsError bool `json:"isError"`
		Content []struct {
			Type      string `json:"type"`
			Text      string `json:"text"`
			MediaType string `json:"media_type"`
			Data      string `json:"data"`
		} `json:"content"`
	}
	if err := json.Unmarshal(raw, &result); err != nil {
		return string(raw), nil, false
	}
	var parts []string
	for _, c := range result.Content {
		switch c.Type {
		case "text":
			parts = append(parts, c.Text)
		case "image":
			contentBlocks = append(contentBlocks, backend.ContentBlock{
				Type: "image",
				ImageSource: &backend.ImageSource{
					Type:      "base64",
					MediaType: c.MediaType,
					Data:      c.Data,
				},
			})
		}
	}
	return strings.Join(parts, "\n"), contentBlocks, result.IsError
}