// Package ollama provides an Adapter that calls a local Ollama server.
package ollama

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strings"
	"time"

	"go.uber.org/zap"

	"iop/apps/node/internal/runtime"
	"iop/packages/go/config"
)

const Name = "ollama"

const runtimeMetadataOpenAIToolCalls = "openai_tool_calls"

type Ollama struct {
	instanceName     string
	baseURL          string
	contextSize      int
	capacity         int
	maxQueue         int
	queueTimeoutMS   int
	requestTimeoutMS int
	client           *http.Client
	logger           *zap.Logger
}

// New creates an Ollama adapter. The optional instanceName is the registry
// instance key used to disambiguate multiple Ollama instances on the same node.
func New(cfg config.OllamaConf, logger *zap.Logger, instanceName ...string) *Ollama {
	baseURL := strings.TrimRight(cfg.BaseURL, "/")
	if baseURL == "" {
		baseURL = "http://localhost:11434"
	}
	name := Name
	if len(instanceName) > 0 && instanceName[0] != "" {
		name = instanceName[0]
	}
	return &Ollama{
		instanceName:     name,
		baseURL:          baseURL,
		contextSize:      cfg.ContextSize,
		capacity:         cfg.Capacity,
		maxQueue:         cfg.MaxQueue,
		queueTimeoutMS:   cfg.QueueTimeoutMS,
		requestTimeoutMS: cfg.RequestTimeoutMS,
		client:           &http.Client{},
		logger:           logger,
	}
}

func (o *Ollama) Name() string { return Name }

func (o *Ollama) Capabilities(ctx context.Context) (runtime.Capabilities, error) {
	probeCtx, cancel := context.WithTimeout(ctx, 2*time.Second)
	defer cancel()
	targets, err := o.fetchTargets(probeCtx)
	status := runtime.ProviderStatusAvailable
	if err != nil {
		status = runtime.ProviderStatusUnavailable
	}
	return runtime.Capabilities{
		AdapterName:      Name,
		InstanceKey:      o.instanceName,
		Targets:          targets,
		MaxConcurrency:   effectiveCapacity(o.capacity, 4),
		MaxQueue:         o.maxQueue,
		QueueTimeoutMS:   o.queueTimeoutMS,
		RequestTimeoutMS: o.requestTimeoutMS,
		ProviderStatus:   runtime.NormalizeProviderStatus(status),
	}, nil
}

func effectiveCapacity(capacity, defaultVal int) int {
	if capacity <= 0 {
		return defaultVal
	}
	return capacity
}

