(*

The MIT License (MIT)

Copyright © 2023 Dean Lee

Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

*)

open System
open System.IO

// note: this program is purely functional except for the lack of monads for handling the stream objects

// byte handling functions

let readAsChars (bytes : byte list) =
    let rec recursion (list : string list) = if list.Length = 1 then list.Head else list.Head + recursion list.Tail
    List.map (fun x -> Convert.ToChar(x : byte)) bytes |> List.map (fun x -> Convert.ToString(x)) |> recursion

let readAsSigned (byte : byte) = if byte < 128uy then int byte else (int byte) - 256

let readAsUnsigned (bytes : byte list) =
    let uints = List.rev bytes |> List.map (fun x -> uint32 x)
    match uints.Length with
    | 1 -> uints.Head
    | 2 -> uints.Head + (uints.[1] <<< 8)
    | 4 -> uints.Head + (uints.[1] <<< 8) + (uints.[2] <<< 16) + (uints.[3] <<< 24)
    | _ -> 0xffffffffu // probably not the best solution from an api perspective

let writeU16 (num : uint16) =
    let byte0 = byte ((num <<< 8) >>> 8)
    let byte8 = byte (num >>> 8)
    [byte8; byte0]
    
let writeU32 (num : uint32) =
    let byte0 = byte ((num <<< 24) >>> 24)
    let byte8 = byte ((num <<< 16) >>> 24)
    let byte16 = byte ((num <<< 8) >>> 24)
    let byte24 = byte (num >>> 24)
    [byte24; byte16; byte8; byte0]

// list handling functions

