open System.Threading.Tasks
open FsToolkit.ErrorHandling
open FsToolkit.ErrorHandling.Operator.Validation

type TaskValidation<'a, 'err> = Task<Validation<'a, 'err>>

[<RequireQualifiedAccess>]
module TaskValidation =
  let inline map f (tv: TaskValidation<'a, 'err>) : TaskValidation<'b, 'err> = TaskResult.map f tv
  let inline bind (f: 'a -> TaskValidation<'b, 'err>) (tv: TaskValidation<'a, 'err>) : TaskValidation<'b, 'err> = TaskResult.bind f tv

  let inline bindValidation (f: 'a -> TaskValidation<'b, 'err>) (tv: Validation<'a, 'err>) : TaskValidation<'b, 'err> =
    tv |> Task.singleton |> TaskResult.bind f

  let inline bindTask (f: 'a -> TaskValidation<'b, 'err>) (tv: Task<'a>) : TaskValidation<'b, 'err> =
    tv |> Task.map Validation.ok |> TaskResult.bind f

  let inline ok value : TaskValidation<'a, 'err> = TaskResult.ok value
  let inline bad value : TaskValidation<'a, 'err> = TaskResult.error value 
  let inline error value : TaskValidation<'a, 'err> = Validation.error value |> Task.singleton

type TaskValidationBuilder() =
  member inline x.Bind(comp, func) = TaskValidation.bind func comp
  member inline _.Delay([<InlineIfLambda>] generator: unit -> TaskValidation<'T, 'TError>) = generator ()
  member inline _.Combine(expr1: TaskValidation<unit, 'err>, expr2: TaskValidation<'a, 'err>) = expr1 |> TaskValidation.bind (fun () -> expr2)
  member inline _.Combine(expr1: Task<unit>, expr2: TaskValidation<'a, 'err>) = expr1 |> TaskValidation.bindTask (fun () -> expr2)
  member inline _.Combine(expr1: Validation<unit, 'err>, expr2: TaskValidation<'a, 'err>) = expr1 |> TaskValidation.bindValidation (fun () -> expr2)
  member inline _.Return(value) = TaskValidation.ok value
  member inline _.ReturnFrom(value: Validation<'a, 'err>) = Task.singleton value
  member inline _.ReturnFrom(value: TaskValidation<'a, 'err>) = value
  member inline _.While
    (
      [<InlineIfLambda>] guard: unit -> bool,
      [<InlineIfLambda>] computation: unit -> TaskValidation<unit, 'TError>
    ) : TaskValidation<unit, 'TError> =
    FSharp.Control.Tasks.Affine.task {
      let mutable fin, result = false, Ok()

      while not fin && guard () do
        match! computation () with
        | Ok x -> x
        | Error _ as e ->
          result <- e
          fin <- true

      return result
    }
  member inline _.For(sequence: #seq<'T>, [<InlineIfLambda>] binder: 'T -> TaskValidation<unit, 'TError>) : TaskValidation<unit, 'TError> =
    FSharp.Control.Tasks.Affine.task {
      use enumerator = sequence.GetEnumerator()
      let mutable fin, result = false, Ok()

      while not fin && enumerator.MoveNext() do
        match! binder enumerator.Current with
        | Ok x -> x
        | Error _ as e ->
          result <- e
          fin <- true

      return result
    }
  member inline _.Zero() = TaskValidation.ok ()
  member inline _.Source(task: TaskValidation<_, _>) : TaskValidation<_, _> = task

// Having members as extensions gives them lower priority in
// overload resolution between Task<_> and TaskValidation<_,_>
// since TaskValidation<_,_> is also a Task<_>
[<AutoOpen>]
module TaskValidationCEExtensionsLower =
  type TaskValidationBuilder with
    member inline this.Source(t: ^TaskLike) : TaskValidation<'T, 'Error> =
      FSharp.Control.Tasks.Affine.task {
        let! r = t
        return Ok r
      }

[<AutoOpen>]
module TaskValidationCEExtensions =
  type TaskValidationBuilder with
    member inline _.Source(choice: Choice<_, _>) : TaskValidation<_, _> = choice |> Result.ofChoice |> Task.singleton
    member inline _.Source(result: Validation<_, _>) : TaskValidation<_, _> = Task.singleton result
    member inline _.Source(s: #seq<_>) = s
    member inline _.Source(asyncComputation: Async<_>) : TaskValidation<_, _> =
      FSharp.Control.Tasks.Affine.task {
        let! r = asyncComputation |> Async.StartAsTask
        return Ok r
      }
    member inline _.Source(task: Task<_>) : Task<Result<_, _>> =
      FSharp.Control.Tasks.Affine.task {
        let! r = task
        return Ok r
      }

[<AutoOpen>]
module TaskValidationCE =
  let taskValidation = TaskValidationBuilder()