func (o *Ollama) Execute(ctx context.Context, spec runtime.ExecutionSpec, sink runtime.EventSink) error {
	model := strings.TrimSpace(spec.Target)
	if model == "" {
		model = stringInput(spec.Input, "model")
	}
	if model == "" {
		return fmt.Errorf("ollama adapter: target/model is required")
	}
	messages := ollamaMessagesFromInput(spec.Input)
	if len(messages) == 0 {
		return fmt.Errorf("ollama adapter: prompt/messages are required")
	}

	if err := sink.Emit(ctx, runtime.RuntimeEvent{
		RunID:     spec.RunID,
		Type:      runtime.EventTypeStart,
		Timestamp: time.Now(),
	}); err != nil {
		return err
	}

	chatReq := ollamaChatRequest{
		Model:    model,
		Messages: messages,
		Stream:   true,
	}
	if options := ollamaOptionsFromInput(spec.Input, o.contextSize); len(options) > 0 {
		chatReq.Options = options
	}
	if v, ok := spec.Input["format"]; ok {
		chatReq.Format = v
	}
	if v, ok := spec.Input["keep_alive"]; ok {
		chatReq.KeepAlive = v
	}
	if v, ok := spec.Input["think"]; ok {
		chatReq.Think = v
	}
	if v, ok := spec.Input["tools"]; ok {
		chatReq.Tools = v
	}
	body, err := json.Marshal(chatReq)
	if err != nil {
		return fmt.Errorf("ollama adapter: marshal request: %w", err)
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, joinURL(o.baseURL, "/api/chat"), bytes.NewReader(body))
	if err != nil {
		return fmt.Errorf("ollama adapter: build request: %w", err)
	}
	req.Header.Set("Content-Type", "application/json")

	o.logger.Info("ollama adapter executing",
		zap.String("run_id", spec.RunID),
		zap.String("target", model),
		zap.String("base_url", o.baseURL),
		zap.Int("context_size", o.contextSize),
	)
	resp, err := o.client.Do(req)
	if err != nil {
		_ = emitError(ctx, sink, spec.RunID, fmt.Sprintf("ollama request failed: %v", err))
		return fmt.Errorf("ollama adapter: request: %w", err)
	}
	defer resp.Body.Close()
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		msg := readLimited(resp.Body, 4096)
		if msg == "" {
			msg = resp.Status
		}
		_ = emitError(ctx, sink, spec.RunID, fmt.Sprintf("ollama returned %s: %s", resp.Status, msg))
		return fmt.Errorf("ollama adapter: non-2xx response: %s", resp.Status)
	}

	decoder := json.NewDecoder(resp.Body)
	outputTokens := 0
	var toolCalls []any
	for {
		var chunk ollamaChatChunk
		if err := decoder.Decode(&chunk); err != nil {
			if err == io.EOF {
				break
			}
			_ = emitError(ctx, sink, spec.RunID, fmt.Sprintf("ollama stream decode failed: %v", err))
			return fmt.Errorf("ollama adapter: decode stream: %w", err)
		}
		if chunk.Error != "" {
			_ = emitError(ctx, sink, spec.RunID, chunk.Error)
			return fmt.Errorf("ollama adapter: %s", chunk.Error)
		}
		if chunk.Message.Thinking != "" {
			if err := sink.Emit(ctx, runtime.RuntimeEvent{
				RunID:     spec.RunID,
				Type:      runtime.EventTypeReasoningDelta,
				Delta:     chunk.Message.Thinking,
				Timestamp: time.Now(),
			}); err != nil {
				return err
			}
		}
		if chunk.Message.Content != "" {
			outputTokens += len(strings.Fields(chunk.Message.Content))
			if err := sink.Emit(ctx, runtime.RuntimeEvent{
				RunID:     spec.RunID,
				Type:      runtime.EventTypeDelta,
				Delta:     chunk.Message.Content,
				Timestamp: time.Now(),
			}); err != nil {
				return err
			}
		}
		if len(chunk.Message.ToolCalls) > 0 {
			toolCalls = append(toolCalls, chunk.Message.ToolCalls...)
		}
		if chunk.Done {
			inputTokens := chunk.PromptEvalCount
			if inputTokens == 0 {
				inputTokens = countMessageWords(messages)
			}
			if chunk.EvalCount > 0 {
				outputTokens = chunk.EvalCount
			}
			var metadata map[string]string
			if len(toolCalls) > 0 {
				if encoded, err := json.Marshal(toolCalls); err == nil {
					metadata = map[string]string{
						"finish_reason":                "tool_calls",
						runtimeMetadataOpenAIToolCalls: string(encoded),
					}
				}
			}
			return sink.Emit(ctx, runtime.RuntimeEvent{
				RunID:   spec.RunID,
				Type:    runtime.EventTypeComplete,
				Message: "ollama chat complete",
				Usage: &runtime.UsageStats{
					InputTokens:  inputTokens,
					OutputTokens: outputTokens,
				},
				Metadata:  metadata,
				Timestamp: time.Now(),
			})
		}
	}
	_ = emitError(ctx, sink, spec.RunID, "ollama stream ended without done=true")
	return fmt.Errorf("ollama adapter: stream ended without done=true")
}