let inline (..>) (element : 'a) (tuple : 'a list * 'a list) = (element :: fst tuple, snd tuple)

let splitList (count : int) (list : 'a list) =
    let rec recursion (count : int) (list : 'a list) = if count = 1 then ([list.Head], list.Tail) else list.Head ..> (recursion (count - 1) list.Tail)
    match count with
    | _ when count < 0 -> failwith "Count cannot be negative."
    | _ when count > list.Length -> failwith "Count must not exceed list length."
    | 0 -> ([], list)
    | _ -> recursion count list

let divList (interval : int) (list : 'a list) =
    let rec recursion (list : 'a list list) =
        let tuple = splitList interval list.Head
        if list.Head.Length <= interval then list else fst tuple :: recursion [snd tuple]
    if interval = 0 then [list] else recursion [list]
    
let diceList (dicer : int list) (list : 'a list) =
    let rec recursion (dicer : int list) (list : 'a list list) =
        let tuple = splitList dicer.Head list.Head
        if dicer.Length = 1 then list else fst tuple :: recursion dicer.Tail [snd tuple]
    if List.sum dicer <> list.Length then failwith "List length must equal sum of dicer values."
    if dicer.Length = 0 then [list] else recursion dicer [list]
    
let processList (processor : 'a list -> 'a list * 'a list) (list : 'a list) =
    let rec recursion (list : 'a list list) =
        let tuple = processor list.Head
        if (snd tuple).Length = 0 then list else fst tuple :: recursion [snd tuple]
    recursion [list]

// stream handling functions

let readStream (stream : Stream) =
    let rec recursion (list : int list) =
        let byteInt = stream.ReadByte()
        if byteInt < 0 then list else byteInt :: recursion list
    stream.Seek(0L, SeekOrigin.Begin) |> ignore
    let result = List.map (fun x -> byte x) (recursion [])
    stream.Seek(0L, SeekOrigin.Begin) |> ignore
    result
    
let writeStream (stream : Stream) (byteList : byte list) =
    let rec recursion (byteList : byte list) = if byteList.Length = 0 then () else stream.WriteByte(byteList.Head); recursion byteList.Tail
    stream.Seek(0L, SeekOrigin.Begin) |> ignore
    let result = recursion byteList
    stream.Seek(0L, SeekOrigin.Begin) |> ignore
    result

// create crctable

let partialDiv (uint : uint32) =
    let rec recursion (acc : uint32) (count : int) =
        let acc = if acc % 2u = 0u then acc >>> 1 else (acc >>> 1) ^^^ 0xedb88320u
        if count = 0 then acc else recursion acc (count - 1)
    recursion uint 7

let crcTable = List.map partialDiv [0u .. 255u] |> List.toArray

// data types

type PsdHeaderData =
    { Signature : string
      Version : uint32
      Channels : uint32
      Height : uint32
      Width : uint32
      Depth : uint32
      ColorMode : uint32 }

    static member make (byteList : byte list) =
        let (header, byteList) = splitList 26 byteList
        let (signature, header) = splitList 4 header
        let signature = readAsChars signature
        
        if signature <> "8BPS" then failwith "Not a valid PSD file."
        
        let header = diceList [2; 6; 2; 4; 4; 2; 2] header |> List.map readAsUnsigned
        let headerData =
            { Signature = signature
              Version = header.Head
              Channels = header.[2]
              Height = header.[3]
              Width = header.[4]
              Depth = header.[5]
              ColorMode = header.[6] }
        (headerData, byteList)

type PsdLayerData =
    { ContentsRect : byte list
      Channels : uint32
      ChannelInfo : byte list
      BlendModeSig : string
      BlendModeKey : string
      Opacity : uint32
      Clipping : uint32
      Flags : byte }

    static member make (byteList : byte list) =
        let (contentsRect, byteList) = splitList 16 byteList
        let (channels, byteList) = splitList 2 byteList
        let channels = readAsUnsigned channels
        let (channelInfo, byteList) = splitList (6 * (int channels)) byteList
        let (rest, byteList) = splitList 12 byteList
        let rest = diceList [4; 4; 1; 1; 1; 1] rest
        let layerData =
            { ContentsRect = contentsRect
              Channels = channels
              ChannelInfo = channelInfo
              BlendModeSig = readAsChars rest.Head
              BlendModeKey = readAsChars rest.[1]
              Opacity = readAsUnsigned rest.[2]
              Clipping = readAsUnsigned rest.[3]
              Flags = rest.[4].Head }
        (layerData, byteList)

type PsdPixelData =
    { Compression : uint32
      Pixels : byte list }
    
    static member make (byteList : byte list) =
        let (compression, byteList) = splitList 2 byteList
        { Compression = readAsUnsigned compression
          Pixels = byteList }
          
type PsdData =
    { HeaderData : PsdHeaderData
      Layers : uint32
      LayerData : PsdLayerData
      ChannelImageData : PsdPixelData
      PixelData : PsdPixelData }
      
[<RequireQualifiedAccess>]
module PsdData =

    let readFile (filepath : string) =
        use source = new FileStream(filepath, FileMode.Open, FileAccess.Read)
        readStream source

    let splitSection (byteList : byte list) =
        let (length, byteList) = splitList 4 byteList
        let length = readAsUnsigned length
        splitList (int length) byteList
      
    let make (filepath : string) =
        let byteList = readFile filepath
        let (headerData, byteList) = PsdHeaderData.make byteList
        
        if headerData.Version <> 1u then failwith "Must be version 1."
        if headerData.Channels <> 4u then failwith "Only 4 channel images supported."
        if headerData.Depth <> 8u then failwith "Only 8 bit images supported."
        if headerData.ColorMode <> 3u then failwith "Only RGBA images supported."
        
        let (whatever, byteList) = splitSection byteList // remove color section
        let (whatever, byteList) = splitSection byteList // remove resource section
        let (layerMaskSec, byteList) = splitSection byteList
        let (layerInfo, whatever) = splitSection layerMaskSec
        let (layers, layerInfo) = splitList 2 layerInfo
        let layers = readAsUnsigned layers
        
        if layers <> 1u then failwith "Multi-layer images not supported."

        let (layerData, layerInfo) = PsdLayerData.make layerInfo
        let (whatever, layerInfo) = splitSection layerInfo // ignore layer name
        let channelImageData = PsdPixelData.make layerInfo
        let pixelData = PsdPixelData.make byteList
        { HeaderData = headerData
          Layers = layers
          LayerData = layerData
          ChannelImageData = channelImageData
          PixelData = pixelData }
      
type GeneralInfo =
    { Height : uint32
      Width : uint32
      Channels : uint32
      Depth : uint32 }

    static member make (psdData : PsdData) =
        { Height = psdData.HeaderData.Height
          Width = psdData.HeaderData.Width
          Channels = psdData.HeaderData.Channels
          Depth = psdData.HeaderData.Depth }
      
type PngHeader =
    { Data : byte list }
    
    static member make (info : GeneralInfo) =
        let height = writeU32 info.Height
        let width = writeU32 info.Width
        { Data = height @ width @ [8uy; 6uy; 0uy; 0uy ;0uy] }
      
type PngImage =
    { Data : byte list }
    
[<RequireQualifiedAccess>]
module PngImage =

    let rec oneToMany (repeats : int) (byte : byte) = if repeats = 0 then [byte] else byte :: oneToMany (repeats - 1) byte

    let decompressA (byteList : byte list) =
        let sbyte = readAsSigned byteList.Head
        match sbyte with
        | _ when sbyte = -128 -> ([byteList.Head], byteList.Tail)
        | _ when sbyte >= 0 -> splitList (int (byteList.Head) + 2) byteList
        | _ -> splitList 2 byteList

    let decompressB (byteList : byte list) =
        let sbyte = readAsSigned byteList.Head
        match sbyte with
        | _ when sbyte = -128 -> []
        | _ when sbyte >= 0 -> byteList.Tail
        | _ -> oneToMany -sbyte byteList.Tail.Head
        
    let decompressUnit (byteList : byte list) = processList decompressA byteList |> List.map decompressB |> List.concat
    
    let decompress (info : GeneralInfo) (byteList : byte list) =
        let height = int info.Height
        let width = int info.Width
        let channels = int info.Channels
        let (byteCounts, byteList) = splitList (channels * height * 2) byteList
        let byteCounts = divList 2 byteCounts |> List.map (fun x -> int (readAsUnsigned x))
        diceList byteCounts byteList |> List.map decompressUnit |> List.concat
        
    let mergeColors (list : byte list list) =
        let rec recursion (red : byte list) (green : byte list) (blue : byte list) (alpha : byte list) =
            let merged = [red.Head; green.Head; blue.Head; alpha.Head]
            match red.Length with
            | 1 -> [merged]
            | _ -> merged :: recursion red.Tail green.Tail blue.Tail alpha.Tail
        recursion list.Head list.[1] list.[2] list.[3] |> List.concat

    let addFilterType (width : int) (channels : int) (byteList : byte list) =
        divList (channels * width) byteList |> List.map (fun x -> 0uy :: x) |> List.concat

    let adler32 (byteList : byte list) =
        let rec recursion (acc : uint32) (acc2 : uint32) (uintList : uint32 list) =
            let acc = (acc - uintList.Head) % 65521u
            let acc2 = acc2 + acc
            if uintList.Length <= 2 then acc2 else recursion acc acc2 uintList.Tail
        let uintList = List.map (fun x -> uint32 x) byteList |> List.rev
        let sum = List.sum uintList
        let a = (sum + 1u) % 65521u
        let b = recursion sum a uintList 
        writeU16 (uint16 b) @ writeU16 (uint16 a)

    let deflate (byteList : byte list) =
        if byteList.Length > 65535 then failwith "Image of this size not yet supported."
        let len = writeU16 (uint16 byteList.Length)
        let nlen = List.map (fun x -> x ^^^ 255uy) len
        [0x78uy; 0xdauy; 128uy;] @ len @ nlen @ byteList @ (adler32 byteList)
    
    let make (info : GeneralInfo) (image : PsdPixelData) =
        let height = int info.Height
        let width = int info.Width
        let channels = int info.Channels
        let byteList =
            match image.Compression with
            | 0u -> image.Pixels
            | 1u -> decompress info image.Pixels
            | _ -> failwith "Only RLE compression supported." ; []
        let result = divList (height * width) byteList |> mergeColors |> addFilterType width channels |> deflate
        { Data = result }
    
type PngChunklet =
    { ChunkType : byte list
      ChunkData : byte list }

type PngChunk =
    { Data : byte list }

    static member computeCrc (byteList : byte list) =
        let rec recursion (crcAcc : uint32) (uintList : uint32 list) =
            let index = int ((crcAcc ^^^ uintList.Head) &&& 255u)
            let crcAcc = crcTable.[index] ^^^ (crcAcc >>> 8)
            if uintList.Tail.Length = 0 then crcAcc else recursion crcAcc uintList.Tail
        let uintList = List.map (fun x -> uint32 x) byteList
        (recursion 0xffffffffu uintList) ^^^ 0xffffffffu
        
    static member make (chunklet : PngChunklet) =
        let length = writeU32 (uint32 chunklet.ChunkData.Length)
        let byteList = chunklet.ChunkType @ chunklet.ChunkData
        let crc = writeU32 (PngChunk.computeCrc byteList)
        { Data = length @ byteList @ crc }

type PngData =
    { Data : byte list }

    static member make (header : PngHeader) (image : PngImage) =
        let signature = [137uy; 80uy; 78uy; 71uy; 13uy; 10uy; 26uy; 10uy]
        let headerChunklet =
            { ChunkType = [73uy; 72uy; 68uy; 82uy]
              ChunkData = header.Data }
        let imageChunklet =
            { ChunkType = [73uy; 68uy; 65uy; 84uy]
              ChunkData = image.Data }
        let endChunklet =
            { ChunkType = [73uy; 69uy; 78uy; 68uy]
              ChunkData = [] }
        let chunks = List.map PngChunk.make [headerChunklet; imageChunklet; endChunklet] |> List.map (fun x -> x.Data)
        { Data = signature @ (List.concat chunks) }

let makePng (filepath : string) (data : PngData) =
    let filepath = filepath.ToCharArray()
    let filepath = Array.toList filepath.[0 .. filepath.Length - 5] |> List.map (fun x -> Convert.ToString(x)) |> List.reduce (fun acc x -> acc + x)
    let filepath = filepath + ".png"
    use dest = new FileStream(filepath, FileMode.Create, FileAccess.Write)
    writeStream dest data.Data
    0

let convert (filepath : string) =
    let psdData = PsdData.make filepath
    let info = GeneralInfo.make psdData
    let pngHeader = PngHeader.make info
    let pngImage = PngImage.make info psdData.PixelData
    let pngData = PngData.make pngHeader pngImage
    makePng filepath pngData

[<EntryPoint>]
let main (args : string []) =
    match args.Length with
    | 1 -> if File.Exists(args.[0]) then convert args.[0] else printfn "File does not exist." ; 1
    | _ -> printfn "This command takes exactly 1 argument." ; 1