use crate::google::protobuf::descriptor::{DescriptorProto, EnumDescriptorProto, FieldDescriptorProto, MessageOptions, OneofDescriptorProto};
use crate::google::protobuf::descriptor::FileDescriptorSet;
use crate::generator::types::map_type_to_rust_builder;

use roto_runtime::ProtoAccessor;
use crate::generator::utils::{to_pascal_case, to_snake_case};
use crate::generator::types::map_type_to_rust_accessor;

pub fn write_enum(enum_proto: &EnumDescriptorProto, output: &mut String) {

    let enum_name = to_pascal_case(enum_proto.name().unwrap());

    output.push_str(&format!(
        "#[derive(Debug, Clone, Copy, PartialEq, Eq)]\n#[repr(i32)]\npub enum {} {{\n",
        enum_name
    ));

    let mut values = enum_proto.value();
    let mut zero_variant_name = None;
    while let Some(val_res) = values.next() {
        let (val_data, _) = val_res.expect("Failed to iterate enum");
        let accessor =
            ProtoAccessor::new(val_data).expect("Failed to parse EnumValueDescriptorProto");
        let (name_bytes, _) = accessor.get_value(1).expect("Enum value name missing");
        let name = std::str::from_utf8(name_bytes).expect("Enum value name invalid utf8");
        let (num_bytes, _) = accessor.get_value(2).expect("Enum value number missing");
        let (num, _) =
            roto_runtime::read_varint(num_bytes).expect("Enum value number invalid varint");

        let pascal_name = to_pascal_case(name);
        if num == 0 {
            zero_variant_name = Some(pascal_name.clone());
        }
        output.push_str(&format!("    {} = {},\n", pascal_name, num));
    }

    if zero_variant_name.is_none() {
        output.push_str("    Unknown = 0,\n");
        zero_variant_name = Some("Unknown".to_string());
    }

    output.push_str("}\n\n");

    output.push_str(&format!(
        "impl {} {{\n    pub fn from_i32(value: i32) -> Self {{\n        match value {{\n",
        enum_name
    ));

    let mut values = enum_proto.value();
    while let Some(val_res) = values.next() {
        let (val_data, _) = val_res.expect("Failed to read enum value");
        let accessor =
            ProtoAccessor::new(val_data).expect("Failed to parse EnumValueDescriptorProto");
        let (name_bytes, _) = accessor.get_value(1).expect("Enum value name missing");
        let name = std::str::from_utf8(name_bytes).expect("Enum value name invalid utf8");
        let (num_bytes, _) = accessor.get_value(2).expect("Enum value number missing");
        let (num, _) =
            roto_runtime::read_varint(num_bytes).expect("Enum value number invalid varint");

        output.push_str(&format!(
            "            {} => {}::{},\n",
            num,
            enum_name,
            to_pascal_case(name)
        ));
    }

    output.push_str(&format!(
        "            _ => {}::{},\n",
        enum_name,
        zero_variant_name.as_ref().unwrap()
    ));
    output.push_str("        }\n    }\n}\n\n");
}
pub fn write_message(msg_proto: &DescriptorProto, output: &mut String) {
    let msg_name = to_pascal_case(msg_proto.name().unwrap());
    let mod_name = to_snake_case(msg_proto.name().unwrap());

    let mut fields_info = Vec::new();
    for field_res in msg_proto.field() {
        let (field_data, _) = field_res.expect("Failed to iterate field");
        let field_proto =
            FieldDescriptorProto::new(field_data).expect("Failed to parse FieldDescriptorProto");
        let field_name = field_proto.name().unwrap();

        let tag = field_proto.number().unwrap();
        let f_type = field_proto.rype().unwrap() as i32;
        let f_label = field_proto.label().unwrap() as i32;
        let oneof_index = field_proto.oneof_index().ok();
        let is_map = field_proto
            .options()
            .map(|opt| {
                MessageOptions::new(opt)
                    .unwrap()
                    .map_entry()
                    .unwrap_or(false)
            })
            .unwrap_or(false);

        fields_info.push((
            field_name.to_string(),
            tag,
            f_type,
            f_label,
            oneof_index,
            is_map,
        ));
    }

    let mut oneofs = Vec::new();
    for o_res in msg_proto.oneof_decl() {
        let (o, _) = o_res.expect("Failed to iterate oneof");
        oneofs.push(o);
    }

    output.push_str(&format!("pub struct {}<'a> {{\n", msg_name));
    output.push_str("    accessor: roto_runtime::ProtoAccessor<'a>,\n");

    for (field_name, _tag, _f_type, f_label, _oneof_index, _is_map) in &fields_info {
        if *f_label == 3 {
            output.push_str(&format!("    {}_start: Option<usize>,\n", field_name));
            output.push_str(&format!("    {}_end: Option<usize>,\n", field_name));
        } else {
            output.push_str(&format!("    {}_offset: Option<usize>,\n", field_name));
        }
    }
    output.push_str("}\n\n");

    output.push_str(&format!("impl<'a> {}<'a> {{\n", msg_name));
    output.push_str("    pub fn new(data: &'a [u8]) -> roto_runtime::Result<Self> {\n");
    output.push_str("        let accessor = roto_runtime::ProtoAccessor::new(data)?;\n");
    for (name, _, _, label, _oneof_index, _is_map) in &fields_info {
        if *label == 3 {
            output.push_str(&format!("        let mut {}_start = None;\n", name));
            output.push_str(&format!("        let mut {}_end = None;\n", name));
        } else {
            output.push_str(&format!("        let mut {}_offset = None;\n", name));
        }
    }

    output.push_str("        for item in accessor.fields() {\n");
    output.push_str("            let (offset, tag, _) = item?;\n");
    for (name, tag, _, label, _oneof_index, _is_map) in &fields_info {
        if *label == 3 {
            output.push_str(&format!("            if tag.field_number == {} {{\n", tag));
            output.push_str(&format!(
                "                if {}_start.is_none() {{ {}_start = Some(offset); }}\n",
                name, name
            ));
            output.push_str(&format!("                {}_end = Some(offset);\n", name));
            output.push_str("            }\n");
        } else {
            output.push_str(&format!(
                "            if tag.field_number == {} {{ {}_offset = Some(offset); }}\n",
                tag, name
            ));
        }
    }
    output.push_str("        }\n\n");

    output.push_str("        Ok(Self {\n");
    output.push_str("            accessor,\n");
    for (name, _, _, label, _oneof_index, _is_map) in &fields_info {
        if *label == 3 {
            output.push_str(&format!("{}_start, {}_end,\n", name, name));
        } else {
            output.push_str(&format!("{}_offset,\n", name));
        }
    }
    output.push_str("        })\n    }\n\n");

    for (field_name, tag, f_type, f_label, _oneof_index, is_map) in &fields_info {
        let (rust_type, logic, default_val) = map_type_to_rust_accessor(*f_type, *f_label, *is_map);
        let safe_name = if field_name == "type" {
            format!("r#{}", field_name)
        } else {
            field_name.clone()
        };

        if *f_label == 3 {
            output.push_str(&format!(
                "    pub fn {}(&self) -> {} {{\n",
                safe_name, rust_type
            ));
            output.push_str(&format!(
                "        match (self.{}_start, self.{}_end) {{\n",
                field_name, field_name
            ));
            if *is_map {
                output.push_str(&format!("            (Some(start), Some(end)) => roto_runtime::MapFieldIterator::new(self.accessor.iter_repeated_range({}, start, end)),\n", tag));
                output.push_str(&format!(
                    "            _ => roto_runtime::MapFieldIterator::new(self.accessor.iter_repeated({})),\n",
                    tag
                ));
            } else {
                output.push_str(&format!("            (Some(start), Some(end)) => self.accessor.iter_repeated_range({}, start, end),\n", tag));
                output.push_str(&format!(
                    "            _ => self.accessor.iter_repeated({}),\n",
                    tag
                ));
            }
            output.push_str("        }\n    }\n\n");
        } else {
            output.push_str(&format!(
                "    pub fn {}(&self) -> roto_runtime::Result<{}> {{\n",
                safe_name, rust_type
            ));
            output.push_str(&format!(
                "        let offset = self.{}_offset.ok_or(roto_runtime::RotoError::FieldNotFound)?;\n",
                field_name
            ));
            output.push_str("        let (bytes, _) = self.accessor.get_value_at(offset)?;\n");
            output.push_str(&format!("        {}\n", logic));
            output.push_str("    }\n\n");

            output.push_str(&format!(
                "    pub fn {}_or_default(&self) -> roto_runtime::Result<{}> {{\n",
                safe_name, rust_type
            ));
            output.push_str(&format!(
                "        self.{}().or(Ok({}))\n",
                safe_name, default_val
            ));
            output.push_str("    }\n\n");

            output.push_str(&format!(
                "    pub fn has_{}(&self) -> bool {{ self.{}_offset.is_some() }}\n\n",
                field_name, field_name
            ));
        }
    }

    for (oneof_index, oneof_proto) in oneofs.iter().enumerate() {
        let oneof_desc =
            OneofDescriptorProto::new(*oneof_proto).expect("Failed to parse OneofDescriptorProto");
        let oneof_name = oneof_desc.name().unwrap();
        let pascal_oneof_name = to_pascal_case(oneof_name);
        let snake_oneof_name = to_snake_case(oneof_name);

        let return_type = format!("{}::{}<'a>", mod_name, pascal_oneof_name);
        let signature = format!(
            "    pub fn which_{}(&self) -> roto_runtime::Result<Option<{}> > {{\n",
            snake_oneof_name, return_type
        );
        output.push_str(&signature);
        for (field_name, _tag, _f_type, _f_label, f_oneof_index, _is_map) in &fields_info {
            if *f_oneof_index == Some(oneof_index as i32) {
                let safe_field_name = if field_name == "type" {
                    format!("r#{}", field_name)
                } else {
                    field_name.clone()
                };
                output.push_str(&format!(
                    "        if self.{}_offset.is_some() {{\n",
                    field_name
                ));
                output.push_str(&format!(
                    "            return Ok(Some({}::{}::{} (self.{}()?)));\n",
                    mod_name, pascal_oneof_name, safe_field_name, safe_field_name
                ));
                output.push_str("        }\n");
            }
        }
        output.push_str("        Ok(None)\n    }\n\n");
    }

    // raw_fields() convenience on the message struct (before closing the impl)
    output.push_str("    pub fn raw_fields(&self) -> roto_runtime::RawFieldIterator<'a> {\n");
    output.push_str("        self.accessor.raw_fields()\n");
    output.push_str("    }\n\n");
    output.push_str("}\n\n");

    // Collect builder field info so we can use it multiple times below.
    // Tuple: (field_name, safe_name, tag, rust_type, write_method)
    let mut builder_fields: Vec<(String, String, u32, String, String)> = Vec::new();
    for field_res in msg_proto.field() {
        let (field_data, _) = field_res.expect("Failed to iterate field");
        let field_proto =
            FieldDescriptorProto::new(field_data).expect("Failed to parse FieldDescriptorProto");
        let field_name = field_proto.name().unwrap().to_string();
        let safe_name = if field_name == "type" {
            format!("r#{}", field_name)
        } else {
            field_name.clone()
        };
        let tag = field_proto.number().unwrap();
        let f_type = field_proto.rype().unwrap() as i32;
        let (rust_type, method) = map_type_to_rust_builder(f_type);
        builder_fields.push((field_name, safe_name, tag as u32, rust_type, method));
    }

    // Builder struct — one `_written: bool` flag per field
    output.push_str(&format!("pub struct {}Builder<'b> {{\n", msg_name));
    output.push_str("    builder: roto_runtime::ProtoBuilder<'b>,\n");
    for (field_name, _, _, _, _) in &builder_fields {
        output.push_str(&format!("    {}_written: bool,\n", field_name));
    }
    output.push_str(&format!("}}\n\nimpl<'b> {}Builder<'b> {{\n", msg_name));

    // Constructor — initialise every flag to false
    output.push_str(&format!(
        "    pub fn builder(buf: &mut [u8]) -> {}Builder<'_> {{\n        {}Builder {{\n",
        msg_name, msg_name
    ));
    output.push_str("            builder: roto_runtime::ProtoBuilder::new(buf),\n");
    for (field_name, _, _, _, _) in &builder_fields {
        output.push_str(&format!("            {}_written: false,\n", field_name));
    }
    output.push_str("        }\n    }\n\n");

    // Per-field setters — mark field as written
    for (field_name, safe_name, tag, rust_type, method) in &builder_fields {
        output.push_str(&format!(
            "    pub fn {}(mut self, value: {}) -> roto_runtime::Result<Self> {{\n        self.builder.{}({}, value)?;\n        self.{}_written = true;\n        Ok(self)\n    }}\n\n",
            safe_name, rust_type, method, tag, field_name
        ));
    }

    // with() — copies unseen fields from an existing message
    output.push_str(&format!(
        "    pub fn with(mut self, msg: &{}<'_>) -> roto_runtime::Result<Self> {{\n",
        msg_name
    ));
    output.push_str("        for item in msg.accessor.raw_fields() {\n");
    output.push_str("            let (field_number, raw_bytes) = item?;\n");
    output.push_str("            let is_written = match field_number {\n");
    for (field_name, _, tag, _, _) in &builder_fields {
        output.push_str(&format!(
            "                {} => self.{}_written,\n",
            tag, field_name
        ));
    }
    output.push_str("                _ => false,\n");
    output.push_str("            };\n");
    output.push_str("            if !is_written {\n");
    output.push_str("                self.builder.write_raw(raw_bytes)?;\n");
    output.push_str("            }\n");
    output.push_str("        }\n");
    output.push_str("        Ok(self)\n");
    output.push_str("    }\n\n");

    output.push_str(&format!("    pub fn finish(self) -> roto_runtime::Result<&'b mut [u8]> {{\n        self.builder.finish()\n    }}\n}}\n\n"));

    output.push_str(&format!(
        pub struct Owned{} {{\n",
        msg_name
    ));
    output.push_str("    pub data: bytes::Bytes,\n");
    output.push_str("}\n\n");

    output.push_str(&format!(
        impl roto_runtime::RotoOwned for Owned{} {{\n",
        msg_name
    ));
    output.push_str(&format!("    type Reader<'a> = {}<'a>;\n", msg_name));
    output.push_str(&format!("    fn reader(&self) -> {}<'_> {{\n", msg_name));
    output.push_str(&format!(
        "        {}::new(&self.data).expect(\"failed to create reader\")\n",
        msg_name
    ));
    output.push_str("    }\n");
    output.push_str("}\n\n");

    output.push_str(&format!(
        impl roto_runtime::RotoMessage for Owned{} {{\n",
        msg_name
    ));
    output.push_str("    fn decode(buf: bytes::Bytes) -> roto_runtime::Result<Self> {\n");
    output.push_str(&format!("        Ok(Owned{} {{ data: buf }})\n", msg_name));
    output.push_str("    }\n\n");
    output.push_str("    fn bytes(&self) -> bytes::Bytes {\n");
    output.push_str("        self.data.clone()\n");
    output.push_str("    }\n");
    output.push_str("}\n\n");

    let mut nested_enums = Vec::new();
    for e_res in msg_proto.enum_type() {
        if let Ok((e, _)) = e_res {
            nested_enums.push(e);
        }
    }
    let mut nested_msgs = Vec::new();
    for m_res in msg_proto.nested_type() {
        if let Ok((m, _)) = m_res {
            nested_msgs.push(m);
        }
    }

    if !nested_enums.is_empty() || !nested_msgs.is_empty() || !oneofs.is_empty() {
        let mod_name = to_snake_case(msg_proto.name().unwrap());
        output.push_str(&format!("pub mod {} {{\n", mod_name));
        for e_data in &nested_enums {
            write_enum(
                &EnumDescriptorProto::new(e_data)
                    .expect("Failed to parse nested EnumDescriptorProto"),
                output,
            );
        }
        for m_data in &nested_msgs {
            write_message(
                &DescriptorProto::new(m_data).expect("Failed to parse nested DescriptorProto"),
                output,
            );
        }

        for (oneof_index, oneof_proto) in oneofs.iter().enumerate() {
            let oneof_desc = OneofDescriptorProto::new(*oneof_proto)
                .expect("Failed to parse OneofDescriptorProto");
            let oneof_name = oneof_desc.name().unwrap();
            let pascal_oneof_name = to_pascal_case(oneof_name);
            output.push_str(&format!("pub enum {}<'a> {{\n", pascal_oneof_name));
            for (field_name, _tag, f_type, f_label, f_oneof_index, _is_map) in &fields_info {
                if *f_oneof_index == Some(oneof_index as i32) {
                    let (rust_type, _, _) = map_type_to_rust_accessor(*f_type, *f_label, *_is_map);
                    let safe_field_name = if field_name == "type" {
                        format!("r#{}", field_name)
                    } else {
                        field_name.clone()
                    };
                    output.push_str(&format!("    {}({}),\n", safe_field_name, rust_type));
                }
            }
            output.push_str("}\n\n");
        }
    }
    if !nested_enums.is_empty() || !nested_msgs.is_empty() || !oneofs.is_empty() {
        output.push_str("}\n\n");
    }
}
pub fn generate_protobuf_code(
    set: &FileDescriptorSet,
    files_to_generate: Option<&[String]>,
    generate_mod_files: bool,
) -> Vec<(String, String)> {
    generate_files_common(
        set,
        files_to_generate,
        generate_mod_files,
        DATA_IMPORTS,
        |file_proto, output| {
            // Enums
            for enum_res in file_proto.enum_type() {
                let (enum_data, _) = enum_res.expect("Failed to iterate enum");
                write_enum(
                    &EnumDescriptorProto::new(enum_data)
                        .expect("Failed to parse EnumDescriptorProto"),
                    output,
                );
            }

            // Messages
            for msg_res in file_proto.message_type() {
                let (msg_data, _) = msg_res.expect("Failed to iterate message");
                write_message(
                    &DescriptorProto::new(msg_data).expect("Failed to parse DescriptorProto"),
                    output,
                );
            }
        },
    )
}






