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

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

        namespace export ReadImageFile
        namespace export ReadImageHeader
        namespace export WriteImageFile

        #
        # Utility procedures.
        #

        proc _GetTypeList {} {
            return [list "P2" "P3" "P5" "P6"]
        }

        proc _GetAsciiList {} {
            return [list true true false false]
        }

        proc _GetChanList {} {
            return [list 1 3 1 3]
        }

        proc _IsAscii { ppmType } {
            set ind [lsearch -exact [_GetTypeList] $ppmType]
            return [lindex [_GetAsciiList] $ind]
        }

        proc _GetNumChans { ppmType } {
            set ind [lsearch -exact [_GetTypeList] $ppmType]
            return [lindex [_GetChanList] $ind]
        }

        proc _GetPpmType { numChans isAscii } {
            set index 0
            foreach chans [_GetChanList] ascii [_GetAsciiList] {
                if { $chans == $numChans && $ascii == $isAscii } {
                    return [lindex [_GetTypeList] $index]
                }
                incr index
            }
            return "InvalidPpmType"
        }

        proc _GetPixelTypeName { maxColors } {
            if { $maxColors < 256 } {
                return "byte"
            } else {
                return "short"
            }
        }

        proc _GetMaxColors { pixelType } {
            if { $pixelType eq "short" } {
                return 65535
            } elseif { $pixelType eq "byte" } {
                return 255
            } else {
                error "Invalid PixelType value: $pixelType (must be byte or short)"
            }
        }

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

        proc _GetScanFormat { pixelType byteOrder } {
            if { $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 or short)"
            }
            return $scanFmt
        }

        #
        # Image header read/write procedures.
        #

        proc _GetImageHeader { fp fileName } {
            # Read the header lines (ASCII): Type, width and height, maximum color value.
            while { ! [info exists maxColor] } {
                if { [gets $fp line] < 0 } {
                    close $fp
                    error "Error while parsing PPM header"
                }
                set value [string trim $line]
                if { [string index $value 0] eq "#" } {
                    continue
                }
                if { ! [info exists type] } {
                    set type $value
                    if { [lsearch -exact [_GetTypeList] $type] < 0 } {
                        close $fp
                        error "Unsupported PPM type $type. Must be \"P2\", \"P3\", \"P5\", \"P6\"."
                    }
                } elseif { ! [info exists width] } {
                    if { 2 != [scan $value "%d %d" width height] } {
                        close $fp
                        error "Cannot parse values for width and height"
                    }
                    if { $width <= 0 } {
                        close $fp
                        error "Invalid width value: $width (must be greater than zero)"
                    }
                    if { $height <= 0 } {
                        close $fp
                        error "Invalid height value: $height (must be greater than zero)"
                    }
                } elseif { ! [info exists maxColor] } {
                    if { 1 != [scan $value "%d" maxColor] } {
                        close $fp
                        error "Cannot parse value for maxColor"
                    }
                    if { $maxColor < 0 } {
                        close $fp
                        error "Invalid value for maximum color: $maxColor (must be greater or equal to zero)"
                    }
                    if { $maxColor >= 65536 } {
                        close $fp
                        error "Invalid value for maximum color: $maxColor (must be less than 65536)"
                    }
                }
            }

            set headerDict [dict create]

            set pixelType [_GetPixelTypeName $maxColor]
            set numChan   [_GetNumChans $type]
            set byteOrder "Motorola"

            # Store the values contained in the header.
            dict set headerDict Magic     "PPM"
            dict set headerDict Width     $width
            dict set headerDict Height    $height
            dict set headerDict NumChan   $numChan
            dict set headerDict ByteOrder $byteOrder
            dict set headerDict ScanOrder "TopDown"
            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

            # Store additional information for reading the image data.
            dict set headerDict IsAscii [_IsAscii $type]

            return $headerDict
        }

        proc ReadImageHeader { ppmImgFile args } {
            # Read the header of a PPM image file.
            #
            # ppmImgFile - File name of the PPM image.
            # args       - Optional format specific arguments. Currently none.
            #
            # See [::pawt::ReadImageHeader] for the default entries of a
            # header dictionary.
            #
            # In addition to the default header dictionary entries, 
            # the following keys are supported by the PPM format:
            # IsAscii - Boolean flag indicating, if the values are stored in ASCII or binary.
            #
            # 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 $ppmImgFile "r"} fp]
            if { $retVal != 0 } {
                error "Cannot open file $ppmImgFile"
            }
            fconfigure $fp -translation binary
            set headerDict [_GetImageHeader $fp [file tail $ppmImgFile]]
            close $fp
            return $headerDict
        }

        proc _PutImageHeader { fp headerDict isAscii } {
            set pixelType [dict get $headerDict PixelType]
            set numChans  [dict get $headerDict NumChan]
            
            set maxColors [_GetMaxColors $pixelType]
            set ppmType   [_GetPpmType $numChans $isAscii]

            set headerStr ""
            append headerStr [format "%s" $ppmType]
            append headerStr "\n"
            append headerStr [format "%d " [dict get $headerDict Width]]
            append headerStr [format "%d"  [dict get $headerDict Height]]
            append headerStr "\n"
            append headerStr [format "%d"  $maxColors]
            append headerStr "\n"
            puts -nonewline $fp $headerStr
        }

        #
        # Image data read/write procedures.
        #

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

            set numBytes [dict get $headerDict NumByte]
            set numPixel [dict get $headerDict NumPixel]
            set numChans [dict get $headerDict NumChan]
            set numValues [expr { $numPixel * $numChans }]
            if { [dict get $headerDict IsAscii] } {
                set fmt [dict get $headerDict ScanFormat]
                set valueList [read $fp]
                if { [llength $valueList] < $numValues } {
                    error "Unexpected EOF in [dict get $headerDict FileName]"
                }
                set count 0
                foreach valueStr $valueList {
                    set trimmedValue [string trimleft $valueStr "0"]
                    if { $trimmedValue eq "" } {
                        set value 0
                    } else {
                        set value $trimmedValue
                    }
                    append valueData [binary format $fmt $value]
                    incr count
                    if { $count >= $numValues } {
                        break
                    }
                }
                dict set myDict Data $valueData
            } else {
                dict set myDict Data [read $fp $numBytes]
            }
        }

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

            set byteOrder [dict get $myDict Header ByteOrder]
            set pixelSize [dict get $myDict Header PixelSize]
            if { $isAscii } {
                set imgAsList [pawt GetImageAsList myDict -scanorder "TopDown"]
                foreach pixel $imgAsList {
                    puts $fp $pixel
                }
            } else {
                if { $byteOrder eq "Intel" } {
                    set scanFmt "cu*"
                    if { $pixelSize == 2 } {
                        set scanFmt "Su*"
                    }
                    set imgAsList [pawt GetImageAsList myDict -scanorder "TopDown"]
                    puts -nonewline $fp [binary format $scanFmt $imgAsList]
                } else {
                    puts -nonewline $fp [dict get $myDict Data]
                }
            }
        }

        proc ReadImageFile { ppmImgFile args } {
            # Read a PPM image file into a dict.
            #
            # ppmImgFile - File name of the PPM image.
            # args       - Optional format specific arguments. Currently none.
            #
            # 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 $ppmImgFile "r"} fp]
            if { $retVal != 0 } {
                error "Cannot open file $ppmImgFile"
            }
            fconfigure $fp -translation binary

            set imgDict [dict create]
            set headerDict [_GetImageHeader $fp [file tail $ppmImgFile]]
            _GetImageData $fp $headerDict imgDict
            dict append imgDict Header $headerDict

            close $fp
            return $imgDict
        }

        proc WriteImageFile { imgDict ppmImgFile args } {
            # Write the values of an image dict into a PPM image file.
            #
            # imgDict    - Image dictionary.
            # ppmImgFile - File name of the image.
            # args       - Optional format specific arguments described below.
            #
            # -datatype <string>  - Store image data as `binary`or `ascii` values.
            #                       If option is not specified, the default value is `binary`.
            # -pixeltype <string> - Store image data as `byte` or `short` values.
            #                       If option is not specified, the `PixelType` value of the
            #                       dictionary is used.
            #
            # * PPM files are always stored in `TopDown` scan order.
            # * Binary PPM files are always stored in `Motorola` byte order.
            #
            # 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 \
                -datatype  "binary" \
                -pixeltype "" \
            ]

            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 isAscii false
            if { [dict get $opts "-datatype"] eq "ascii" } {
                set isAscii true
            }

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

            _PutImageHeader $fp [dict get $imgDict Header] $isAscii
            _PutImageData   $fp imgDict $isAscii
            close $fp
        }
    }
}