func (o *Ollama) ProbeProvider(ctx context.Context, target string) (runtime.ProviderProbeResult, error) {
	probeCtx, cancel := context.WithTimeout(ctx, 2*time.Second)
	defer cancel()

	targets, err := o.fetchTargets(probeCtx)

	result := runtime.ProviderProbeResult{
		AdapterName: Name,
		InstanceKey: o.instanceName,
		Target:      target,
	}

	if err != nil {
		result.Status = runtime.NormalizeProviderStatus(runtime.ProviderStatusUnavailable)
		result.Detail = err.Error()
		return result, nil
	}

	result.Targets = targets

	if target == "" {
		result.Status = runtime.NormalizeProviderStatus(runtime.ProviderStatusAvailable)
		return result, nil
	}

	found := false
	for _, t := range targets {
		if t == target || (strings.Index(target, ":") == -1 && t == target+":latest") {
			found = true
			break
		}
	}

	if found {
		result.Status = runtime.NormalizeProviderStatus(runtime.ProviderStatusAvailable)
	} else {
		result.Status = runtime.NormalizeProviderStatus(runtime.ProviderStatusUnavailable)
		result.Detail = fmt.Sprintf("target model %q not found in provider models", target)
	}

	return result, nil
}

func (o *Ollama) fetchTargets(ctx context.Context) ([]string, error) {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, joinURL(o.baseURL, "/api/tags"), nil)
	if err != nil {
		return nil, fmt.Errorf("build request: %w", err)
	}
	resp, err := o.client.Do(req)
	if err != nil {
		return nil, fmt.Errorf("request failed: %w", err)
	}
	defer resp.Body.Close()
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		return nil, fmt.Errorf("status code %d", resp.StatusCode)
	}
	var tags ollamaTagsResponse
	if err := json.NewDecoder(resp.Body).Decode(&tags); err != nil {
		return nil, fmt.Errorf("decode response: %w", err)
	}
	out := make([]string, 0, len(tags.Models))
	for _, model := range tags.Models {
		if model.Name != "" {
			out = append(out, model.Name)
		}
	}
	return out, nil
}

func (o *Ollama) HandleCommand(ctx context.Context, req runtime.CommandRequest) (runtime.CommandResponse, error) {
	if req.Type != runtime.CommandTypeOllamaAPI {
		return runtime.CommandResponse{}, fmt.Errorf("ollama adapter: unsupported command %q", req.Type)
	}
	method := strings.ToUpper(strings.TrimSpace(req.Metadata["ollama_method"]))
	if method == "" {
		method = http.MethodGet
	}
	if method != http.MethodGet && method != http.MethodPost && method != http.MethodDelete {
		return runtime.CommandResponse{}, fmt.Errorf("ollama adapter: unsupported method %q", method)
	}
	path := strings.TrimSpace(req.Metadata["ollama_path"])
	if !strings.HasPrefix(path, "/api/") {
		return runtime.CommandResponse{}, fmt.Errorf("ollama adapter: unsupported path %q", path)
	}
	body := req.Metadata["ollama_body"]
	var reader io.Reader
	if body != "" {
		reader = strings.NewReader(body)
	}
	httpReq, err := http.NewRequestWithContext(ctx, method, joinURLRequestURI(o.baseURL, path), reader)
	if err != nil {
		return runtime.CommandResponse{}, fmt.Errorf("ollama adapter: build passthrough request: %w", err)
	}
	if body != "" {
		httpReq.Header.Set("Content-Type", "application/json")
	}
	resp, err := o.client.Do(httpReq)
	if err != nil {
		return runtime.CommandResponse{}, fmt.Errorf("ollama adapter: passthrough request: %w", err)
	}
	defer resp.Body.Close()
	respBody := readLimited(resp.Body, 16<<20)
	return runtime.CommandResponse{
		RequestID: req.RequestID,
		Type:      req.Type,
		Adapter:   req.Adapter,
		Target:    req.Target,
		SessionID: req.SessionID,
		Result: map[string]string{
			"status_code":  fmt.Sprintf("%d", resp.StatusCode),
			"content_type": resp.Header.Get("Content-Type"),
			"body":         respBody,
		},
	}, nil
}

