package service

import (
	"context"
	"fmt"
	"strconv"
	"time"

	toki "git.toki-labs.com/toki/proto-socket/go"
	"google.golang.org/protobuf/types/known/structpb"

	edgeevents "iop/apps/edge/internal/events"
	edgenode "iop/apps/edge/internal/node"
	eventpkg "iop/packages/events"
	iop "iop/proto/gen/iop"
)

const (
	DefaultSessionID  = "default"
	DefaultTimeoutSec = 30
)

type Service struct {
	registry *edgenode.Registry
	events   *edgeevents.Bus
}

func New(registry *edgenode.Registry, events *edgeevents.Bus) *Service {
	return &Service{registry: registry, events: events}
}

func (s *Service) ListNodes() []*edgenode.NodeEntry {
	return s.registry.All()
}

// NodeSnapshot is a surface-neutral view of a registered node, suitable for
// CLI, HTTP, or other transports. Label is the short display label (node0, node1).
type NodeSnapshot struct {
	NodeID string
	Alias  string
	Label  string
}

// ListNodeSnapshots returns surface-neutral node descriptors derived from the
// current registry contents.
func (s *Service) ListNodeSnapshots() []NodeSnapshot {
	entries := s.registry.All()
	out := make([]NodeSnapshot, 0, len(entries))
	for _, entry := range entries {
		out = append(out, NodeEntrySnapshot(entry))
	}
	return out
}

// NodeEntrySnapshot converts a registry entry into the surface-neutral DTO.
func NodeEntrySnapshot(entry *edgenode.NodeEntry) NodeSnapshot {
	if entry == nil {
		return NodeSnapshot{}
	}
	return NodeSnapshot{
		NodeID: entry.NodeID,
		Alias:  entry.Alias,
		Label:  nodeLabel(entry),
	}
}

func (s *Service) ResolveNode(ref string) (*edgenode.NodeEntry, error) {
	return s.registry.Resolve(ref)
}

// ResolveNodeSnapshot returns a surface-neutral DTO for the resolved node.
func (s *Service) ResolveNodeSnapshot(ref string) (NodeSnapshot, error) {
	entry, err := s.ResolveNode(ref)
	if err != nil {
		return NodeSnapshot{}, err
	}
	return NodeEntrySnapshot(entry), nil
}

type SubmitRunRequest struct {
	NodeRef    string
	RunID      string
	Adapter    string
	Target     string
	SessionID  string
	Prompt     string
	Input      map[string]any
	Background bool
	TimeoutSec int
	Metadata   map[string]string
}

// RunDispatch describes a dispatched run in surface-neutral terms. It is the
// metadata any caller (console, HTTP, future RPC) needs after submission.
type RunDispatch struct {
	RunID      string
	NodeID     string
	NodeLabel  string
	Adapter    string
	Target     string
	SessionID  string
	Background bool
	TimeoutSec int
}

// RunStream carries asynchronous events for a dispatched foreground run.
// Background runs leave both channels nil.
type RunStream struct {
	Events     <-chan *iop.RunEvent
	NodeEvents <-chan *iop.EdgeNodeEvent
}

// RunHandle is the legacy combined surface kept for existing callers; it
// embeds the surface-neutral DTOs so HTTP/API code can consume RunDispatch
// or RunStream directly without depending on this struct.
type RunHandle struct {
	RunDispatch
	RunStream
	close func()
}

func (h *RunHandle) Close() {
	if h != nil && h.close != nil {
		h.close()
	}
}

func (h *RunHandle) WaitTimeout() time.Duration {
	if h == nil {
		return time.Duration(DefaultTimeoutSec+5) * time.Second
	}
	return time.Duration(normalizeTimeoutSec(h.TimeoutSec)+5) * time.Second
}

