package service

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"io"
	"path"
	"strings"

	"iop/packages/go/config"
	iop "iop/proto/gen/iop"
)

const (
	InternalWorkspaceToolRead    = "workspace_read"
	InternalWorkspaceToolList    = "workspace_list"
	InternalWorkspaceToolWrite   = "workspace_write"
	InternalWorkspaceToolDelete  = "workspace_delete"
	InternalWorkspaceToolCommand = "workspace_command"
)

var (
	ErrSingleRequestInternalToolUnavailable = errors.New("single-request internal tool continuation is unavailable")
	ErrSingleRequestInternalToolInvalidCall = errors.New("single-request internal tool call is invalid")
	ErrSingleRequestInternalToolDenied      = errors.New("single-request internal tool capability is denied")
	ErrSingleRequestInternalToolBudget      = errors.New("single-request internal tool budget is exhausted")
	ErrSingleRequestInternalToolFailed      = errors.New("single-request internal tool execution failed")
)

// InternalWorkspaceToolCall is the closed, service-owned tool-call envelope
// emitted by an internal provider stage. Arguments are decoded only by the
// operation-specific decoder below; caller-facing tool codecs are not involved.
type InternalWorkspaceToolCall struct {
	RequestID  string
	StageID    string
	ToolCallID string
	Name       string
	Arguments  json.RawMessage
}

// Clone returns an independent copy so an executor cannot mutate an admitted
// call after the coordinator has accepted it.
func (c *InternalWorkspaceToolCall) Clone() *InternalWorkspaceToolCall {
	if c == nil {
		return nil
	}
	return &InternalWorkspaceToolCall{
		RequestID:  c.RequestID,
		StageID:    c.StageID,
		ToolCallID: c.ToolCallID,
		Name:       c.Name,
		Arguments:  append(json.RawMessage(nil), c.Arguments...),
	}
}

// InternalWorkspaceToolResult contains only bounded typed fields accepted from
// the private Node wire. Raw Node error text and tool arguments are omitted.
type InternalWorkspaceToolResult struct {
	RequestID  string
	StageID    string
	ToolCallID string
	Status     string
	ErrorCode  string
	Content    []byte
	Entries    []string
	Stdout     []byte
	Stderr     []byte
	ExitCode   int32
	Truncated  bool
	DurationMS int64
}

// Clone returns an independent result suitable for one continuation delivery.
func (r InternalWorkspaceToolResult) Clone() InternalWorkspaceToolResult {
	r.Content = append([]byte(nil), r.Content...)
	r.Entries = append([]string(nil), r.Entries...)
	r.Stdout = append([]byte(nil), r.Stdout...)
	r.Stderr = append([]byte(nil), r.Stderr...)
	return r
}

// ValidateInternalWorkspaceToolCall applies the service-owned closed decoder
// without executing the call. Private provider stages use it to reject or
// boundedly correct malformed model output before the coordinator enters an
// internal-tool state.
func ValidateInternalWorkspaceToolCall(call *InternalWorkspaceToolCall) error {
	_, err := decodeInternalWorkspaceToolCall(call.Clone())
	return err
}

// SingleRequestToolContinuation is optional. An executor that emits an
// internal workspace call must implement it so the coordinator can deliver the
// correlated Node result without involving an HTTP caller.
type SingleRequestToolContinuation interface {
	ContinueInternalTool(context.Context, InternalWorkspaceToolResult) error
}

type internalWorkspacePathArguments struct {
	RelativePath *string `json:"relative_path"`
}

type internalWorkspaceWriteArguments struct {
	RelativePath *string `json:"relative_path"`
	Content      *string `json:"content"`
}

type internalWorkspaceCommandArguments struct {
	CommandID   *string           `json:"command_id"`
	Environment map[string]string `json:"environment,omitempty"`
}

