module WebSockets

open System
open System.IO
open System.Net
open System.Net.Sockets
open System.Security
open System.Text
open System.Threading

type TcpListener with
    member this.AsyncAcceptTcpClient() =
        Async.FromBeginEnd(this.BeginAcceptTcpClient, this.EndAcceptTcpClient)

#nowarn "25"
type HttpRequest =
  { headers:Map<string, string>;
    body:byte[] } with
      static member fromStream(stream:NetworkStream) =
          let bytes = Array.create 4096 0uy
          let bytesReadCount = stream.Read(bytes, 0, bytes.Length)
          let header = Encoding.UTF8.GetString(bytes.[0..bytesReadCount - 9])
          let lines = header.Split([|"\r"; "\n"|], StringSplitOptions.RemoveEmptyEntries)
          let kvs = lines.[1..lines.Length - 1]
                    |> Array.map (fun s -> s.Split([|":"|], 2, StringSplitOptions.RemoveEmptyEntries))
                    |> Array.map (function [|k; v|] -> (k.Trim(), v.Trim()) | _ -> ("Unknown", "Unknown"))
          { headers = new Map<string, string>(kvs);
            body = bytes.[bytesReadCount - 8 .. bytesReadCount - 1] }

let (|WebSocketRequest|_|) (r:HttpRequest) =
    if r.headers.["Upgrade"] = "WebSocket" && r.headers.["Connection"] = "Upgrade"
    then Some (r.headers.["Sec-WebSocket-Key1"], r.headers.["Sec-WebSocket-Key2"], r.body)
    else None

let computeKey (k:string) =
    let aux (sb:StringBuilder, cnt:int) = function
        | ' ' -> (sb, cnt + 1)
        | c -> if Char.IsDigit(c) then (sb.Append(c), cnt) else (sb, cnt)
    let (n, c) = k.ToCharArray() |> Array.fold aux (new StringBuilder(), 0)
    Int64.Parse(n.ToString()) / int64(c) |> int

let keyBytes (k:int) = k |> BitConverter.GetBytes |> Array.rev

let challenge (k1:string) (k2:string) (body:byte[]) =
    let k1 = computeKey k1
    let k2 = computeKey k2
    use md5 = Cryptography.MD5.Create()
    md5.ComputeHash(Array.concat([keyBytes k1; keyBytes k2; body]))

let handleClient (stream:NetworkStream) = async {
    try
        while true do
            let buf = Encoding.UTF8.GetBytes("hello")
            stream.WriteByte(0x00uy)
            stream.Write(buf, 0, buf.Length)
            stream.WriteByte(0xffuy)
            do! Async.Sleep(1000)
    with
        | :? IOException -> stream.Close()
}

let handleRequest (client:TcpClient) = async {
    let stream = client.GetStream()
    let request = HttpRequest.fromStream(stream)
    match request with
    | WebSocketRequest (k1, k2, body) ->
        let response = "HTTP/1.1 101 WebSocket Protocol Handshake\r\n" +
                       "Upgrade: WebSocket\r\n" +
                       "Connection: Upgrade\r\n" +
                       "Sec-WebSocket-Origin: " + request.headers.["Origin"] + "\r\n" +
                       "Sec-WebSocket-Location: ws://localhost:8080/test\r\n\r\n"
        let buf = Array.append (Encoding.UTF8.GetBytes(response)) (challenge k1 k2 body)
        do! stream.AsyncWrite(buf, 0, buf.Length)
        stream.WriteByte(0x0auy)
        let buf = Encoding.UTF8.GetBytes("hello")
        stream.WriteByte(0x00uy)
        stream.Write(buf, 0, buf.Length)
        stream.WriteByte(0xffuy)
        client.Close()
        //return! handleClient stream
    | _ -> client.Close()
}

let startServer port =
    let listener = new TcpListener(IPAddress.Any, port)
    let stop = ref false

    let main = async {
        listener.Start(50)
        printfn "Server started..."
        while not !stop do
            let! client = listener.AsyncAcceptTcpClient()
            Async.Start <| handleRequest client
    }

    Async.Start(main)
    { new IDisposable with
        member this.Dispose() = stop := true }

let main () =
    use server = startServer 8080
    Console.ReadKey() |> ignore

do main ()