"""
Golang code generator for LinkML schemas.

Based on the PydanticGenerator architecture.
"""

import logging
import os
import re
from dataclasses import dataclass, field
from pathlib import Path

import click
from jinja2 import ChoiceLoader, Environment, FileSystemLoader

from linkml._version import __version__
from linkml.generators.golanggen.build import FieldResult, StructResult
from linkml.generators.golanggen.template import (
    GolangConstant,
    GolangEnum,
    GolangField,
    GolangModule,
    GolangStruct,
    Import,
    Imports,
)
from linkml.generators.oocodegen import OOCodeGenerator
from linkml.utils.generator import shared_arguments
from linkml_runtime.linkml_model.meta import ClassDefinition, EnumDefinition, SlotDefinition
from linkml_runtime.utils.formatutils import camelcase, underscore
from linkml_runtime.utils.schemaview import SchemaView

logger = logging.getLogger(__name__)


# Type mapping from LinkML base types to Go types
TYPE_MAP = {
    "str": "string",
    "string": "string",
    "int": "int",
    "integer": "int",
    "Bool": "bool",
    "boolean": "bool",
    "Decimal": "float64",
    "decimal": "float64",
    "float": "float64",
    "double": "float64",
    "XSDDate": "time.Time",
    "XSDDateTime": "time.Time",
    "date": "time.Time",
    "datetime": "time.Time",
    "time": "time.Time",
    "uri": "string",
    "uriorcurie": "string",
    "ncname": "string",
}

# Go primitive types eligible for pointer promotion via ``nullable_primitives``
GO_PRIMITIVE_TYPES = frozenset(
    {
        "string",
        "int",
        "int8",
        "int16",
        "int32",
        "int64",
        "float32",
        "float64",
        "bool",
    }
)


