# Copyright: 2021-2023 Paul Obermeier (obermeier@poSoft.de)
# Distributed under BSD license.

namespace eval pawt {
    namespace eval raw {
        variable ns [namespace current]

        namespace export ReadImageHeader
        namespace export ReadImageFile
        namespace export WriteImageFile

        #
        # Utility procedures.
        #

        proc _GetNativeByteOrder {} {
            if { $::tcl_platform(byteOrder) eq "littleEndian" } {
                return "Intel"
            } else {
                return "Motorola"
            }
        }

        proc _GetPixelSize { pixelType } {
            if { $pixelType eq "double" } {
                return 8
            } elseif { $pixelType eq "float" } {
                return 4
            } elseif { $pixelType eq "int" } {
                return 4
            } elseif { $pixelType eq "short" } {
                return 2
            } elseif { $pixelType eq "byte" } {
                return 1
            } else {
                error "Invalid PixelType value: $pixelType (must be byte, short, int, float or double)"
            }
        }

        proc _GetScanFormat { pixelType byteOrder } {
            if { $pixelType eq "double" } {
                if { $byteOrder eq "Intel" } {
                    set scanFmt "q"
                } else {
                    set scanFmt "Q"
                }
            } elseif { $pixelType eq "float" } {
                if { $byteOrder eq "Intel" } {
                    set scanFmt "r"
                } else {
                    set scanFmt "R"
                }
            } elseif { $pixelType eq "int" } {
                if { $byteOrder eq "Intel" } {
                    set scanFmt "iu"
                } else {
                    set scanFmt "Iu"
                }
            } elseif { $pixelType eq "short" } {
                if { $byteOrder eq "Intel" } {
                    set scanFmt "su"
                } else {
                    set scanFmt "Su"
                }
            } elseif { $pixelType eq "byte" } {
                set scanFmt "cu"
            } else {
                error "Invalid PixelType value: $pixelType (must be byte, short, int, float or double)"
            }
            return $scanFmt
        }

        proc _FormatHeaderLine { fmt val } {
            set HEADLEN 20
            set headerLine [format $fmt $val]
            while { [string length $headerLine] < [expr { $HEADLEN -1 }] } {
                append headerLine " "
            }
            append headerLine "\n"
            return $headerLine
        }

        #
        # Image header read/write procedures.
        #

        proc _GetImageHeader { fp fileName } {
            if { [gets $fp line] >= 0 } {
                scan $line "Magic=%s" magic
                if { $magic ne "RAW" } {
                    error "Invalid Magic value: $magic (must be RAW)"
                }
            } else {
                error "Error while trying to parse Magic keyword"
            }
            if { [gets $fp line] >= 0 } {
                scan $line "Width=%d" width
                if { $width <= 0 } {
                    error "Invalid Width value: $width (must be greater than zero)"
                }
            } else {
                error "Error while trying to parse Width keyword"
            }
            if { [gets $fp line] >= 0 } {
                scan $line "Height=%d" height
                if { $height <= 0 } {
                    error "Invalid Height value: $height (must be greater than zero)"
                }
            } else {
                error "Error while trying to parse Height keyword"
            }
            if { [gets $fp line] >= 0 } {
                scan $line "NumChan=%d" numChan
                if { $numChan <= 0 || $numChan > 4 } {
                    error "Invalid NumChan value: $numChan (must be in 1..4)"
                }
            } else {
                error "Error while trying to parse NumChan keyword"
            }
            if { [gets $fp line] >= 0 } {
                scan $line "ByteOrder=%s" byteOrder
                if { $byteOrder ne "Intel" && $byteOrder ne "Motorola" } {
                    error "Invalid ByteOrder value: $byteOrder (must be Intel or Motorola)"
                }
            } else {
                error "Error while trying to parse ByteOrder keyword"
            }
            if { [gets $fp line] >= 0 } {
                scan $line "ScanOrder=%s" scanOrder
                if { $scanOrder ne "TopDown" && $scanOrder ne "BottomUp" } {
                    error "Invalid ScanOrder value: $scanOrder (must be TopDown or BottomUp)"
                }
            } else {
                error "Error while trying to parse ScanOrder keyword"
            }
            if { [gets $fp line] >= 0 } {
                scan $line "PixelType=%s" pixelType
                if { $pixelType ne "byte" && $pixelType ne "short" && $pixelType ne "int" && \
                     $pixelType ne "float" && $pixelType ne "double" } {
                    error "Invalid PixelType value: $pixelType (must be byte, short, int, float or double)"
                }
            } else {
                error "Error while trying to parse PixelType keyword"
            }
            set headerDict [dict create]

            # Store the values contained in the header.
            dict set headerDict Magic     $magic
            dict set headerDict Width     $width
            dict set headerDict Height    $height
            dict set headerDict NumChan   $numChan
            dict set headerDict ByteOrder $byteOrder
            dict set headerDict ScanOrder $scanOrder
            dict set headerDict PixelType $pixelType

            # Calculate derived information from the header values.
            set pixelSize [_GetPixelSize $pixelType]
            set scanFmt   [_GetScanFormat $pixelType $byteOrder]
            dict set headerDict PixelSize  $pixelSize
            dict set headerDict ScanFormat $scanFmt
            dict set headerDict NumPixel   [expr { $width * $height }]
            dict set headerDict NumByte    [expr { $width * $height * $numChan * $pixelSize }]
            dict set headerDict FileName   $fileName

            return $headerDict
        }

