package templates

import (
	"fmt"
	"go/ast"
	"go/types"
	"maps"
	"reflect"
	"sort"
	"strconv"
	"strings"

	. "github.com/dave/jennifer/jen" //nolint:stylecheck
)

func (ps *parseState) parseGoStruct(t *types.Struct, named *types.Named) (*parsedObjectType, error) {
	spec := &parsedObjectType{
		goType: t,
	}

	if named == nil {
		return nil, fmt.Errorf("struct types must be named")
	}
	spec.name = named.Obj().Name()
	if spec.name == "" {
		return nil, fmt.Errorf("struct types must be named")
	}

	// We don't support extending objects from outside this module, so we will
	// be skipping it. But first we want to verify the user isn't adding methods
	// to it (in which case we error out).
	objectIsDaggerGenerated := ps.isDaggerGenerated(named.Obj())

	goFuncTypes := []*types.Func{}
	methodSet := types.NewMethodSet(types.NewPointer(named))
	for i := 0; i < methodSet.Len(); i++ {
		methodObj := methodSet.At(i).Obj()

		if ps.isDaggerGenerated(methodObj) {
			// We don't care about pre-existing methods on core types or objects from dependency modules.
			continue
		}
		if objectIsDaggerGenerated {
			return nil, fmt.Errorf("cannot define methods on objects from outside this module")
		}

		goFuncType, ok := methodObj.(*types.Func)
		if !ok {
			return nil, fmt.Errorf("expected method to be a func, got %T", methodObj)
		}

		if !goFuncType.Exported() {
			continue
		}

		goFuncTypes = append(goFuncTypes, goFuncType)
	}
	if objectIsDaggerGenerated {
		return nil, nil
	}
	sort.Slice(goFuncTypes, func(i, j int) bool {
		return goFuncTypes[i].Pos() < goFuncTypes[j].Pos()
	})

	for _, goFuncType := range goFuncTypes {
		funcTypeSpec, err := ps.parseGoFunc(named, goFuncType)
		if err != nil {
			return nil, fmt.Errorf("failed to parse method %s: %w", goFuncType.Name(), err)
		}
		spec.methods = append(spec.methods, funcTypeSpec)
	}

	// get the comment above the struct (if any)
	astSpec, err := ps.astSpecForNamedType(named)
	if err != nil {
		return nil, fmt.Errorf("failed to find decl for named type %s: %w", spec.name, err)
	}
	spec.doc = astSpec.Doc.Text()

	astStructType, ok := astSpec.Type.(*ast.StructType)
	if !ok {
		return nil, fmt.Errorf("expected type spec to be a struct, got %T", astSpec.Type)
	}

	// Fill out the static fields of the struct (if any)
	astFields := unpackASTFields(astStructType.Fields)
	for i := 0; i < t.NumFields(); i++ {
		field := t.Field(i)
		if !field.Exported() {
			continue
		}

		fieldSpec := &fieldSpec{goType: field.Type()}
		if _, optional, err := ps.isOptionalWrapper(fieldSpec.goType); err != nil {
			return nil, err
		} else if optional {
			return nil, fmt.Errorf("optional type wrapper not allowed in struct field %s", field.Name())
		}
		fieldSpec.typeSpec, err = ps.parseGoTypeReference(fieldSpec.goType, nil, false)
		if err != nil {
			return nil, fmt.Errorf("failed to parse field type: %w", err)
		}

		fieldSpec.goName = field.Name()
		fieldSpec.name = fieldSpec.goName

		// override the name with the json tag if it was set - otherwise, we
		// end up asking for a name that we won't unmarshal correctly
		tag := reflect.StructTag(t.Tag(i))
		if dt := tag.Get("json"); dt != "" {
			dt, _, _ = strings.Cut(dt, ",")
			if dt == "-" {
				continue
			}
			fieldSpec.name = dt
		}

		docPragmas, docComment := parsePragmaComment(astFields[i].Doc.Text())
		linePragmas, lineComment := parsePragmaComment(astFields[i].Comment.Text())
		comment := strings.TrimSpace(docComment)
		if comment == "" {
			comment = strings.TrimSpace(lineComment)
		}
		pragmas := make(map[string]string)
		maps.Copy(pragmas, docPragmas)
		maps.Copy(pragmas, linePragmas)
		if v, ok := pragmas["private"]; ok {
			if v == "" {
				fieldSpec.isPrivate = true
			} else {
				fieldSpec.isPrivate, _ = strconv.ParseBool(v)
			}
		}

		fieldSpec.doc = comment

		spec.fields = append(spec.fields, fieldSpec)
	}

	if ps.isMainModuleObject(spec.name) && ps.constructor != nil {
		spec.constructor, err = ps.parseGoFunc(nil, ps.constructor)
		if err != nil {
			return nil, fmt.Errorf("failed to parse constructor: %w", err)
		}
	}

	return spec, nil
}