func (s *Service) SubmitRun(_ context.Context, req SubmitRunRequest) (*RunHandle, error) {
	entry, err := s.ResolveNode(req.NodeRef)
	if err != nil {
		return nil, err
	}

	runReq, runID, err := BuildRunRequest(req)
	if err != nil {
		return nil, err
	}

	var runEvents <-chan *iop.RunEvent
	var unregisterRun func()
	var nodeEvents <-chan *iop.EdgeNodeEvent
	var unregisterNode func()
	if !runReq.GetBackground() {
		if s.events == nil {
			return nil, fmt.Errorf("event bus is not configured")
		}
		runEvents, unregisterRun = s.events.SubscribeRun(runID, 4096)
		nodeEvents, unregisterNode = s.events.SubscribeNode(entry.NodeID, 16)
	}

	if err := entry.Client.Send(runReq); err != nil {
		if unregisterRun != nil {
			unregisterRun()
		}
		if unregisterNode != nil {
			unregisterNode()
		}
		return nil, err
	}

	return &RunHandle{
		RunDispatch: RunDispatch{
			RunID:      runID,
			NodeID:     entry.NodeID,
			NodeLabel:  nodeLabel(entry),
			Adapter:    runReq.GetAdapter(),
			Target:     runReq.GetTarget(),
			SessionID:  runReq.GetSessionId(),
			Background: runReq.GetBackground(),
			TimeoutSec: int(runReq.GetTimeoutSec()),
		},
		RunStream: RunStream{
			Events:     runEvents,
			NodeEvents: nodeEvents,
		},
		close: func() {
			if unregisterRun != nil {
				unregisterRun()
			}
			if unregisterNode != nil {
				unregisterNode()
			}
		},
	}, nil
}

type CancelRunRequest struct {
	NodeRef   string
	RunID     string
	Adapter   string
	Target    string
	SessionID string
}

func BuildCancelRunRequest(req CancelRunRequest) *iop.CancelRequest {
	return &iop.CancelRequest{
		RunId:     req.RunID,
		Adapter:   req.Adapter,
		Target:    req.Target,
		SessionId: NormalizeSessionID(req.SessionID),
		Action:    iop.CancelAction_CANCEL_ACTION_CANCEL_RUN,
	}
}

func (s *Service) CancelRun(_ context.Context, req CancelRunRequest) (CommandResult, error) {
	entry, err := s.ResolveNode(req.NodeRef)
	if err != nil {
		return CommandResult{}, err
	}
	cancelReq := BuildCancelRunRequest(req)
	if err := entry.Client.Send(cancelReq); err != nil {
		return CommandResult{}, err
	}
	return CommandResult{
		NodeID:    entry.NodeID,
		NodeLabel: nodeLabel(entry),
		SessionID: cancelReq.GetSessionId(),
	}, nil
}

type TerminateSessionRequest struct {
	NodeRef   string
	Adapter   string
	Target    string
	SessionID string
}

// CommandResult is the surface-neutral acknowledgement for one-shot node
// commands (terminate-session, future control RPCs).
type CommandResult struct {
	NodeID    string
	NodeLabel string
	SessionID string
}

// TerminateSessionResult is retained as an alias for backward compatibility.
type TerminateSessionResult = CommandResult

func (s *Service) TerminateSession(_ context.Context, req TerminateSessionRequest) (TerminateSessionResult, error) {
	entry, err := s.ResolveNode(req.NodeRef)
	if err != nil {
		return TerminateSessionResult{}, err
	}

	sessionID := NormalizeSessionID(req.SessionID)
	cancelReq := &iop.CancelRequest{
		Adapter:   req.Adapter,
		Target:    req.Target,
		SessionId: sessionID,
		Action:    iop.CancelAction_CANCEL_ACTION_TERMINATE_SESSION,
	}
	if err := entry.Client.Send(cancelReq); err != nil {
		return TerminateSessionResult{}, err
	}
	return TerminateSessionResult{
		NodeID:    entry.NodeID,
		NodeLabel: nodeLabel(entry),
		SessionID: sessionID,
	}, nil
}

type UsageStatusRequest struct {
	NodeRef    string
	Adapter    string
	Target     string
	SessionID  string
	TimeoutSec int
}

// UsageStatusView is the surface-neutral DTO for the usage-status command.
// proto types are retained for the embedded AgentUsageStatus payload.
type UsageStatusView struct {
	NodeID      string
	NodeLabel   string
	Target      string
	SessionID   string
	UsageStatus *iop.AgentUsageStatus
}

// UsageStatusResult is retained as an alias for backward compatibility.
type UsageStatusResult = UsageStatusView