        proc _CreateHeaderDictFromOptions { fileName args } {
            set headerDict [dict create]
            set useHeader  true
            set width      128
            set height     128
            set numChan    1
            set scanOrder  "TopDown"
            set byteOrder  "Intel"
            set pixelType  "byte"
            foreach { key val } $args {
                switch -exact -- $key {
                    "-useheader" {
                        set useHeader $val
                    }
                    "-width" {
                        set width $val
                    }
                    "-height" {
                        set height $val
                    }
                    "-nchan" -
                    "-numchan" {
                        set numChan $val
                    }
                    "-scanorder" {
                        set scanOrder $val
                    }
                    "-byteorder" {
                        set byteOrder $val
                    }
                    "-pixeltype" {
                        set pixelType $val
                    }
                }
            }

            if { ! $useHeader } {
                # Store the values contained in the header.
                dict set headerDict Magic     "RAW"
                dict set headerDict Width     $width
                dict set headerDict Height    $height
                dict set headerDict NumChan   $numChan
                dict set headerDict ByteOrder $byteOrder
                dict set headerDict ScanOrder $scanOrder
                dict set headerDict PixelType $pixelType

                # Calculate derived information from the header values.
                set pixelSize [_GetPixelSize $pixelType]
                set scanFmt   [_GetScanFormat $pixelType $byteOrder]
                dict set headerDict PixelSize  $pixelSize
                dict set headerDict ScanFormat $scanFmt
                dict set headerDict NumPixel   [expr { $width * $height }]
                dict set headerDict NumByte    [expr { $width * $height * $numChan * $pixelSize }]
                dict set headerDict FileName   $fileName
            }
            return $headerDict
        }

        proc ReadImageHeader { rawImgFile args } {
            # Read the header of a RAW image file.
            #
            # rawImgFile - File name of the RAW image.
            # args       - Optional format specific arguments. Currently none.
            #
            # See [::pawt::ReadImageHeader] for the default entries of a
            # header dictionary.
            #
            # Returns the header information as a dictionary.
            # If the file does not exist or the header contains invalid information,
            # an error is thrown.
            #
            # See also: ReadImageFile WriteImageFile

            set retVal [catch {open $rawImgFile "r"} fp]
            if { $retVal != 0 } {
                error "Cannot open file $rawImgFile"
            }
            fconfigure $fp -translation binary
            set headerDict [_GetImageHeader $fp [file tail $rawImgFile]]
            close $fp
            return $headerDict
        }

        proc _PutImageHeader { fp headerDict scanOrder } {
            set headerStr ""
            append headerStr [_FormatHeaderLine "Magic=%s"     "RAW"]
            append headerStr [_FormatHeaderLine "Width=%d"     [dict get $headerDict Width]]
            append headerStr [_FormatHeaderLine "Height=%d"    [dict get $headerDict Height]]
            append headerStr [_FormatHeaderLine "NumChan=%d"   [dict get $headerDict NumChan]]
            append headerStr [_FormatHeaderLine "ByteOrder=%s" [dict get $headerDict ByteOrder]]
            if { $scanOrder eq "" } {
                append headerStr [_FormatHeaderLine "ScanOrder=%s" [dict get $headerDict ScanOrder]]
            } else {
                append headerStr [_FormatHeaderLine "ScanOrder=%s" $scanOrder]
            }
            append headerStr [_FormatHeaderLine "PixelType=%s" [dict get $headerDict PixelType]]
            puts -nonewline $fp $headerStr
        }

        #
        # Image data read/write procedures.
        #

        proc _GetImageData { fp headerDict dataDict skipBytes } {
            upvar 1 $dataDict myDict

            if { $skipBytes > 0 } {
                read $fp $skipBytes
            }
            set numBytes [dict get $headerDict NumByte]
            dict set myDict Data [read $fp $numBytes]
        }