type parsedObjectType struct {
	name string
	doc  string

	fields      []*fieldSpec
	methods     []*funcTypeSpec
	constructor *funcTypeSpec

	goType *types.Struct
}

var _ NamedParsedType = &parsedObjectType{}

func (spec *parsedObjectType) TypeDefCode() (*Statement, error) {
	withObjectArgsCode := []Code{
		Lit(spec.name),
	}
	withObjectOptsCode := []Code{}
	if spec.doc != "" {
		withObjectOptsCode = append(withObjectOptsCode, Id("Description").Op(":").Lit(strings.TrimSpace(spec.doc)))
	}
	if len(withObjectOptsCode) > 0 {
		withObjectArgsCode = append(withObjectArgsCode, Id("TypeDefWithObjectOpts").Values(withObjectOptsCode...))
	}

	typeDefCode := Qual("dag", "TypeDef").Call().Dot("WithObject").Call(withObjectArgsCode...)

	for _, method := range spec.methods {
		fnTypeDefCode, err := method.TypeDefCode()
		if err != nil {
			return nil, fmt.Errorf("failed to convert method %s to function def: %w", method.name, err)
		}
		typeDefCode = dotLine(typeDefCode, "WithFunction").Call(Add(Line(), fnTypeDefCode))
	}

	for _, field := range spec.fields {
		if field.isPrivate {
			continue
		}

		fieldTypeDefCode, err := field.typeSpec.TypeDefCode()
		if err != nil {
			return nil, fmt.Errorf("failed to convert field type: %w", err)
		}
		withFieldArgsCode := []Code{
			Lit(field.name),
			fieldTypeDefCode,
		}
		if field.doc != "" {
			withFieldArgsCode = append(withFieldArgsCode,
				Id("TypeDefWithFieldOpts").Values(
					Id("Description").Op(":").Lit(field.doc),
				))
		}
		typeDefCode = dotLine(typeDefCode, "WithField").Call(withFieldArgsCode...)
	}

	if spec.constructor != nil {
		fnTypeDefCode, err := spec.constructor.TypeDefCode()
		if err != nil {
			return nil, fmt.Errorf("failed to convert constructor to function def: %w", err)
		}
		typeDefCode = dotLine(typeDefCode, "WithConstructor").Call(Add(Line(), fnTypeDefCode))
	}

	return typeDefCode, nil
}

func (spec *parsedObjectType) GoType() types.Type {
	return spec.goType
}

func (spec *parsedObjectType) GoSubTypes() []types.Type {
	var subTypes []types.Type
	for _, method := range spec.methods {
		subTypes = append(subTypes, method.GoSubTypes()...)
	}
	for _, field := range spec.fields {
		if field.isPrivate {
			continue
		}
		subTypes = append(subTypes, field.typeSpec.GoSubTypes()...)
	}
	if spec.constructor != nil {
		subTypes = append(subTypes, spec.constructor.GoSubTypes()...)
	}
	return subTypes
}

func (spec *parsedObjectType) Name() string {
	return spec.name
}

