package software.amazon.smithy.go.codegen;

import static software.amazon.smithy.go.codegen.GoWriter.emptyGoTemplate;
import static software.amazon.smithy.go.codegen.GoWriter.goTemplate;

import java.util.Map;
import java.util.Optional;
import software.amazon.smithy.go.codegen.util.ShapeUtil;
import software.amazon.smithy.model.loader.Prelude;
import software.amazon.smithy.model.shapes.MemberShape;
import software.amazon.smithy.model.shapes.ServiceShape;
import software.amazon.smithy.model.shapes.Shape;
import software.amazon.smithy.model.shapes.ShapeId;
import software.amazon.smithy.model.traits.Trait;
import software.amazon.smithy.utils.SmithyInternalApi;
import software.amazon.smithy.utils.StringUtils;

@SmithyInternalApi
public class SchemaGenerator implements Writable {
    private final GoCodegenContext ctx;
    private final Shape shape;

    public SchemaGenerator(GoCodegenContext ctx, Shape shape) {
        this.ctx = ctx;
        this.shape = shape;
    }

    public static String getSchemaName(Shape shape, ServiceShape service) {
        if (Prelude.isPublicPreludeShape(shape)) {
            return "smithyprelude." + shape.getId().getName();
        }
        var name = modelShapeID(shape).getName(service);
        return ShapeUtil.isExported(shape)
                ? name
                : "_" + name;
    }

    public static String getMemberSchemaName(Shape shape, MemberShape member, ServiceShape service) {
        return String.format("%s_%s", getSchemaName(shape, service), member.getMemberName());
    }

    // modelShapeID returns the original model shape ID, preferring the
    // Synthetic trait's archetype (if any) so that synthetic clones of Input /
    // Output inline shapes carry the Smithy model's shape ID rather than the
    // Go-codegen synthetic ID.
    private static ShapeId modelShapeID(Shape shape) {
        Optional<Synthetic> synth = shape.getTrait(Synthetic.class);
        if (synth.isPresent() && synth.get().getArchetype().isPresent()) {
            return synth.get().getArchetype().get();
        }
        return shape.getId();
    }

    // Returns the schema reference for use outside the schemas package (e.g. in
    // serde generators). Prelude shapes are already package-qualified, local
    // shapes get the "schemas." prefix. Stub synthetics (no archetype) return
    // "nil" since no schema is generated for them.
    public static String getSchemaRef(Shape shape, ServiceShape service) {
        if (CodegenUtils.isStubSynthetic(shape)) {
            return "nil";
        }
        if (Prelude.isPublicPreludeShape(shape)) {
            return getSchemaName(shape, service);
        }
        return "schemas." + getSchemaName(shape, service);
    }

    public static String getMemberSchemaRef(Shape shape, MemberShape member, ServiceShape service) {
        if (Prelude.isPublicPreludeShape(shape)) {
            return getSchemaName(shape, service) + "_" + member.getMemberName();
        }
        return "schemas." + getMemberSchemaName(shape, member, service);
    }

    @Override
    public void accept(GoWriter writer) {
        if (Prelude.isPublicPreludeShape(shape)) {
            return;
        }
        var service = ctx.service();
        var modelID = modelShapeID(shape);
        writer.addUseImports(SmithyGoDependency.SMITHY);
        writer.writeGoTemplate("""
                        var $ident:L = smithy.NewSchema(smithy.ShapeID{
                            Namespace: $namespace:S,
                            Name: $name:S,
                        }, smithy.ShapeType$type:L, $numMembers:L$traits:W)
                        $memberVarDecls:W
                        """,
                Map.of(
                        "ident", getSchemaName(shape, service),
                        "namespace", modelID.getNamespace(),
                        "name", modelID.getName(),
                        "type", StringUtils.capitalize(shape.getType().toString()),
                        "numMembers", shape.members().size(),
                        "traits", renderVariadicTraits(shape.getAllTraits().values()),
                        "memberVarDecls", Writable.map(shape.members(), this::renderMemberVarDecl, true)
                ));
    }

    public void acceptMembersInit(GoWriter writer) {
        if (shape.members().isEmpty() || Prelude.isPublicPreludeShape(shape)) {
            return;
        }
        writer.writeGoTemplate("""
                        $memberAddAndAssign:W
                        """,
                Map.of(
                        "memberAddAndAssign", Writable.map(shape.members(), this::renderMemberAddAndAssign, true)
                ));
    }

    private Writable renderMemberVarDecl(MemberShape member) {
        return goTemplate("""
                        var $ident:L *smithy.Schema
                        """,
                Map.of("ident", getMemberSchemaName(shape, member, ctx.service())));
    }

    private Writable renderMemberAddAndAssign(MemberShape member) {
        var target = ctx.model().expectShape(member.getTarget());
        var service = ctx.service();
        var memberTraits = member.getAllTraits().values();

        if (Prelude.isPublicPreludeShape(target)) {
            return goTemplate("""
                            $preludeImport:D$ident:L = $schema:L.AddMember($name:S, $target:L$traits:W)
                            """,
                    Map.of(
                            "preludeImport", SmithyGoDependency.SMITHY_PRELUDE,
                            "ident", getMemberSchemaName(shape, member, service),
                            "schema", getSchemaName(shape, service),
                            "name", member.getMemberName(),
                            "target", getSchemaName(target, service),
                            "traits", renderVariadicTraits(memberTraits)
                    ));
        }

        return goTemplate("""
                        $ident:L = $schema:L.AddMember($name:S, $target:L$traits:W)
                        """,
                Map.of(
                        "ident", getMemberSchemaName(shape, member, service),
                        "schema", getSchemaName(shape, service),
                        "name", member.getMemberName(),
                        "target", getSchemaName(target, service),
                        "traits", renderVariadicTraits(memberTraits)
                ));
    }

    private Writable renderVariadicTraits(java.util.Collection<Trait> traits) {
        if (traits.isEmpty()) {
            return emptyGoTemplate();
        }
        return goTemplate(", $W", Writable.map(traits, this::renderVariadicTrait));
    }

    private Writable renderTraitMapEntry(Trait trait) {
        var generator = DefaultTraitGenerators.forTrait(trait.toShapeId());
        return generator != null
                ? goTemplate("$S: $W,", trait.toShapeId().toString(), generator.render(trait))
                : emptyGoTemplate();
    }

    private Writable renderVariadicTrait(Trait trait) {
        var generator = DefaultTraitGenerators.forTrait(trait.toShapeId());
        return generator != null
                ? goTemplate("$W,", generator.render(trait))
                : emptyGoTemplate();
    }
}
