import json
import sys
import textwrap
from dataclasses import asdict, dataclass
from dataclasses import field as dataclass_field
from enum import Enum

from google.protobuf import descriptor_pb2
from google.protobuf.compiler import plugin_pb2

from mlflow.protos import databricks_pb2


class Visibility(Enum):
    PUBLIC = "public"
    INTERNAL = "internal"
    PUBLIC_UNDOCUMENTED = "public_undocumented"
    PUBLIC_UNDOCUMENTED_READ_ONLY = "public_undocumented_read_only"


@dataclass
class ProtoMessageField:
    description: str
    field_name: str
    field_default: str | None
    entity_type: str
    field_type: str
    full_path: list[str]
    visibility: str
    since: str
    deprecated: bool
    repeated: bool
    validate_required: bool
    oneof: list["ProtoMessageField"] = dataclass_field(default_factory=list)


@dataclass
class ProtoMessage:
    name: str
    full_path: list[str]
    description: str
    visibility: str
    fields: list[ProtoMessageField]
    enums: list["ProtoEnum"]
    messages: list["ProtoMessage"]


@dataclass
class ProtoEnumValue:
    value: str
    full_path: list[str]
    visibility: str
    description: str


@dataclass
class ProtoEnum:
    name: str
    description: str
    full_path: list[str]
    values: list[ProtoEnumValue]
    visibility: str


@dataclass
class DatabricksRpcOptionsDescription:
    path: str | None = None
    method: str | None = None
    visibility: str = "internal"
    since_major: int | None = None
    since_minor: int | None = None
    error_codes: list[int] | None = None
    rpc_doc_title: str = ""

    def __post_init__(self):
        if self.error_codes is None:
            self.error_codes = []


@dataclass
class ProtoServiceMethod:
    name: str
    full_path: list[str]
    request_full_path: list[str]
    response_full_path: list[str]
    description: str
    rpc_options: DatabricksRpcOptionsDescription | None


@dataclass
class ProtoService:
    name: str
    full_path: list[str]
    description: str
    visibility: str
    methods: list[ProtoServiceMethod]


@dataclass
class ProtoTopComment:
    content: str
    visibility: str


@dataclass
class ProtoFileElement:
    comment: ProtoTopComment | None = None
    enum: ProtoEnum | None = None
    message: ProtoMessage | None = None
    service: ProtoService | None = None


@dataclass
class ProtoFile:
    filename: str
    requested_visibility: str
    content: list[ProtoFileElement]


@dataclass
class ProtoAllContent:
    requested_visibility: str
    files: list[ProtoFile]