func ollamaMessagesFromInput(input map[string]any) []ollamaMessage {
	if raw, ok := input["messages"].([]any); ok {
		out := make([]ollamaMessage, 0, len(raw))
		for _, item := range raw {
			m, ok := item.(map[string]any)
			if !ok {
				continue
			}
			role, _ := m["role"].(string)
			content, _ := m["content"].(string)
			thinking, _ := m["thinking"].(string)
			toolName, _ := m["tool_name"].(string)
			role = strings.TrimSpace(role)
			content = strings.TrimSpace(content)
			if role == "" || (content == "" && thinking == "") {
				continue
			}
			out = append(out, ollamaMessage{
				Role:      role,
				Content:   content,
				Thinking:  thinking,
				Images:    anySlice(m["images"]),
				ToolCalls: anySlice(m["tool_calls"]),
				ToolName:  toolName,
			})
		}
		if len(out) > 0 {
			return out
		}
	}
	prompt := strings.TrimSpace(stringInput(input, "prompt"))
	if prompt == "" {
		return nil
	}
	return []ollamaMessage{{Role: "user", Content: prompt}}
}

func ollamaOptionsFromInput(input map[string]any, contextSize int) map[string]any {
	options := map[string]any{}
	if raw, ok := input["options"].(map[string]any); ok {
		for k, v := range raw {
			options[k] = v
		}
	}
	if contextSize > 0 {
		options["num_ctx"] = contextSize // Edge-owned: always overrides request
	}
	if len(options) == 0 {
		return nil
	}
	return options
}

func anySlice(v any) []any {
	if items, ok := v.([]any); ok {
		return items
	}
	return nil
}

func stringInput(input map[string]any, key string) string {
	if input == nil {
		return ""
	}
	if v, ok := input[key].(string); ok {
		return v
	}
	return ""
}

func countMessageWords(messages []ollamaMessage) int {
	n := 0
	for _, msg := range messages {
		n += len(strings.Fields(msg.Content))
	}
	return n
}

func emitError(ctx context.Context, sink runtime.EventSink, runID, msg string) error {
	return sink.Emit(ctx, runtime.RuntimeEvent{
		RunID:     runID,
		Type:      runtime.EventTypeError,
		Error:     msg,
		Timestamp: time.Now(),
	})
}

func joinURL(baseURL, path string) string {
	u, err := url.Parse(baseURL)
	if err != nil {
		return strings.TrimRight(baseURL, "/") + path
	}
	u.Path = strings.TrimRight(u.Path, "/") + path
	return u.String()
}

func joinURLRequestURI(baseURL, requestURI string) string {
	u, err := url.Parse(baseURL)
	if err != nil {
		return strings.TrimRight(baseURL, "/") + requestURI
	}
	rel, err := url.Parse(requestURI)
	if err != nil {
		u.Path = strings.TrimRight(u.Path, "/") + requestURI
		return u.String()
	}
	u.Path = strings.TrimRight(u.Path, "/") + rel.Path
	u.RawQuery = rel.RawQuery
	return u.String()
}

func readLimited(r io.Reader, limit int64) string {
	b, _ := io.ReadAll(io.LimitReader(r, limit))
	return strings.TrimSpace(string(b))
}

type ollamaChatRequest struct {
	Model     string          `json:"model"`
	Messages  []ollamaMessage `json:"messages"`
	Stream    bool            `json:"stream"`
	Options   map[string]any  `json:"options,omitempty"`
	Format    any             `json:"format,omitempty"`
	KeepAlive any             `json:"keep_alive,omitempty"`
	Think     any             `json:"think,omitempty"`
	Tools     any             `json:"tools,omitempty"`
}

type ollamaMessage struct {
	Role      string `json:"role"`
	Content   string `json:"content"`
	Thinking  string `json:"thinking,omitempty"`
	Images    []any  `json:"images,omitempty"`
	ToolCalls []any  `json:"tool_calls,omitempty"`
	ToolName  string `json:"tool_name,omitempty"`
}

type ollamaChatChunk struct {
	Message         ollamaMessage `json:"message"`
	Done            bool          `json:"done"`
	PromptEvalCount int           `json:"prompt_eval_count"`
	EvalCount       int           `json:"eval_count"`
	Error           string        `json:"error"`
}

type ollamaTagsResponse struct {
	Models []struct {
		Name string `json:"name"`
	} `json:"models"`
}