package logging

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

	"google.golang.org/grpc"
	"google.golang.org/grpc/status"
	"google.golang.org/protobuf/encoding/protojson"
	"google.golang.org/protobuf/proto"
	"google.golang.org/protobuf/reflect/protoreflect"
)

func UnaryClientInterceptor(local, remote string) grpc.UnaryClientInterceptor {
	return func(ctx context.Context, method string, req, reply any, cc *grpc.ClientConn, invoke grpc.UnaryInvoker, opts ...grpc.CallOption) error {
		attrs := rpcIdentity(method, req)
		started := time.Now()
		slog.InfoContext(ctx, "RPC request", append(attrs, "direction", local+"->"+remote, "stage", "invoke")...)
		rpcBody(ctx, "RPC request data", req, append(attrs, "direction", local+"->"+remote, "stage", "invoke")...)
		err := invoke(ctx, method, req, reply, cc, opts...)
		result := append(attrs, "direction", remote+"->"+local, "duration_ms", time.Since(started).Milliseconds(), "rpc_status", status.Code(err).String())
		if err != nil {
			slog.ErrorContext(ctx, "RPC failed", append(result, "stage", "response_failed", "error", err)...)
			return err
		}
		slog.InfoContext(ctx, "RPC response", append(result, "stage", "response_received")...)
		rpcBody(ctx, "RPC response data", reply, append(result, "stage", "response_received")...)
		return nil
	}
}

func UnaryServerInterceptor(local, remote string) grpc.UnaryServerInterceptor {
	return func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handle grpc.UnaryHandler) (any, error) {
		attrs := rpcIdentity(info.FullMethod, req)
		started := time.Now()
		slog.InfoContext(ctx, "RPC request", append(attrs, "direction", remote+"->"+local, "stage", "request_received")...)
		rpcBody(ctx, "RPC request data", req, append(attrs, "direction", remote+"->"+local, "stage", "request_received")...)
		reply, err := handle(ctx, req)
		result := append(attrs, "direction", local+"->"+remote, "duration_ms", time.Since(started).Milliseconds(), "rpc_status", status.Code(err).String())
		if err != nil {
			slog.ErrorContext(ctx, "RPC failed", append(result, "stage", "response_failed", "error", err)...)
			return reply, err
		}
		// Returning from a handler is not evidence of delivery to the caller.
		slog.InfoContext(ctx, "RPC response", append(result, "stage", "response_ready")...)
		rpcBody(ctx, "RPC response data", reply, append(result, "stage", "response_ready")...)
		return reply, nil
	}
}

func rpcIdentity(method string, v any) []any {
	attrs := []any{"transport", "grpc", "method", method}
	p, ok := v.(proto.Message)
	if !ok || p == nil {
		return attrs
	}
	var collect func(protoreflect.Message)
	collect = func(m protoreflect.Message) {
		m.Range(func(f protoreflect.FieldDescriptor, v protoreflect.Value) bool {
			switch string(f.Name()) {
			case "request_id", "trace_id", "operation_id", "dispatcher_id", "tenant_id", "agent_id", "cell_id", "boot_id", "source_event_id", "task_id", "call_id", "execution_id":
				if !f.IsList() && !f.IsMap() {
					attrs = append(attrs, string(f.Name()), v.Interface())
				}
			case "meta":
				if f.Kind() == protoreflect.MessageKind {
					collect(v.Message())
				}
			}
			return true
		})
	}
	collect(p.ProtoReflect())
	return attrs
}

func rpcBody(ctx context.Context, message string, v any, attrs ...any) {
	if !slog.Default().Enabled(ctx, slog.LevelDebug) {
		return
	}
	p, ok := v.(proto.Message)
	if !ok || p == nil {
		return
	}
	body, err := rpcJSON(p)
	if err != nil {
		slog.DebugContext(ctx, message, append(attrs, "redaction_error", err.Error(), "body", "[OMITTED: RPC projection failed]")...)
		return
	}
	Body(ctx, message, body, append(attrs, "wire_encoding", "protobuf", "log_encoding", "protobuf_json_with_embedded_json")...)
}

// JSON payload bytes are observed as JSON instead of opaque base64. All other
// byte fields are represented only by length. No request object is modified.
func rpcJSON(p proto.Message) ([]byte, error) {
	encoded, err := (protojson.MarshalOptions{UseProtoNames: true}).Marshal(p)
	if err != nil {
		return nil, err
	}
	var fields map[string]json.RawMessage
	if err = json.Unmarshal(encoded, &fields); err != nil {
		return nil, err
	}
	var projectionErr error
	p.ProtoReflect().Range(func(f protoreflect.FieldDescriptor, v protoreflect.Value) bool {
		if f.IsList() || f.IsMap() {
			return true
		}
		name := string(f.Name())
		if f.Kind() == protoreflect.BytesKind {
			raw := v.Bytes()
			if strings.HasSuffix(name, "_json") {
				safe, e := RedactJSON(raw)
				if e != nil {
					fields[name], _ = json.Marshal("[OMITTED: invalid embedded JSON]")
					fields[name+"_redaction_error"], _ = json.Marshal(e.Error())
				} else if len(raw) > 0 {
					fields[name] = json.RawMessage(safe)
				}
			} else {
				fields[name], _ = json.Marshal(fmt.Sprintf("[BINARY: %d bytes]", len(raw)))
			}
		} else if f.Kind() == protoreflect.MessageKind {
			nested, e := rpcJSON(v.Message().Interface())
			if e != nil {
				projectionErr = e
				return false
			}
			fields[name] = nested
		}
		return true
	})
	if projectionErr != nil {
		return nil, projectionErr
	}
	return json.Marshal(fields)
}