class ProtobufDocGenerator:
    def get_field_type_name(self, field: descriptor_pb2.FieldDescriptorProto) -> str:
        type_names = {
            descriptor_pb2.FieldDescriptorProto.TYPE_DOUBLE: "DOUBLE",
            descriptor_pb2.FieldDescriptorProto.TYPE_FLOAT: "FLOAT",
            descriptor_pb2.FieldDescriptorProto.TYPE_INT64: "INT64",
            descriptor_pb2.FieldDescriptorProto.TYPE_UINT64: "UINT64",
            descriptor_pb2.FieldDescriptorProto.TYPE_INT32: "INT32",
            descriptor_pb2.FieldDescriptorProto.TYPE_FIXED64: "FIXED64",
            descriptor_pb2.FieldDescriptorProto.TYPE_FIXED32: "FIXED32",
            descriptor_pb2.FieldDescriptorProto.TYPE_BOOL: "BOOL",
            descriptor_pb2.FieldDescriptorProto.TYPE_STRING: "STRING",
            descriptor_pb2.FieldDescriptorProto.TYPE_BYTES: "BYTES",
            descriptor_pb2.FieldDescriptorProto.TYPE_UINT32: "UINT32",
            descriptor_pb2.FieldDescriptorProto.TYPE_SFIXED32: "SFIXED32",
            descriptor_pb2.FieldDescriptorProto.TYPE_SFIXED64: "SFIXED64",
            descriptor_pb2.FieldDescriptorProto.TYPE_SINT32: "SINT32",
            descriptor_pb2.FieldDescriptorProto.TYPE_SINT64: "SINT64",
        }

        if field.type in type_names:
            return type_names[field.type]
        elif field.type in (
            descriptor_pb2.FieldDescriptorProto.TYPE_MESSAGE,
            descriptor_pb2.FieldDescriptorProto.TYPE_ENUM,
        ):
            # Remove leading dot if present
            type_name = field.type_name
            return type_name.removeprefix(".")
        else:
            return "unknown"

    def get_visibility(self, options) -> Visibility:
        # For now, return PUBLIC as default
        # In a real implementation, this would check custom options
        return Visibility.PUBLIC

    def get_validate_required(self, options) -> bool:
        """Extract validate_required value from field options.

        Checks if the field has the validate_required extension set.
        """
        if options.HasExtension(databricks_pb2.validate_required):
            return options.Extensions[databricks_pb2.validate_required]
        return False

    def extract_rpc_options(self, options) -> DatabricksRpcOptionsDescription | None:
        """Extract RPC options from method options.

        This requires databricks_pb2 to be available for full extraction.
        When running as a protoc plugin, the extension data is present in the
        file descriptors but needs the compiled extension definitions to parse.
        """
        # Check if the method has the Databricks RPC extension
        if options.HasExtension(databricks_pb2.rpc):
            # Extract the RPC options
            rpc_ext = options.Extensions[databricks_pb2.rpc]

            # Extract endpoint information
            path = None
            method = None
            since_major = None
            since_minor = None

            if rpc_ext.endpoints:
                # Use the first endpoint for now
                endpoint = rpc_ext.endpoints[0]
                path = endpoint.path if endpoint.HasField("path") else None
                method = endpoint.method if endpoint.HasField("method") else None

                if endpoint.HasField("since") and endpoint.since:
                    since_major = endpoint.since.major if endpoint.since.HasField("major") else None
                    since_minor = endpoint.since.minor if endpoint.since.HasField("minor") else None

            # Extract visibility
            visibility = "internal"  # default
            if rpc_ext.HasField("visibility"):
                visibility_enum = rpc_ext.visibility
                # Map the enum value to string
                if visibility_enum == databricks_pb2.PUBLIC:
                    visibility = "public"
                elif visibility_enum == databricks_pb2.INTERNAL:
                    visibility = "internal"
                elif visibility_enum == databricks_pb2.PUBLIC_UNDOCUMENTED:
                    visibility = "public_undocumented"

            # Extract error codes
            error_codes = list(rpc_ext.error_codes) if rpc_ext.error_codes else []

            # Extract RPC doc title
            rpc_doc_title = rpc_ext.rpc_doc_title if rpc_ext.HasField("rpc_doc_title") else ""

            return DatabricksRpcOptionsDescription(
                path=path,
                method=method,
                visibility=visibility,
                since_major=since_major,
                since_minor=since_minor,
                error_codes=error_codes,
                rpc_doc_title=rpc_doc_title,
            )

    def get_full_path_for_file(
        self, file: descriptor_pb2.FileDescriptorProto, name: str
    ) -> list[str]:
        path_parts = []
        if file.package:
            path_parts.extend(file.package.split("."))
        path_parts.append(name)
        return path_parts

    def get_full_path_for_nested(self, parent_path: list[str], name: str) -> list[str]:
        return parent_path + [name]

    def get_documentation(self, source_location: descriptor_pb2.SourceCodeInfo.Location) -> str:
        if source_location and source_location.leading_comments:
            return textwrap.dedent(source_location.leading_comments).strip()
        return ""

    def find_source_location(
        self, source_info: descriptor_pb2.SourceCodeInfo, path: list[int]
    ) -> descriptor_pb2.SourceCodeInfo.Location | None:
        for location in source_info.location:
            if list(location.path) == path:
                return location
        return None

    def process_field(
        self,
        field: descriptor_pb2.FieldDescriptorProto,
        parent_path: list[str],
        field_index: int,
        source_info: descriptor_pb2.SourceCodeInfo,
        message_path: list[int],
    ) -> ProtoMessageField:
        # Build source location path for this field
        field_path = message_path + [2, field_index]  # 2 = field in message
        location = self.find_source_location(source_info, field_path)

        field_type = self.get_field_type_name(field)

        default_value = None
        if field.HasField("default_value"):
            default_value = field.default_value

        return ProtoMessageField(
            description=self.get_documentation(location) if location else "",
            field_name=field.name,
            field_default=default_value,
            entity_type=str(field.type),
            field_type=field_type,
            full_path=self.get_full_path_for_nested(parent_path, field.name),
            visibility=self.get_visibility(field.options).value,
            since="",
            deprecated=field.options.deprecated if field.options.HasField("deprecated") else False,
            repeated=field.label == descriptor_pb2.FieldDescriptorProto.LABEL_REPEATED,
            validate_required=self.get_validate_required(field.options),
            oneof=[],
        )

    def process_enum_value(
        self,
        value: descriptor_pb2.EnumValueDescriptorProto,
        parent_path: list[str],
        value_index: int,
        source_info: descriptor_pb2.SourceCodeInfo,
        enum_path: list[int],
    ) -> ProtoEnumValue:
        # Build source location path for this enum value
        value_path = enum_path + [2, value_index]  # 2 = value in enum
        location = self.find_source_location(source_info, value_path)

        return ProtoEnumValue(
            value=value.name,
            full_path=self.get_full_path_for_nested(parent_path, value.name),
            visibility=self.get_visibility(value.options).value,
            description=self.get_documentation(location) if location else "",
        )

    def process_enum(
        self,
        enum: descriptor_pb2.EnumDescriptorProto,
        parent_path: list[str],
        enum_index: int,
        source_info: descriptor_pb2.SourceCodeInfo,
        parent_path_numbers: list[int],
        is_nested: bool = False,
    ) -> ProtoEnum:
        # Build source location path for this enum
        enum_path = parent_path_numbers + [4, enum_index] if is_nested else [5, enum_index]

        location = self.find_source_location(source_info, enum_path)

        full_path = self.get_full_path_for_nested(parent_path, enum.name)

        values = []
        for i, value in enumerate(enum.value):
            values.append(self.process_enum_value(value, full_path, i, source_info, enum_path))

        return ProtoEnum(
            name=enum.name,
            description=self.get_documentation(location) if location else "",
            full_path=full_path,
            values=values,
            visibility=self.get_visibility(enum.options).value,
        )

    def process_message(
        self,
        msg: descriptor_pb2.DescriptorProto,
        parent_path: list[str],
        msg_index: int,
        source_info: descriptor_pb2.SourceCodeInfo,
        parent_path_numbers: list[int] | None = None,
        is_nested: bool = False,
    ) -> ProtoMessage:
        # Build source location path for this message
        message_path = parent_path_numbers + [3, msg_index] if is_nested else [4, msg_index]

        location = self.find_source_location(source_info, message_path)

        full_path = self.get_full_path_for_nested(parent_path, msg.name)

        fields = []
        # Process regular fields
        for i, proto_field in enumerate(msg.field):
            if not proto_field.HasField("oneof_index"):  # Skip oneof fields for now
                fields.append(
                    self.process_field(proto_field, full_path, i, source_info, message_path)
                )

        # Process oneofs
        for oneof_index, oneof in enumerate(msg.oneof_decl):
            oneof_fields = []
            for i, proto_field in enumerate(msg.field):
                if proto_field.HasField("oneof_index") and proto_field.oneof_index == oneof_index:
                    oneof_fields.append(
                        self.process_field(proto_field, full_path, i, source_info, message_path)
                    )

            if oneof_fields:
                oneof_field = ProtoMessageField(
                    description="",
                    field_name=oneof.name,
                    field_default=None,
                    entity_type="oneof",
                    field_type="oneof",
                    full_path=self.get_full_path_for_nested(full_path, oneof.name),
                    visibility=self.get_visibility(oneof.options).value,
                    since="",
                    deprecated=False,
                    repeated=False,
                    validate_required=False,
                    oneof=oneof_fields,
                )
                fields.append(oneof_field)

        # Process nested enums
        enums = []
        for i, enum in enumerate(msg.enum_type):
            enums.append(
                self.process_enum(enum, full_path, i, source_info, message_path, is_nested=True)
            )

        # Process nested messages
        messages = []
        for i, nested in enumerate(msg.nested_type):
            messages.append(
                self.process_message(
                    nested, full_path, i, source_info, message_path, is_nested=True
                )
            )

        return ProtoMessage(
            name=msg.name,
            full_path=full_path,
            description=self.get_documentation(location) if location else "",
            visibility=self.get_visibility(msg.options).value,
            fields=fields,
            enums=enums,
            messages=messages,
        )

    def process_method(
        self,
        method: descriptor_pb2.MethodDescriptorProto,
        parent_path: list[str],
        method_index: int,
        source_info: descriptor_pb2.SourceCodeInfo,
        service_path: list[int],
    ) -> ProtoServiceMethod:
        # Build source location path for this method
        method_path = service_path + [2, method_index]  # 2 = method in service
        location = self.find_source_location(source_info, method_path)

        # Remove leading dots from type names
        input_type = method.input_type
        input_type = input_type.removeprefix(".")
        output_type = method.output_type
        output_type = output_type.removeprefix(".")

        input_path = input_type.split(".")
        output_path = output_type.split(".")

        # Extract RPC options from custom extensions
        rpc_options = self.extract_rpc_options(method.options)

        return ProtoServiceMethod(
            name=method.name,
            full_path=self.get_full_path_for_nested(parent_path, method.name),
            request_full_path=input_path,
            response_full_path=output_path,
            description=self.get_documentation(location) if location else "",
            rpc_options=rpc_options,
        )

    def process_service(
        self,
        service: descriptor_pb2.ServiceDescriptorProto,
        parent_path: list[str],
        service_index: int,
        source_info: descriptor_pb2.SourceCodeInfo,
    ) -> ProtoService:
        # Build source location path for this service
        service_path = [6, service_index]  # 6 = service at file level
        location = self.find_source_location(source_info, service_path)

        full_path = self.get_full_path_for_nested(parent_path, service.name)

        methods = []
        for i, method in enumerate(service.method):
            methods.append(self.process_method(method, full_path, i, source_info, service_path))

        return ProtoService(
            name=service.name,
            full_path=full_path,
            description=self.get_documentation(location) if location else "",
            visibility=self.get_visibility(service.options).value,
            methods=methods,
        )

    def process_file(
        self, file: descriptor_pb2.FileDescriptorProto, requested_vis: Visibility
    ) -> ProtoFile:
        elements = []

        # Base path from package
        base_path = file.package.split(".") if file.package else []

        # Get source code info for documentation
        source_info = file.source_code_info if file.HasField("source_code_info") else None
        if not source_info:
            source_info = descriptor_pb2.SourceCodeInfo()

        # Process top-level messages
        for i, msg in enumerate(file.message_type):
            elements.append(
                ProtoFileElement(message=self.process_message(msg, base_path, i, source_info))
            )

        # Process top-level enums
        for i, enum in enumerate(file.enum_type):
            elements.append(
                ProtoFileElement(
                    enum=self.process_enum(enum, base_path, i, source_info, [], is_nested=False)
                )
            )

        # Process services
        for i, service in enumerate(file.service):
            elements.append(
                ProtoFileElement(service=self.process_service(service, base_path, i, source_info))
            )

        return ProtoFile(
            filename=file.name, requested_visibility=requested_vis.value, content=elements
        )