// decodeInternalWorkspaceToolCall performs closed per-operation JSON decoding.
// It rejects unknown fields, duplicate object keys, trailing values, malformed
// identities, non-canonical paths, and every operation outside the five-name
// internal schema. Returned errors never contain raw argument data.
func decodeInternalWorkspaceToolCall(call *InternalWorkspaceToolCall) (*iop.WorkspaceToolRequest, error) {
	if call == nil || !validInternalToolIdentity(call.RequestID) || !validInternalToolIdentity(call.StageID) ||
		!validInternalToolIdentity(call.ToolCallID) || call.Name == "" || strings.TrimSpace(call.Name) != call.Name ||
		len(call.Arguments) == 0 || len(call.Arguments) > config.MaxSingleRequestOutputBytes {
		return nil, ErrSingleRequestInternalToolInvalidCall
	}

	req := &iop.WorkspaceToolRequest{
		RequestId:  call.RequestID,
		StageId:    call.StageID,
		ToolCallId: call.ToolCallID,
	}
	switch call.Name {
	case InternalWorkspaceToolRead, InternalWorkspaceToolList, InternalWorkspaceToolDelete:
		var args internalWorkspacePathArguments
		if strictDecodeInternalToolArguments(call.Arguments, &args) != nil || args.RelativePath == nil ||
			!validInternalWorkspacePath(*args.RelativePath) {
			return nil, ErrSingleRequestInternalToolInvalidCall
		}
		switch call.Name {
		case InternalWorkspaceToolRead:
			req.Operation = iop.WorkspaceOperation_WORKSPACE_OPERATION_READ
		case InternalWorkspaceToolList:
			req.Operation = iop.WorkspaceOperation_WORKSPACE_OPERATION_LIST
		case InternalWorkspaceToolDelete:
			req.Operation = iop.WorkspaceOperation_WORKSPACE_OPERATION_DELETE
		}
		req.Input = &iop.WorkspaceToolRequest_RelativePath{RelativePath: *args.RelativePath}
	case InternalWorkspaceToolWrite:
		var args internalWorkspaceWriteArguments
		if strictDecodeInternalToolArguments(call.Arguments, &args) != nil || args.RelativePath == nil || args.Content == nil ||
			!validInternalWorkspacePath(*args.RelativePath) || *args.RelativePath == "." {
			return nil, ErrSingleRequestInternalToolInvalidCall
		}
		req.Operation = iop.WorkspaceOperation_WORKSPACE_OPERATION_WRITE
		req.Input = &iop.WorkspaceToolRequest_Write{Write: &iop.WorkspaceWriteInput{
			RelativePath: *args.RelativePath,
			Content:      []byte(*args.Content),
		}}
	case InternalWorkspaceToolCommand:
		var args internalWorkspaceCommandArguments
		if strictDecodeInternalToolArguments(call.Arguments, &args) != nil || args.CommandID == nil ||
			!validInternalToolIdentity(*args.CommandID) {
			return nil, ErrSingleRequestInternalToolInvalidCall
		}
		for name := range args.Environment {
			if !validEnvironmentName(name) {
				return nil, ErrSingleRequestInternalToolInvalidCall
			}
		}
		req.Operation = iop.WorkspaceOperation_WORKSPACE_OPERATION_COMMAND
		req.Input = &iop.WorkspaceToolRequest_CommandId{CommandId: *args.CommandID}
		req.Environment = cloneStringMap(args.Environment)
	default:
		return nil, ErrSingleRequestInternalToolInvalidCall
	}
	return req, nil
}

func strictDecodeInternalToolArguments(raw json.RawMessage, target any) error {
	if err := validateUniqueInternalToolJSON(raw); err != nil {
		return ErrSingleRequestInternalToolInvalidCall
	}
	decoder := json.NewDecoder(bytes.NewReader(raw))
	decoder.DisallowUnknownFields()
	if err := decoder.Decode(target); err != nil {
		return ErrSingleRequestInternalToolInvalidCall
	}
	var trailing any
	if err := decoder.Decode(&trailing); err != io.EOF {
		return ErrSingleRequestInternalToolInvalidCall
	}
	return nil
}

func validateUniqueInternalToolJSON(raw json.RawMessage) error {
	decoder := json.NewDecoder(bytes.NewReader(raw))
	decoder.UseNumber()
	if err := consumeUniqueInternalToolJSONValue(decoder); err != nil {
		return err
	}
	if _, err := decoder.Token(); err != io.EOF {
		return ErrSingleRequestInternalToolInvalidCall
	}
	return nil
}

