"""Dissector generation for binary protocols."""

from __future__ import annotations

from pathlib import Path
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from oscura.analyzers.binary.core.results import BinaryAnalysisResult, Field


class DissectorGenerator:
    """Generate protocol dissectors from binary analysis results.

    Supports Wireshark Lua and Kaitai Struct formats.

    Example:
        >>> generator = DissectorGenerator()
        >>> generator.generate_wireshark_lua(results, "protocol.lua")
        >>> generator.generate_kaitai_struct(results, "protocol.ksy")
    """

    def generate_wireshark_lua(self, result: BinaryAnalysisResult, output_path: str | Path) -> None:
        """Generate Wireshark Lua dissector.

        Args:
            result: Binary analysis result.
            output_path: Path to save Lua dissector.

        Example:
            >>> generator = DissectorGenerator()
            >>> generator.generate_wireshark_lua(results, "custom.lua")
        """
        if not result.structure or not result.structure.has_messages:
            # Create minimal dissector
            lua_code = """-- No structure detected
-- This is a placeholder dissector

local proto = Proto("custom", "Custom Protocol")

function proto.dissector(buffer, pinfo, tree)
    pinfo.cols.protocol = "CUSTOM"
    local subtree = tree:add(proto, buffer(), "Custom Protocol (No structure detected)")
end

DissectorTable.get("wtap_encap"):add(wtap.USER0, proto)
"""
            Path(output_path).write_text(lua_code)
            return

        # Generate full dissector
        fields = result.structure.fields
        message_length = result.structure.message_length

        lua_code = f"""-- Auto-generated Wireshark dissector
-- Generated by Oscura Binary Analysis
-- Message length: {message_length} bytes

local proto = Proto("custom", "Custom Binary Protocol")

-- Field definitions
"""

        # Add field definitions
        for field in fields:
            field_name = field.name.replace(" ", "_").lower()
            lua_code += f'proto.fields.{field_name} = ProtoField.bytes("custom.{field_name}", "{field.name}", base.HEX)\n'

        lua_code += """
-- Dissector function
function proto.dissector(buffer, pinfo, tree)
    local length = buffer:len()
    if length == 0 then return end

    pinfo.cols.protocol = proto.name

    local subtree = tree:add(proto, buffer(), "Custom Binary Protocol")

"""

        # Add field parsing
        for field in fields:
            field_name = field.name.replace(" ", "_").lower()
            lua_code += f"    -- {field.name} ({field.field_type.value})\n"
            lua_code += f"    if length >= {field.offset + field.length} then\n"
            lua_code += f"        subtree:add(proto.fields.{field_name}, buffer({field.offset}, {field.length}))\n"
            lua_code += "    end\n\n"

        lua_code += """end

-- Register dissector
DissectorTable.get("wtap_encap"):add(wtap.USER0, proto)

-- Also register for UDP port (example)
-- local udp_table = DissectorTable.get("udp.port")
-- udp_table:add(12345, proto)
"""

        Path(output_path).write_text(lua_code)

    def generate_kaitai_struct(self, result: BinaryAnalysisResult, output_path: str | Path) -> None:
        """Generate Kaitai Struct YAML definition.

        Args:
            result: Binary analysis result.
            output_path: Path to save Kaitai Struct definition.

        Example:
            >>> generator = DissectorGenerator()
            >>> generator.generate_kaitai_struct(results, "protocol.ksy")
        """
        if not result.structure or not result.structure.has_messages:
            # Create minimal definition
            ksy_content = """meta:
  id: custom_protocol
  title: Custom Binary Protocol
  endian: le

doc: |
  No structure detected. This is a placeholder definition.

seq: []
"""
            Path(output_path).write_text(ksy_content)
            return

        # Generate full definition
        fields = result.structure.fields

        ksy_content = f"""meta:
  id: custom_protocol
  title: Custom Binary Protocol
  endian: le

doc: |
  Auto-generated Kaitai Struct definition from Oscura Binary Analysis.
  Message length: {result.structure.message_length} bytes
  Messages detected: {result.structure.message_count}

seq:
"""

        # Add field definitions
        for field in fields:
            field_name = field.name.replace(" ", "_").lower()
            ksy_type = self._map_field_to_kaitai_type(field)

            ksy_content += f"  - id: {field_name}\n"
            ksy_content += f"    type: {ksy_type}\n"
            ksy_content += f"    doc: '{field.field_type.value} field at offset {field.offset}'\n"

        Path(output_path).write_text(ksy_content)

    def _map_field_to_kaitai_type(self, field: Field) -> str:
        """Map field to Kaitai Struct type.

        Args:
            field: Field to map.

        Returns:
            Kaitai type string.
        """
        # Map by field length
        if field.length == 1:
            return "u1"
        elif field.length == 2:
            return "u2"
        elif field.length == 4:
            return "u4"
        elif field.length == 8:
            return "u8"
        else:
            # Variable or custom length
            return f"bytes({field.length})" if field.length < 1024 else "bytes_eos"
