package protoset

import (
	"context"
	"errors"
	"fmt"
	"sync"

	protocodec "go.unistack.org/micro-codec-proto/v3"
	"google.golang.org/protobuf/reflect/protodesc"
	"google.golang.org/protobuf/reflect/protoreflect"
	"google.golang.org/protobuf/reflect/protoregistry"
	"google.golang.org/protobuf/types/descriptorpb"
	"google.golang.org/protobuf/types/dynamicpb"
)

var errNotFound = errors.New("file descriptor not found")

type ProtoSet struct {
	mu    sync.Mutex
	files map[string]*protoregistry.Files
}

func NewProtoSet() *ProtoSet {
	return &ProtoSet{
		mu:    sync.Mutex{},
		files: make(map[string]*protoregistry.Files, 0),
	}
}

func (p *ProtoSet) GetMessage(addr, pkg, svc, mth string) (protoreflect.Message, protoreflect.Message, error) {
	if addr == "" || svc == "" || mth == "" || pkg == "" {
		return nil, nil, errors.New("addr or service name is empty")
	}
	p.mu.Lock()
	pfile, ok := p.files[addr+"|"+svc]
	p.mu.Unlock()
	if !ok || pfile == nil {
		return nil, nil, errNotFound
	}

	pdesc, err := pfile.FindDescriptorByName(protoreflect.FullName(pkg + "." + svc))
	if err != nil {
		return nil, nil, fmt.Errorf("failed to find service %s.%s, err: %w", pkg, svc, err)
	}

	sdesc, ok := pdesc.(protoreflect.ServiceDescriptor)
	if !ok {
		return nil, nil, fmt.Errorf("failed to find service " + pkg + "." + svc)
	}

	mdesc := sdesc.Methods().ByName(protoreflect.Name(mth))
	if mdesc == nil {
		return nil, nil, fmt.Errorf("unknown method " + mth)
	}

	req := dynamicpb.NewMessageType(mdesc.Input()).New()
	rsp := dynamicpb.NewMessageType(mdesc.Output()).New()

	return req, rsp, nil
}

func (p *ProtoSet) AddProtoset(addr, svc string, data []byte) error {
	fdset := &descriptorpb.FileDescriptorSet{}
	if err := protocodec.NewCodec().Unmarshal(data, fdset); err != nil {
		return fmt.Errorf("failed to unmarshal protoset file: %w", err)
	}

	pfileoptions := protodesc.FileOptions{AllowUnresolvable: true}
	pfiles, err := pfileoptions.NewFiles(fdset)
	if err != nil {
		return fmt.Errorf("failed to use protoset file, err: %w", err)
	}

	p.mu.Lock()
	p.files[addr+"|"+svc] = pfiles
	p.mu.Unlock()

	return nil
}

func (p *ProtoSet) AddReflection(ctx context.Context, service string, addr string) error {
	return nil
}