/*
Extra generated code needed for the object implementation.

Right now, this is just an UnmarshalJSON method. This is needed because objects may have fields
of an interface type, which the JSON unmarshaller can't handle on its own. Instead, this custom
unmarshaller will unmarshal the JSON into a struct where all the fields are concrete types,
including the underlying concrete struct implementation of any interface fields.

After it unmarshals into that, it copies the fields to the real object fields, handling any
special cases around interface conversions (e.g. converting a slice of structs to a slice of
interfaces).

e.g.:

	func (r *Test) UnmarshalJSON(bs []byte) error {
		var concrete struct {
			Iface          *customIfaceImpl
			IfaceList      []*customIfaceImpl
			OtherIfaceList []*otherIfaceImpl
		}
		err := json.Unmarshal(bs, &concrete)
		if err != nil {
			return err
		}
		r.Iface = concrete.Iface.toIface()
		r.IfaceList = convertSlice(concrete.IfaceList, (*customIfaceImpl).toIface)
		r.OtherIfaceList = convertSlice(concrete.OtherIfaceList, (*otherIfaceImpl).toIface)
		return nil
	}
*/
func (spec *parsedObjectType) ImplementationCode() (*Statement, error) {
	concreteFields := make([]Code, 0, len(spec.fields))
	setFieldCodes := make([]*Statement, 0, len(spec.fields))
	for _, field := range spec.fields {
		fieldTypeCode, err := spec.concreteFieldTypeCode(field.typeSpec)
		if err != nil {
			return nil, fmt.Errorf("failed to generate field type code: %w", err)
		}
		fieldCode := Id(field.goName).Add(fieldTypeCode)
		if field.goName != field.name {
			fieldCode.Tag(map[string]string{"json": field.name})
		}
		concreteFields = append(concreteFields, fieldCode)

		setFieldCode, err := spec.setFieldsFromConcreteStructCode(field)
		if err != nil {
			return nil, fmt.Errorf("failed to generate set field code: %w", err)
		}
		setFieldCodes = append(setFieldCodes, setFieldCode)
	}

	return Func().Params(Id("r").Op("*").Id(spec.name)).
		Id("UnmarshalJSON").
		Params(Id("bs").Id("[]byte")).
		Params(Id("error")).
		BlockFunc(func(g *Group) {
			g.Var().Id("concrete").Struct(concreteFields...)
			g.Id("err").Op(":=").Id("json").Dot("Unmarshal").Call(Id("bs"), Op("&").Id("concrete"))
			g.If(Id("err").Op("!=").Nil()).Block(Return(Id("err")))

			for _, setFieldCode := range setFieldCodes {
				g.Add(setFieldCode)
			}

			g.Return(Nil())
		}), nil
}

/*
The code for the type of a field in the concrete struct unmarshalled into. Mainly needs to handle
interface types, which need to be converted to their concrete struct implementations.
*/
func (spec *parsedObjectType) concreteFieldTypeCode(typeSpec ParsedType) (*Statement, error) {
	s := Empty()
	switch typeSpec := typeSpec.(type) {
	case *parsedPrimitiveType:
		if typeSpec.isPtr {
			s.Op("*")
		}
		if typeSpec.alias != "" {
			s.Id(typeSpec.alias)
		} else {
			s.Id(typeSpec.GoType().String())
		}

	case *parsedSliceType:
		fieldTypeCode, err := spec.concreteFieldTypeCode(typeSpec.underlying)
		if err != nil {
			return nil, fmt.Errorf("failed to generate slice field type code: %w", err)
		}
		s.Index().Add(fieldTypeCode)

	case *parsedObjectTypeReference:
		if typeSpec.isPtr {
			s.Op("*")
		}
		s.Id(typeSpec.name)

	case *parsedIfaceTypeReference:
		s.Op("*").Id(formatIfaceImplName(typeSpec.name))

	default:
		return nil, fmt.Errorf("unsupported concrete field type %T", typeSpec)
	}

	return s, nil
}

/*
The code for setting the fields of the real object from the concrete struct unmarshalled into. e.g.:

	r.Iface = concrete.Iface.toIface()
	r.IfaceList = convertSlice(concrete.IfaceList, (*customIfaceImpl).toIface)
*/
func (spec *parsedObjectType) setFieldsFromConcreteStructCode(field *fieldSpec) (*Statement, error) {
	s := Empty()
	switch typeSpec := field.typeSpec.(type) {
	case *parsedPrimitiveType, *parsedObjectTypeReference:
		s.Id("r").Dot(field.goName).Op("=").Id("concrete").Dot(field.goName)

	case *parsedSliceType:
		switch underlyingTypeSpec := typeSpec.underlying.(type) {
		case *parsedIfaceTypeReference:
			s.Id("r").Dot(field.goName).Op("=").Id("convertSlice").Call(
				Id("concrete").Dot(field.goName),
				Parens(Op("*").Id(formatIfaceImplName(underlyingTypeSpec.name))).Dot("toIface"),
			)
		default:
			s.Id("r").Dot(field.goName).Op("=").Id("concrete").Dot(field.goName)
		}

	case *parsedIfaceTypeReference:
		s.Id("r").Dot(field.goName).Op("=").Id("concrete").Dot(field.goName).Dot("toIface").Call()

	default:
		return nil, fmt.Errorf("unsupported field type %T", typeSpec)
	}

	return s, nil
}

type fieldSpec struct {
	name     string
	doc      string
	typeSpec ParsedType

	// isPrivate is true if the field is marked with the +private pragma
	isPrivate bool
	// goName is the name of the field in the Go struct. It may be different than name if the user changed the name of the field via a json tag
	goName string

	goType types.Type
}
