package main

import (
	"context"
	"encoding/json"
	"fmt"
	"log/slog"
	"os"
	"time"

	"github.com/google/uuid"
	"github.com/openai/openai-go"
	"github.com/openai/openai-go/option"
	"github.com/openai/openai-go/shared"
)

const maxIterations = 15

type agent struct {
	client   *mcpClient
	system   string
	provider *openai.Client
	model    string
	store    *store
	agentID  uuid.UUID
	reqOpts  []option.RequestOption
}

func newAgent(ctx context.Context, mcpClient *mcpClient, st *store, agentSlug string) (*agent, error) {
	system := loadSoul()
	apiKey := os.Getenv("OPENROUTER_API_KEY")
	model := os.Getenv("NOMOS_MODEL")
	if model == "" {
		model = "deepseek/deepseek-v4-flash"
	}

	provider := openai.NewClient(
		option.WithBaseURL("https://openrouter.ai/api/v1"),
		option.WithAPIKey(apiKey),
	)

	agentID := st.resolveAgentID(ctx, agentSlug)
	if agentID == uuid.Nil {
		slog.Warn("nomos: agent entity not found; tool-call activity will not be logged", "slug", agentSlug)
	}

	// OpenRouter provider routing. data_collection=deny pins to zero-data-
	// retention providers (privacy: conversations + tool results transit
	// OpenRouter); require_parameters ensures the routed provider actually
	// supports tool calling. NOMOS_PROVIDER_SORT (price|throughput|latency)
	// and Exacto tool-accuracy routing are opt-in — the latter via a model
	// suffix in NOMOS_MODEL (e.g. "deepseek/deepseek-v4-flash:exacto"), so an
	// unsupported value never silently breaks the confirmed routing below.
	providerRouting := map[string]any{
		"data_collection":    "deny",
		"require_parameters": true,
	}
	if sort := os.Getenv("NOMOS_PROVIDER_SORT"); sort != "" {
		providerRouting["sort"] = sort
	}
	reqOpts := []option.RequestOption{option.WithJSONSet("provider", providerRouting)}

	return &agent{
		client:   mcpClient,
		system:   system,
		provider: &provider,
		model:    model,
		store:    st,
		agentID:  agentID,
		reqOpts:  reqOpts,
	}, nil
}

func loadSoul() string {
	paths := []string{"/app/nomos/SOUL.md", "nomos/SOUL.md"}
	for _, p := range paths {
		if data, err := os.ReadFile(p); err == nil {
			return string(data)
		}
	}
	return `You are Nomos, the steward of the oikos — the AI agent for the hubris homelab.
You have access to MCP tools to query topology, health, knowledge, and request
gated mutations through request_execution. Be concise. Prefer tools over guessing.`
}

type toolDef struct {
	Name        string         `json:"name"`
	Description string         `json:"description"`
	InputSchema map[string]any `json:"inputSchema"`
}

type agentEvent struct {
	Type       string `json:"type"`
	Data       any    `json:"data,omitempty"`
	SessionID  string `json:"session_id,omitempty"`
	Iteration  int    `json:"iteration,omitempty"`
}