func consumeUniqueInternalToolJSONValue(decoder *json.Decoder) error {
	token, err := decoder.Token()
	if err != nil {
		return err
	}
	delim, ok := token.(json.Delim)
	if !ok {
		return nil
	}
	switch delim {
	case '{':
		keys := make(map[string]struct{})
		for decoder.More() {
			keyToken, err := decoder.Token()
			if err != nil {
				return err
			}
			key, ok := keyToken.(string)
			if !ok {
				return ErrSingleRequestInternalToolInvalidCall
			}
			if _, duplicate := keys[key]; duplicate {
				return ErrSingleRequestInternalToolInvalidCall
			}
			keys[key] = struct{}{}
			if err := consumeUniqueInternalToolJSONValue(decoder); err != nil {
				return err
			}
		}
		end, err := decoder.Token()
		if err != nil || end != json.Delim('}') {
			return ErrSingleRequestInternalToolInvalidCall
		}
	case '[':
		for decoder.More() {
			if err := consumeUniqueInternalToolJSONValue(decoder); err != nil {
				return err
			}
		}
		end, err := decoder.Token()
		if err != nil || end != json.Delim(']') {
			return ErrSingleRequestInternalToolInvalidCall
		}
	default:
		return ErrSingleRequestInternalToolInvalidCall
	}
	return nil
}

func validInternalToolIdentity(value string) bool {
	return value != "" && strings.TrimSpace(value) == value && !strings.ContainsRune(value, 0)
}

func validInternalWorkspacePath(value string) bool {
	if value == "" || len(value) > 4096 || strings.Contains(value, "\\") || strings.ContainsRune(value, 0) ||
		path.IsAbs(value) || path.Clean(value) != value || value == ".iop" || strings.HasPrefix(value, ".iop/") {
		return false
	}
	return value == "." || (value != ".." && !strings.HasPrefix(value, "../"))
}

func validEnvironmentName(value string) bool {
	if value == "" {
		return false
	}
	for index, r := range value {
		if (r >= 'A' && r <= 'Z') || (r >= 'a' && r <= 'z') || r == '_' || (index > 0 && r >= '0' && r <= '9') {
			continue
		}
		return false
	}
	return true
}

func cloneStringMap(input map[string]string) map[string]string {
	if input == nil {
		return nil
	}
	output := make(map[string]string, len(input))
	for key, value := range input {
		output[key] = value
	}
	return output
}

func internalWorkspaceToolResult(response *iop.WorkspaceToolResponse) InternalWorkspaceToolResult {
	return InternalWorkspaceToolResult{
		RequestID:  response.GetRequestId(),
		StageID:    response.GetStageId(),
		ToolCallID: response.GetToolCallId(),
		Status:     internalWorkspaceStatus(response.GetStatus()),
		ErrorCode:  internalWorkspaceErrorCode(response.GetErrorCode()),
		Content:    append([]byte(nil), response.GetContent()...),
		Entries:    append([]string(nil), response.GetEntries()...),
		Stdout:     append([]byte(nil), response.GetStdout()...),
		Stderr:     append([]byte(nil), response.GetStderr()...),
		ExitCode:   response.GetExitCode(),
		Truncated:  response.GetTruncated(),
		DurationMS: response.GetDurationMs(),
	}
}

func internalWorkspaceStatus(status iop.WorkspaceStatus) string {
	switch status {
	case iop.WorkspaceStatus_WORKSPACE_STATUS_SUCCESS:
		return "success"
	case iop.WorkspaceStatus_WORKSPACE_STATUS_ERROR:
		return "error"
	case iop.WorkspaceStatus_WORKSPACE_STATUS_TIMEOUT:
		return "timeout"
	case iop.WorkspaceStatus_WORKSPACE_STATUS_CANCELLED:
		return "cancelled"
	case iop.WorkspaceStatus_WORKSPACE_STATUS_UNSUPPORTED:
		return "unsupported"
	default:
		return "invalid"
	}
}

func internalWorkspaceErrorCode(code iop.WorkspaceErrorCode) string {
	switch code {
	case iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_UNSPECIFIED:
		return ""
	case iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_NOT_READY:
		return "not_ready"
	case iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_UNSUPPORTED:
		return "unsupported"
	case iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_INVALID_REQUEST:
		return "invalid_request"
	case iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_NOT_FOUND:
		return "not_found"
	case iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_TIMEOUT:
		return "timeout"
	case iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_CANCELLED:
		return "cancelled"
	case iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_INTERNAL:
		return "internal"
	default:
		return "invalid"
	}
}

func internalWorkspaceToolOutputBytes(result InternalWorkspaceToolResult) int {
	total := len(result.Content) + len(result.Stdout) + len(result.Stderr)
	for _, entry := range result.Entries {
		total += len(entry)
	}
	return total
}