func (s *Service) UsageStatus(_ context.Context, req UsageStatusRequest) (UsageStatusResult, error) {
	entry, err := s.ResolveNode(req.NodeRef)
	if err != nil {
		return UsageStatusResult{}, err
	}

	commandReq := BuildUsageStatusRequest(req.Adapter, req.Target, req.SessionID, req.TimeoutSec)
	resp, err := toki.SendRequestTyped[*iop.NodeCommandRequest, *iop.NodeCommandResponse](
		&entry.Client.Communicator,
		commandReq,
		StatusWaitTimeout(commandReq),
	)
	if err != nil {
		return UsageStatusResult{}, fmt.Errorf("transport error: %w", err)
	}
	if resp.GetError() != "" {
		return UsageStatusResult{}, fmt.Errorf("node reported error: %s", resp.GetError())
	}

	return UsageStatusResult{
		NodeID:      entry.NodeID,
		NodeLabel:   nodeLabel(entry),
		Target:      commandReq.GetTarget(),
		SessionID:   commandReq.GetSessionId(),
		UsageStatus: resp.GetUsageStatus(),
	}, nil
}

func BuildUsageStatusRequest(adapter, targetName, sessionID string, timeoutSec int) *iop.NodeCommandRequest {
	return buildNodeCommandRequest(iop.NodeCommandType_NODE_COMMAND_TYPE_USAGE_STATUS, "status", adapter, targetName, sessionID, timeoutSec)
}

// BuildNodeCommandRequest builds a NodeCommandRequest for any supported type.
// idPrefix is used to namespace request_id (e.g. "status", "caps", "sessions").
func BuildNodeCommandRequest(cmdType iop.NodeCommandType, idPrefix, adapter, targetName, sessionID string, timeoutSec int) *iop.NodeCommandRequest {
	return buildNodeCommandRequest(cmdType, idPrefix, adapter, targetName, sessionID, timeoutSec)
}

func buildNodeCommandRequest(cmdType iop.NodeCommandType, idPrefix, adapter, targetName, sessionID string, timeoutSec int) *iop.NodeCommandRequest {
	if idPrefix == "" {
		idPrefix = "cmd"
	}
	return &iop.NodeCommandRequest{
		RequestId:  fmt.Sprintf("%s-%d", idPrefix, time.Now().UnixNano()),
		Type:       cmdType,
		Adapter:    adapter,
		Target:     targetName,
		SessionId:  NormalizeSessionID(sessionID),
		TimeoutSec: int32(normalizeTimeoutSec(timeoutSec)),
	}
}

func StatusWaitTimeout(req *iop.NodeCommandRequest) time.Duration {
	return time.Duration(normalizeTimeoutSec(int(req.GetTimeoutSec()))+5) * time.Second
}

// NodeCommandRequestSpec is the surface-neutral input for ops console node
// commands (capabilities, session_list, transport_status). UsageStatus keeps
// its own typed request shape.
type NodeCommandRequestSpec struct {
	NodeRef    string
	Adapter    string
	Target     string
	SessionID  string
	TimeoutSec int
}

// NodeCommandView is the surface-neutral result for non-usage-status node
// commands. Result mirrors the proto map and is empty when the node returned
// no payload.
type NodeCommandView struct {
	NodeID    string
	NodeLabel string
	Adapter   string
	Target    string
	SessionID string
	Type      iop.NodeCommandType
	Result    map[string]string
}

// Capabilities dispatches a CAPABILITIES node command and returns the result map.
func (s *Service) Capabilities(_ context.Context, req NodeCommandRequestSpec) (NodeCommandView, error) {
	return s.sendNodeCommand(req, iop.NodeCommandType_NODE_COMMAND_TYPE_CAPABILITIES, "caps")
}

// SessionList dispatches a SESSION_LIST node command and returns the result map.
func (s *Service) SessionList(_ context.Context, req NodeCommandRequestSpec) (NodeCommandView, error) {
	return s.sendNodeCommand(req, iop.NodeCommandType_NODE_COMMAND_TYPE_SESSION_LIST, "sessions")
}

// TransportStatus dispatches a TRANSPORT_STATUS node command and returns the result map.
func (s *Service) TransportStatus(_ context.Context, req NodeCommandRequestSpec) (NodeCommandView, error) {
	return s.sendNodeCommand(req, iop.NodeCommandType_NODE_COMMAND_TYPE_TRANSPORT_STATUS, "transport")
}

type OllamaAPIRequest struct {
	NodeRef    string
	Adapter    string
	Target     string
	Method     string
	Path       string
	Body       string
	TimeoutSec int
}

type OllamaAPIView struct {
	StatusCode  int
	ContentType string
	Body        string
}