@dataclass
class GolangGenerator(OOCodeGenerator):
    """
    Generates Golang code from a LinkML schema.

    This generator creates idiomatic Go structs with JSON tags from LinkML class definitions.
    It supports:
    - Struct generation from classes
    - Const blocks from enums
    - Inheritance via struct embedding
    - Proper type mappings
    - Multivalued fields as slices
    - Required vs optional fields
    """

    # ClassVars
    generatorname = os.path.basename(__file__)
    generatorversion = "0.2.0"
    valid_formats = ["go", "golang"]
    file_extension = "go"

    # ObjectVars
    package_name: str | None = None
    """Override the package name. If None, derived from schema name."""

    gen_slots: bool = True
    """Generate struct fields"""

    sort_imports: bool = True
    """Sort imports by group (stdlib, thirdparty, local)"""

    add_json_tags: bool = True
    """Add JSON struct tags to fields"""

    use_time_package: bool = True
    """Import time package when needed for date/datetime types"""

    alphabetical_sort: bool = False
    """Sort structs and enums alphabetically by name for deterministic output."""

    nullable_primitives: bool = True
    """
    Use pointer types for optional primitive fields (e.g. ``*string``, ``*int``, ``*bool``).

    When True (default), optional primitives and ``time.Time`` fields become
    pointers so that ``omitempty`` only omits ``nil``, preserving intentional
    zero values.

    When False, optional primitives and ``time.Time`` use bare Go types.
    The ``omitzero`` JSON struct tag (Go 1.24+) is added alongside
    ``omitempty`` so that zero values — including ``time.Time{}`` — are
    correctly omitted during JSON marshalling.
    """

    named_slot_types: bool = False
    """
    Generate named Go types for parent slots that have ``is_a`` children.

    When enabled, a slot like ``signature`` with children ``owner_signature``
    and ``witness_signature`` produces ``type Signature string`` and all child
    slots use ``Signature`` instead of the bare ``string``.  This adds
    compile-time type safety between slot families.  Named slot types follow
    the same pointer rules as regular primitives.
    """

    template_dir: str | Path | None = None
    """
    Override templates for each GolangTemplateModel.

    Directory with templates that override the default templates.
    If a matching template is not found in the override directory,
    the default templates will be used.
    """

    # Private attributes
    _needs_time: bool = field(default=False, init=False)
    _type_defs: dict[str, tuple[str, str | None]] = field(default_factory=dict, init=False)
    _parent_slot_names: set[str] = field(default_factory=set, init=False)

    # Go zero-value defaults by type
    GO_TYPE_DEFAULTS = {
        "string": '""',
        "int": "0",
        "int8": "0",
        "int16": "0",
        "int32": "0",
        "int64": "0",
        "float32": "0.0",
        "float64": "0.0",
        "bool": "false",
        "time.Time": "time.Time{}",
    }

    def __post_init__(self):
        super().__post_init__()

    def default_value_for_type(self, typ: str) -> str:
        """
        Return the Go zero value for a given type.

        Args:
            typ: The Go type string

        Returns:
            The zero value literal for that type
        """
        return self.GO_TYPE_DEFAULTS.get(typ, "nil")

    def _resolve_underlying_type(self, go_type: str) -> str:
        """
        Resolve a Go type through type definitions to get the underlying primitive type.

        Named slot types (e.g. ``Signature``) and reference types (e.g. ``PersonId``)
        are backed by a primitive.  This method returns the underlying type so callers
        can check whether it is a primitive or ``time.Time``.

        Args:
            go_type: The Go type string, possibly a named type.

        Returns:
            The underlying Go type (e.g. ``"string"``, ``"int"``, ``"time.Time"``).
        """
        td = self._type_defs.get(go_type)
        return td[0] if td else go_type

    def _reference_type_for_class(self, class_name: str) -> str:
        """
        Return a Go type name for referencing a class by its identifier.

        When a slot with a class range is not inlined, the serialised form
        contains only the identifier value, not the full object.  This method
        creates a named Go type definition (e.g. ``type PersonId string``)
        and returns its name so it can be used as the field type.

        Args:
            class_name: The LinkML class name whose identifier type is needed.

        Returns:
            The Go type-definition name (e.g. ``"PersonId"``).
        """
        sv = self.schemaview
        id_slot = sv.get_identifier_slot(class_name)
        if id_slot is None:
            logger.warning(f"No identifier slot for {class_name}, falling back to string")
            go_type = "string"
        elif id_slot.range and id_slot.range in sv.all_types():
            t = sv.get_type(id_slot.range)
            if t.base and t.base in TYPE_MAP:
                go_type = TYPE_MAP[t.base]
            else:
                go_type = "string"
        else:
            go_type = "string"

        alias_name = f"{camelcase(class_name)}Id"
        self._type_defs[alias_name] = (go_type, None)
        return alias_name

    def _named_type_for_slot(self, slot_name: str) -> str:
        """
        Return a named Go type for a parent slot that has ``is_a`` children.

        Creates a type definition (e.g. ``type Signature string``) and returns
        its name so child slots can use it instead of the bare primitive.  The
        definition is cached in :attr:`_type_defs` so each parent slot produces
        at most one type definition.

        Args:
            slot_name: The LinkML slot name (must be a parent slot).

        Returns:
            The Go type-definition name (e.g. ``"Signature"``).
        """
        type_name = camelcase(slot_name)
        if type_name in self._type_defs:
            return type_name

        sv = self.schemaview
        slot = sv.get_slot(slot_name)
        slot_range = slot.range if slot.range else "string"

        if slot_range in sv.all_types():
            t = sv.get_type(slot_range)
            if t.base and t.base in TYPE_MAP:
                go_type = TYPE_MAP[t.base]
            else:
                go_type = "string"
        else:
            go_type = "string"

        description = slot.description if slot.description else None
        self._type_defs[type_name] = (go_type, description)
        return type_name

    @staticmethod
    def sort_classes(clist: list[ClassDefinition]) -> list[ClassDefinition]:
        """
        Sort classes such that if C is a child of P then C appears after P in the list
        """
        clist = list(clist)
        slist = []  # sorted
        while len(clist) > 0:
            can_add = False
            for i in range(len(clist)):
                candidate = clist[i]
                can_add = False
                if candidate.is_a:
                    candidates = [candidate.is_a] + candidate.mixins
                else:
                    candidates = candidate.mixins

                if not candidates:
                    can_add = True
                else:
                    if set(candidates) <= set([p.name for p in slist]):
                        can_add = True
                if can_add:
                    slist = slist + [candidate]
                    del clist[i]
                    break
            if not can_add:
                raise ValueError(f"could not find suitable element in {clist} that does not ref {slist}")
        return slist

    def generate_struct(self, cls: ClassDefinition) -> StructResult:
        """
        Generate a Go struct from a LinkML class definition.

        Args:
            cls: The class definition to convert

        Returns:
            StructResult containing the generated struct and any imports needed
        """
        struct_name = camelcase(cls.name)

        # Handle embedded structs for inheritance
        embedded_structs = []
        if cls.is_a:
            embedded_structs.append(camelcase(cls.is_a))
        if cls.mixins:
            embedded_structs.extend([camelcase(m) for m in cls.mixins])

        go_struct = GolangStruct(
            name=struct_name,
            description=cls.description if cls.description else None,
            embedded_structs=embedded_structs if embedded_structs else None,
        )

        result = StructResult(struct=go_struct, source=cls)

        # Generate fields — only direct slots, inherited ones come via embedding
        if self.gen_slots:
            direct = bool(embedded_structs)
            slots = [
                self.schemaview.induced_slot(sn, cls.name)
                for sn in self.schemaview.class_slots(cls.name, direct=direct)
            ]

            field_results = []
            for slot in slots:
                field_result = self.generate_field(slot, cls)
                field_results.append(field_result)
                result = result.merge(field_result)

            fields = {fr.field.name: fr.field for fr in field_results}
            result.struct.fields = fields if fields else None

        return result

    def generate_field(self, slot: SlotDefinition, cls: ClassDefinition) -> FieldResult:
        """
        Generate a Go struct field from a LinkML slot definition.

        Args:
            slot: The slot definition to convert
            cls: The parent class definition

        Returns:
            FieldResult containing the generated field
        """
        slot_alias = slot.alias if slot.alias else slot.name
        go_name = camelcase(slot_alias)
        json_name = underscore(slot_alias)

        go_type = self.generate_go_type(slot, cls)

        # When pointers are not used (nullable_primitives=False), add the
        # ``omitzero`` JSON tag for optional primitive and time.Time fields.
        # ``omitzero`` (Go 1.24+) correctly omits zero values including
        # ``time.Time{}`` which ``omitempty`` alone does not handle.
        omitzero = False
        if not self.nullable_primitives:
            is_optional = not ((slot.required or False) or (slot.identifier or False) or (slot.key or False))
            if is_optional and not slot.multivalued:
                underlying = self._resolve_underlying_type(go_type)
                if underlying in GO_PRIMITIVE_TYPES or underlying == "time.Time":
                    omitzero = True

        field = GolangField(
            name=slot.name,
            go_name=go_name,
            json_name=json_name,
            type=go_type,
            required=slot.required or False,
            identifier=slot.identifier or False,
            key=slot.key or False,
            description=slot.description if slot.description else None,
            pattern=slot.pattern if slot.pattern else None,
            omitzero=omitzero,
        )

        result = FieldResult(field=field, source=slot)
        return result

    def generate_go_type(self, slot: SlotDefinition, cls: ClassDefinition) -> str:
        """
        Generate the Go type for a slot.

        Args:
            slot: The slot definition
            cls: The parent class definition

        Returns:
            The Go type string (e.g., "string", "[]Person", "*Address")
        """
        sv = self.schemaview
        slot_range = slot.range

        # Determine base type
        if slot_range in sv.all_classes():
            if sv.is_inlined(slot):
                # Inlined: embed the full struct
                base_type = camelcase(slot_range)
            else:
                # Referenced: use the identifier type
                base_type = self._reference_type_for_class(slot_range)
            # Use pointer for optional complex types (single-valued only;
            # slices are already nil-able and the pointer is added before
            # the multivalued check below)
            if not slot.multivalued and not (slot.required or slot.identifier or slot.key):
                base_type = f"*{base_type}"
        elif slot_range in sv.all_enums():
            # Reference to enum
            base_type = camelcase(slot_range)
        elif slot_range in sv.all_types():
            # Built-in or custom type
            t = sv.get_type(slot_range)
            if t.base and t.base in TYPE_MAP:
                base_type = TYPE_MAP[t.base]
                # Track if we need time package
                if base_type == "time.Time":
                    self._needs_time = True
            else:
                logger.warning(f"Unknown type base: {t.name}, defaulting to string")
                base_type = "string"
        else:
            # Default to string
            base_type = "string"

        # If the slot's is_a parent is a parent slot, use the named type.
        # If the slot itself is a parent slot, also use its own named type.
        if slot.is_a and slot.is_a in self._parent_slot_names:
            base_type = self._named_type_for_slot(slot.is_a)
        elif slot.name in self._parent_slot_names:
            base_type = self._named_type_for_slot(slot.name)

        # Resolve the underlying Go type for pointer-promotion checks.
        # Named slot types backed by primitives follow the same rules.
        underlying_type = self._resolve_underlying_type(base_type)

        # Pointer promotion for optional fields.
        #
        # When ``nullable_primitives`` is True (default), optional primitives
        # and ``time.Time`` become pointers so ``omitempty`` only omits
        # ``nil``, preserving intentional zero values.
        #
        # When ``nullable_primitives`` is False, bare types are used instead
        # and the ``omitzero`` JSON tag (Go 1.24+) is added in
        # :meth:`generate_field` to correctly omit zero values including
        # ``time.Time{}``.
        is_optional = not (slot.required or slot.identifier or slot.key)
        if self.nullable_primitives and not slot.multivalued and is_optional:
            if underlying_type == "time.Time" or underlying_type in GO_PRIMITIVE_TYPES:
                base_type = f"*{base_type}"

        # Handle multivalued slots
        if slot.multivalued:
            return f"[]{base_type}"

        return base_type

    def generate_enum(self, enum: EnumDefinition) -> GolangEnum:
        """
        Generate a Go const block from a LinkML enum definition.

        Args:
            enum: The enum definition to convert

        Returns:
            GolangEnum model
        """
        enum_name = camelcase(enum.name)
        constants = {}

        if enum.permissible_values:
            for pv_name, pv in enum.permissible_values.items():
                const_name = camelcase(pv_name)
                # Use the text value if available, otherwise the name
                const_value = pv.text if pv.text else pv_name

                constants[pv_name] = GolangConstant(
                    name=f"{enum_name}{const_name}",
                    value=const_value,
                    description=pv.description if pv.description else None,
                )

        return GolangEnum(
            name=enum_name,
            type="string",  # Default to string type for enums
            description=enum.description if enum.description else None,
            values=constants,
        )

    def render(self) -> GolangModule:
        """
        Render the schema to a GolangModule model.

        Returns:
            GolangModule containing all generated structs and enums
        """
        sv: SchemaView = self.schemaview
        self._type_defs = {}
        self._needs_time = False

        # Precompute which slots have is_a children (parent slots get named types)
        self._parent_slot_names = set()
        if self.named_slot_types:
            for slot_name in sv.all_slots():
                if sv.slot_children(slot_name):
                    self._parent_slot_names.add(slot_name)

        # Determine package name
        package_name = self.package_name
        if package_name is None:
            schema_name = sv.schema.name
            # Extract package name from schema name (e.g., "kitchen_sink" -> "kitchen")
            if "_" in schema_name:
                package_name = schema_name[: schema_name.find("_")]
            else:
                package_name = schema_name
            # Make it valid Go package name (lowercase, alphanumeric + underscore)
            package_name = re.sub(r"[^a-z0-9_]", "", package_name.lower())

        # Generate enums
        enums = {}
        for enum_def in sv.all_enums().values():
            enum_model = self.generate_enum(enum_def)
            enums[enum_model.name] = enum_model

        # Generate structs
        structs = {}
        all_classes = list(sv.all_classes().values())
        sorted_classes = self.sort_classes(all_classes)

        for cls in sorted_classes:
            struct_result = self.generate_struct(cls)
            structs[struct_result.struct.name] = struct_result.struct

        # Add type definitions for referenced (non-inlined) class identifiers
        # and named slot types.
        for td_name, (go_type, description) in self._type_defs.items():
            structs[td_name] = GolangStruct(
                name=td_name,
                description=description,
                is_type_alias=True,
                type_alias_value=go_type,
            )

        # Sort alphabetically for deterministic output
        if self.alphabetical_sort:
            enums = dict(sorted(enums.items()))
            structs = dict(sorted(structs.items()))

        # Build imports
        imports = Imports()
        if self._needs_time and self.use_time_package:
            imports += Import(module="time")

        if self.sort_imports:
            imports.sort()

        module = GolangModule(
            package_name=package_name,
            imports=imports,
            enums=enums,
            structs=structs,
        )

        return module

    def serialize(self, rendered_module: GolangModule | None = None) -> str:
        """
        Serialize the schema to a Go source code string.

        Args:
            rendered_module: Optional pre-rendered module

        Returns:
            Go source code as a string
        """
        if rendered_module is not None:
            module = rendered_module
        else:
            module = self.render()

        return module.render(environment=self._template_environment())

    def _template_environment(self) -> Environment:
        """Get the Jinja2 template environment.

        If :attr:`template_dir` is set, templates in that directory take
        priority over the built-in defaults via a :class:`jinja2.ChoiceLoader`.
        """
        env = GolangModule.environment()
        if self.template_dir is not None:
            loader = ChoiceLoader([FileSystemLoader(self.template_dir), env.loader])
            env.loader = loader
        return env