class ProtocPlugin:
    """Protoc plugin implementation."""

    def __init__(self):
        self.generator = ProtobufDocGenerator()

    def process_request(
        self, request: plugin_pb2.CodeGeneratorRequest
    ) -> plugin_pb2.CodeGeneratorResponse:
        response = plugin_pb2.CodeGeneratorResponse()

        files = []

        # Process each file that was requested to be generated
        for file_name in request.file_to_generate:
            # Find the file descriptor
            file_descriptor = None
            for proto_file in request.proto_file:
                if proto_file.name == file_name:
                    file_descriptor = proto_file
                    break

            if file_descriptor:
                # Process the file
                proto_file = self.generator.process_file(file_descriptor, Visibility.PUBLIC)
                files.append(proto_file)

        # Generate documentation
        doc_content = ProtoAllContent(requested_visibility=Visibility.PUBLIC.value, files=files)

        # Generate doc_public.json
        doc_file = response.file.add()
        doc_file.name = "protos.json"
        doc_file.content = json.dumps(asdict(doc_content), indent=2)

        return response


def main():
    # Protoc plugin mode
    # Read CodeGeneratorRequest from stdin
    data = sys.stdin.buffer.read()
    request = plugin_pb2.CodeGeneratorRequest.FromString(data)

    # Process request
    plugin = ProtocPlugin()
    response = plugin.process_request(request)

    # Write response to stdout
    sys.stdout.buffer.write(response.SerializeToString())


if __name__ == "__main__":
    main()