func (s *Service) OllamaAPI(_ context.Context, req OllamaAPIRequest) (OllamaAPIView, error) {
	entry, err := s.ResolveNode(req.NodeRef)
	if err != nil {
		return OllamaAPIView{}, err
	}
	commandReq := buildNodeCommandRequest(iop.NodeCommandType_NODE_COMMAND_TYPE_OLLAMA_API, "ollama", req.Adapter, req.Target, "", req.TimeoutSec)
	commandReq.Metadata = map[string]string{
		"ollama_method": req.Method,
		"ollama_path":   req.Path,
		"ollama_body":   req.Body,
	}
	resp, err := toki.SendRequestTyped[*iop.NodeCommandRequest, *iop.NodeCommandResponse](
		&entry.Client.Communicator,
		commandReq,
		StatusWaitTimeout(commandReq),
	)
	if err != nil {
		return OllamaAPIView{}, fmt.Errorf("transport error: %w", err)
	}
	if resp.GetError() != "" {
		return OllamaAPIView{}, fmt.Errorf("node reported error: %s", resp.GetError())
	}
	statusCode, _ := strconv.Atoi(resp.GetResult()["status_code"])
	if statusCode == 0 {
		statusCode = 200
	}
	return OllamaAPIView{
		StatusCode:  statusCode,
		ContentType: resp.GetResult()["content_type"],
		Body:        resp.GetResult()["body"],
	}, nil
}

func (s *Service) sendNodeCommand(req NodeCommandRequestSpec, cmdType iop.NodeCommandType, idPrefix string) (NodeCommandView, error) {
	entry, err := s.ResolveNode(req.NodeRef)
	if err != nil {
		return NodeCommandView{}, err
	}
	commandReq := buildNodeCommandRequest(cmdType, idPrefix, req.Adapter, req.Target, req.SessionID, req.TimeoutSec)
	resp, err := toki.SendRequestTyped[*iop.NodeCommandRequest, *iop.NodeCommandResponse](
		&entry.Client.Communicator,
		commandReq,
		StatusWaitTimeout(commandReq),
	)
	if err != nil {
		return NodeCommandView{}, fmt.Errorf("transport error: %w", err)
	}
	if resp.GetError() != "" {
		return NodeCommandView{}, fmt.Errorf("node reported error: %s", resp.GetError())
	}
	return NodeCommandView{
		NodeID:    entry.NodeID,
		NodeLabel: nodeLabel(entry),
		Adapter:   commandReq.GetAdapter(),
		Target:    commandReq.GetTarget(),
		SessionID: commandReq.GetSessionId(),
		Type:      resp.GetType(),
		Result:    resp.GetResult(),
	}, nil
}

func NormalizeSessionID(id string) string {
	if id == "" {
		return DefaultSessionID
	}
	return id
}

func NewRunID() string {
	return fmt.Sprintf("manual-%d", time.Now().UnixNano())
}

func IsNodeDisconnected(event *iop.EdgeNodeEvent) bool {
	return event.GetType() == eventpkg.TypeNodeDisconnected
}

func BuildRunRequest(req SubmitRunRequest) (*iop.RunRequest, string, error) {
	inputMap := make(map[string]any, len(req.Input)+1)
	for k, v := range req.Input {
		inputMap[k] = v
	}
	if req.Prompt != "" {
		if _, ok := inputMap["prompt"]; !ok {
			inputMap["prompt"] = req.Prompt
		}
	}
	input, err := structpb.NewStruct(inputMap)
	if err != nil {
		return nil, "", err
	}
	runID := req.RunID
	if runID == "" {
		runID = NewRunID()
	}
	metadata := make(map[string]string, len(req.Metadata))
	for k, v := range req.Metadata {
		metadata[k] = v
	}
	return &iop.RunRequest{
		RunId:       runID,
		Adapter:     req.Adapter,
		Target:      req.Target,
		SessionId:   NormalizeSessionID(req.SessionID),
		SessionMode: iop.RunSessionMode_RUN_SESSION_MODE_CREATE_IF_MISSING,
		Background:  req.Background,
		Input:       input,
		TimeoutSec:  int32(normalizeTimeoutSec(req.TimeoutSec)),
		Metadata:    metadata,
	}, runID, nil
}

func normalizeTimeoutSec(timeoutSec int) int {
	if timeoutSec <= 0 {
		return DefaultTimeoutSec
	}
	return timeoutSec
}

func nodeLabel(entry *edgenode.NodeEntry) string {
	return entry.DisplayLabel()
}