_TEMPLATE_NAMES = [
    "module.go.jinja",
    "struct.go.jinja",
    "field.go.jinja",
    "enum.go.jinja",
    "imports.go.jinja",
]


@shared_arguments(GolangGenerator)
@click.option(
    "--package-name",
    type=str,
    help="Override the Go package name (default: derived from schema name)",
)
@click.option(
    "--alphabetical-sort/--no-alphabetical-sort",
    default=False,
    show_default=True,
    help="Sort structs and enums alphabetically for deterministic output.",
)
@click.option(
    "--nullable-primitives/--no-nullable-primitives",
    default=True,
    show_default=True,
    help="Use pointer types for optional primitives so omitempty only triggers on nil.",
)
@click.option(
    "--named-slot-types/--no-named-slot-types",
    default=False,
    show_default=True,
    help="Generate named Go types for parent slots with is_a children for type safety.",
)
@click.option(
    "--template-dir",
    type=click.Path(),
    help="""
Optional jinja2 template directory to use for Go code generation.

Pass a directory containing templates with the same name as any of the default
GolangTemplateModel templates to override them. The given directory will be
searched for matching templates, and use the default templates as a fallback
if an override is not found.

Available templates to override:

\b
"""
    + "\n".join(["- " + name for name in _TEMPLATE_NAMES]),
)
@click.version_option(__version__, "-V", "--version")
@click.command(name="golang")
def cli(
    yamlfile,
    package_name: str | None = None,
    alphabetical_sort: bool = False,
    nullable_primitives: bool = True,
    named_slot_types: bool = False,
    template_dir: str | None = None,
    **args,
):
    """
    Generate Golang structs from a LinkML schema.

    This generator produces idiomatic Go code with:
    - Structs for classes
    - Const blocks for enums
    - JSON tags for serialization
    - Struct embedding for inheritance
    """
    if template_dir is not None:
        if not Path(template_dir).exists():
            raise FileNotFoundError(f"The template directory {template_dir} does not exist!")

    gen = GolangGenerator(
        yamlfile,
        package_name=package_name,
        alphabetical_sort=alphabetical_sort,
        nullable_primitives=nullable_primitives,
        named_slot_types=named_slot_types,
        template_dir=template_dir,
        **args,
    )
    print(gen.serialize())


if __name__ == "__main__":
    cli()