func (a *agent) chat(ctx context.Context, sessionID, message string, emit func(agentEvent)) {
	correlationID := uuid.New().String()

	tools, err := a.buildTools()
	if err != nil {
		emit(agentEvent{Type: "error", Data: fmt.Sprintf("build tools: %v", err), SessionID: sessionID})
		return
	}

	// Rebuild conversation context from persisted history so sessions are
	// multi-turn. The current user turn is saved by the HTTP handler before
	// this runs, so it is already included in the history for real sessions.
	// Intermediate tool_use/tool_result pairs are not replayed (their ids
	// must match exactly or the API rejects them); prior final answers carry
	// the salient context. Ephemeral sessions (no store) fall back to the
	// single incoming message.
	messages := []openai.ChatCompletionMessageParamUnion{openai.SystemMessage(a.system)}
	history, _ := a.store.getMessages(ctx, sessionID)
	for _, m := range history {
		text := extractText(m.Content)
		switch m.Role {
		case "user":
			messages = append(messages, openai.UserMessage(text))
		case "assistant":
			if text != "" {
				messages = append(messages, openai.AssistantMessage(text))
			}
		}
	}
	if len(history) == 0 {
		messages = append(messages, openai.UserMessage(message))
	}

	for i := 0; i < maxIterations; i++ {
		params := openai.ChatCompletionNewParams{
			Model:    openai.ChatModel(a.model),
			Messages: messages,
			Tools:    tools,
		}

		// Stream the completion, emitting token deltas as they arrive. The
		// accumulator reassembles the full message (content + tool calls) for
		// the loop's control flow.
		stream := a.provider.Chat.Completions.NewStreaming(ctx, params, a.reqOpts...)
		acc := openai.ChatCompletionAccumulator{}
		for stream.Next() {
			chunk := stream.Current()
			acc.AddChunk(chunk)
			if len(chunk.Choices) > 0 {
				if delta := chunk.Choices[0].Delta.Content; delta != "" {
					emit(agentEvent{Type: "text_delta", Data: delta, SessionID: sessionID, Iteration: i + 1})
				}
			}
		}
		if err := stream.Err(); err != nil {
			emit(agentEvent{Type: "error", Data: fmt.Sprintf("llm: %v", err), SessionID: sessionID})
			return
		}
		if len(acc.Choices) == 0 {
			emit(agentEvent{Type: "error", Data: "no choices in response", SessionID: sessionID})
			return
		}

		msg := acc.Choices[0].Message

		if len(msg.ToolCalls) == 0 {
			emit(agentEvent{Type: "text", Data: msg.Content, SessionID: sessionID})
			emit(agentEvent{Type: "done", Data: map[string]any{
				"session_id":     sessionID,
				"usage":          acc.Usage,
				"correlation_id": correlationID,
				"iterations":     i + 1,
			}, SessionID: sessionID})
			return
		}

		slog.Info("nomos: tool calls", "count", len(msg.ToolCalls), "iter", i+1, "correlation", correlationID)

		messages = append(messages, msg.ToParam())

		for _, tc := range msg.ToolCalls {
			var args map[string]any
			if err := json.Unmarshal([]byte(tc.Function.Arguments), &args); err != nil {
				args = map[string]any{}
			}

			emit(agentEvent{
				Type:      "tool_use",
				Data:      map[string]any{"name": tc.Function.Name, "args": args, "id": tc.ID},
				SessionID: sessionID,
				Iteration: i + 1,
			})

			start := time.Now()
			result, callErr := a.client.callTool(tc.Function.Name, args)
			elapsed := int(time.Since(start).Milliseconds())

			inputJSON, _ := json.Marshal(args)
			inputStr := string(inputJSON)

			if callErr != nil {
				a.store.logActivity(ctx, a.agentID, sessionID, tc.Function.Name, inputStr, callErr.Error(), elapsed, false, correlationID)

				emit(agentEvent{
					Type:      "tool_result",
					Data:      map[string]any{"name": tc.Function.Name, "error": callErr.Error(), "id": tc.ID},
					SessionID: sessionID,
					Iteration: i + 1,
				})
				messages = append(messages, openai.ToolMessage(callErr.Error(), tc.ID))
				slog.Error("nomos: tool error", "tool", tc.Function.Name, "error", callErr, "ms", elapsed)
				continue
			}

			resultJSON, _ := json.Marshal(result)
			a.store.logActivity(ctx, a.agentID, sessionID, tc.Function.Name, inputStr, string(resultJSON), elapsed, true, correlationID)

			emit(agentEvent{
				Type:      "tool_result",
				Data:      map[string]any{"name": tc.Function.Name, "result": result, "id": tc.ID},
				SessionID: sessionID,
				Iteration: i + 1,
			})
			messages = append(messages, openai.ToolMessage(string(resultJSON), tc.ID))
			slog.Info("nomos: tool success", "tool", tc.Function.Name, "ms", elapsed)
		}
	}

	emit(agentEvent{Type: "text", Data: "Agent loop: max iterations reached without final answer.", SessionID: sessionID})
	emit(agentEvent{Type: "done", Data: map[string]any{
		"session_id":     sessionID,
		"correlation_id": correlationID,
		"iterations":     maxIterations,
	}, SessionID: sessionID})
}

// extractText pulls the "text" field from a persisted message's JSONB content.
func extractText(content json.RawMessage) string {
	var m struct {
		Text string `json:"text"`
	}
	if err := json.Unmarshal(content, &m); err != nil {
		return ""
	}
	return m.Text
}

func (a *agent) buildTools() ([]openai.ChatCompletionToolParam, error) {
	defs, err := a.client.listToolsFull()
	if err != nil {
		return nil, err
	}

	var tools []openai.ChatCompletionToolParam
	for _, d := range defs {
		params := shared.FunctionParameters(d.InputSchema)
		if params == nil {
			params = shared.FunctionParameters{"type": "object", "properties": map[string]any{}}
		}

		tools = append(tools, openai.ChatCompletionToolParam{
			Type: "function",
			Function: shared.FunctionDefinitionParam{
				Name:        d.Name,
				Description: openai.String(d.Description),
				Parameters:  params,
			},
		})
	}
	return tools, nil
}

func (c *mcpClient) listToolsFull() ([]toolDef, error) {
	resp, err := c.doRequest("tools/list", map[string]any{})
	if err != nil {
		return nil, err
	}
	var tr struct {
		Tools []struct {
			Name        string         `json:"name"`
			Description string         `json:"description"`
			InputSchema map[string]any `json:"inputSchema"`
		} `json:"tools"`
	}
	if err := json.Unmarshal(resp.Result, &tr); err != nil {
		return nil, err
	}
	out := make([]toolDef, len(tr.Tools))
	for i, t := range tr.Tools {
		out[i] = toolDef{
			Name:        t.Name,
			Description: t.Description,
			InputSchema: t.InputSchema,
		}
	}
	return out, nil
}