        proc _PutImageData { fp imgDict scanOrder } {
            upvar 1 $imgDict myDict

            set headerDict [dict get $myDict Header]
            if { $scanOrder eq "" || [dict get $headerDict ScanOrder] eq $scanOrder } {
                puts -nonewline $fp [dict get $myDict Data]
            } else {
                set numChan   [dict get $headerDict NumChan]
                set pixelSize [dict get $headerDict PixelSize]
                set width     [dict get $headerDict Width]
                set height    [dict get $headerDict Height]

                set numBytesPerRow [expr { $width * $pixelSize * $numChan }]
                set pos [expr { $width * ($height - 1) * $pixelSize * $numChan } ]
                for { set row 0 } { $row < $height } { incr row } {
                    set rowVal [string range [dict get $myDict Data] $pos [expr { $pos + $numBytesPerRow -1 }]]
                    puts -nonewline $fp $rowVal
                    incr pos [expr { -$numBytesPerRow }]
                }
            }
        }

        proc ReadImageFile { rawImgFile args } {
            # Read a RAW image file into a dictionary.
            #
            # rawImgFile - File name of the RAW image.
            # args       - Optional format specific arguments described below.
            #
            # -useheader <bool>   - Use RAW image header.
            # -skipbytes <int>    - Skip number of bytes before reading image data.
            #                       Only valid, if `-useheader` is false.
            # -width <int>        - Width of the image in pixel.
            # -height <int>       - Height of the image in pixel.
            # -numchan <int>      - Number of channels contained in the image.
            # -nchan <int>        - Number of channels contained in the image.
            # -byteorder <string> - `Motorola` for big-endian, `Intel` for little-endian architecture.
            # -scanorder <string> - `TopDown` or `BottomUp`.
            # -pixeltype <string> - `byte`   for 1-byte unsigned integers, 
            #                       `short`  for 2-byte unsigned integers,
            #                       `int`    for 4-byte unsigned integers,
            #                       `float`  for 4-byte single precision floating point values.
            #                       `double` for 8-byte double precision floating point values.
            #
            # The above specified format options can be used to interpret any binary data as an image.
            # Set `-useheader`to false and specify the other parameters accordingly.
            #
            # Get the image data as a dictionary containing 2 keys: `Header` and `Data`.
            # * See [::pawt::ReadImageHeader] for the entries of the `Header` key.
            # * See [::pawt::ReadImageFile] for the entries of the `Data` key.
            #
            # Returns the image information as a dictionary.
            # If the file does not exist or no supported image format can be detected,
            # an error is thrown.
            #
            # See also: ReadImageHeader WriteImageFile

            set retVal [catch {open $rawImgFile "r"} fp]
            if { $retVal != 0 } {
                error "Cannot open file $rawImgFile"
            }
            fconfigure $fp -translation binary

            set imgDict [dict create]
            if { [llength $args] == 0 } {
                set headerDict [_GetImageHeader $fp [file tail $rawImgFile]]
            } else {
                set headerDict [_CreateHeaderDictFromOptions $rawImgFile {*}$args]
                if { [dict size $headerDict] == 0 } {
                    # Option -useheader in args was set to true, so read the header.
                    set headerDict [_GetImageHeader $fp [file tail $rawImgFile]]
                }
            }
            set useHeader true
            set skipBytes 0

            foreach { key val } $args {
                switch -exact -- $key {
                    "-useheader" {
                        set useHeader $val
                    }
                    "-skipbytes" {
                        set skipBytes [expr int($val)]
                    }
                }
            }
            if { $skipBytes < 0 || $useHeader } {
                set skipBytes 0
            }
            _GetImageData $fp $headerDict imgDict $skipBytes
            dict append imgDict Header $headerDict

            close $fp
            return $imgDict
        }

        proc WriteImageFile { imgDict rawImgFile args } {
            # Write the values of an image dict into a RAW image file.
            #
            # imgDict    - Image dictionary.
            # rawImgFile - File name of the image.
            # args       - Optional format specific arguments described below.
            #
            # -useheader <bool>   - Write RAW image header.
            #                       If option is not specified, a RAW image header is written.
            # -scanorder <string> - Store image data in row order `TopDown` or `BottomUp`.
            #                       If option is not specified, the `ScanOrder` value of 
            #                       the dictionary is used.
            #
            # See [::pawt::WriteImageFile] for general informations regarding writing images.
            #
            # Returns no value. If invalid options are specified or the file could not
            # be written, an error is thrown.
            #
            # See also: ReadImageFile ReadImageHeader

            set opts [dict create \
                -useheader true \
                -scanorder "" \
            ]

            foreach { key value } $args {
                if { [dict exists $opts $key] } {
                    if { $value eq "" } {
                        error "WriteImageFile: No value specified for key \"$key\"."
                    }
                    dict set opts $key $value
                } else {
                    error "WriteImageFile: Unknown option \"$key\" specified."
                }
            }

            set retVal [catch {open $rawImgFile "w"} fp]
            if { $retVal != 0 } {
                error "Cannot open output file $rawImgFile"
            }
            fconfigure $fp -translation binary

            if { [dict get $opts "-useheader"] } {
                _PutImageHeader $fp [dict get $imgDict Header] [dict get $opts "-scanorder"]
            }
            _PutImageData $fp imgDict [dict get $opts "-scanorder"]
            close $fp
        }
    }
}
