{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE LinearTypes #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE NoImplicitPrelude #-}
{-# OPTIONS_HADDOCK hide #-}

module Control.Optics.Linear.Internal
  ( -- * Types
    Optic_ (..),
    Optic,
    Iso,
    Iso',
    Lens,
    Lens',
    Prism,
    Prism',
    Traversal,
    Traversal',

    -- * Composing optics
    (.>),

    -- * Common optics
    swap,
    assoc,
    _1,
    _2,
    _Left,
    _Right,
    _Just,
    _Nothing,
    traversed,

    -- * Using optics
    get,
    set,
    gets,
    setSwap,
    match,
    build,
    over,
    overU,
    traverseOf,
    traverseOfU,
    toListOf,
    lengthOf,
    reifyLens,
    withIso,
    withLens,
    withPrism,

    -- * Constructing optics
    iso,
    lens,
    prism,
    traversal,
  )
where

import qualified Control.Arrow as NonLinear
import qualified Control.Functor.Linear as Control
import Data.Bifunctor.Linear (SymmetricMonoidal)
import qualified Data.Bifunctor.Linear as Bifunctor
import Data.Functor.Compose hiding (getCompose)
import Data.Functor.Linear
import qualified Data.Profunctor.Kleisli.Linear as Linear
import Data.Profunctor.Linear
import Data.Void
import GHC.Exts (FUN)
import GHC.Types
import Prelude.Linear
import qualified Prelude

newtype Optic_ arr s t a b = Optical (a `arr` b -> s `arr` t)

type Optic c s t a b =
  forall arr. c arr => Optic_ arr s t a b

type Iso s t a b = Optic Profunctor s t a b

type Iso' s a = Iso s s a a

type Lens s t a b = Optic (Strong (,) ()) s t a b

type Lens' s a = Lens s s a a

type Prism s t a b = Optic (Strong Either Void) s t a b

type Prism' s a = Prism s s a a

type Traversal s t a b = Optic Wandering s t a b

type Traversal' s a = Traversal s s a a

swap :: SymmetricMonoidal m u => Iso (a `m` b) (c `m` d) (b `m` a) (d `m` c)
swap = iso Bifunctor.swap Bifunctor.swap

assoc :: SymmetricMonoidal m u => Iso (a `m` (b `m` c)) (d `m` (e `m` f)) ((a `m` b) `m` c) ((d `m` e) `m` f)
assoc = iso Bifunctor.lassoc Bifunctor.rassoc

(.>) :: Optic_ arr s t a b -> Optic_ arr a b x y -> Optic_ arr s t x y
Optical f .> Optical g = Optical (f Prelude.. g)

infixr 9 .> -- same fixity as lens..>

lens :: (s %1 -> (a, b %1 -> t)) -> Lens s t a b
lens k = Optical $ \f -> dimap k (\(x, g) -> g $ x) (first f)

prism :: (b %1 -> t) -> (s %1 -> Either t a) -> Prism s t a b
prism b s = Optical $ \f -> dimap s (either id id) (second (rmap b f))

traversal :: (forall f. Control.Applicative f => (a %1 -> f b) -> s %1 -> f t) -> Traversal s t a b
traversal trav = Optical $ wander trav

_1 :: Lens (a, c) (b, c) a b
_1 = Optical first

_2 :: Lens (c, a) (c, b) a b
_2 = Optical second

_Left :: Prism (Either a c) (Either b c) a b
_Left = Optical first

_Right :: Prism (Either c a) (Either c b) a b
_Right = Optical second

_Just :: Prism (Maybe a) (Maybe b) a b
_Just = prism Just (maybe (Left Nothing) Right)

_Nothing :: Prism' (Maybe a) ()
_Nothing = prism (\() -> Nothing) Left

traversed :: Traversable t => Traversal (t a) (t b) a b
traversed = Optical $ wander traverse

over :: Optic_ (FUN 'One) s t a b -> (a %1 -> b) -> s %1 -> t
over (Optical l) f = l f

traverseOf :: Optic_ (Linear.Kleisli f) s t a b -> (a %1 -> f b) -> s %1 -> f t
traverseOf (Optical l) f = Linear.runKleisli (l (Linear.Kleisli f))

toListOf :: Optic_ (NonLinear.Kleisli (Const [a])) s t a b -> s -> [a]
toListOf l = gets l (\a -> [a])

get :: Optic_ (NonLinear.Kleisli (Const a)) s t a b -> s -> a
get l = gets l Prelude.id

gets :: Optic_ (NonLinear.Kleisli (Const r)) s t a b -> (a -> r) -> s -> r
gets (Optical l) f s = getConst' (NonLinear.runKleisli (l (NonLinear.Kleisli (Const Prelude.. f))) s)

set :: Optic_ (->) s t a b -> b -> s -> t
set (Optical l) x = l (const x)

setSwap :: Optic_ (Linear.Kleisli (Compose (FUN 'One b) ((,) a))) s t a b -> s %1 -> b %1 -> (a, t)
setSwap (Optical l) s = getCompose (Linear.runKleisli (l (Linear.Kleisli (\a -> Compose (\b -> (a, b))))) s)

match :: Optic_ (Market a b) s t a b -> s %1 -> Either t a
match (Optical l) = Prelude.snd (runMarket (l (Market id Right)))

build :: Optic_ (Linear.CoKleisli (Const b)) s t a b -> b %1 -> t
build (Optical l) x = Linear.runCoKleisli (l (Linear.CoKleisli getConst')) (Const x)

-- XXX: move this to Prelude

-- | Linearly typed patch for the newtype deconstructor. (Temporary until
-- inference will get this from the newtype declaration.)
getConst' :: Const a b %1 -> a
getConst' (Const x) = x

lengthOf :: MultIdentity r => Optic_ (NonLinear.Kleisli (Const (Sum r))) s t a b -> s -> r
lengthOf l s =
  (gets l (const (Sum one)) s) & \case
    Sum r -> r

-- XXX: the below two functions will be made redundant with multiplicity
-- polymorphism on over and traverseOfU
overU :: Optic_ (->) s t a b -> (a -> b) -> s -> t
overU (Optical l) f = l f

traverseOfU :: Optic_ (NonLinear.Kleisli f) s t a b -> (a -> f b) -> s -> f t
traverseOfU (Optical l) f = NonLinear.runKleisli (l (NonLinear.Kleisli f))

iso :: (s %1 -> a) -> (b %1 -> t) -> Iso s t a b
iso f g = Optical (dimap f g)

withIso :: Optic_ (Exchange a b) s t a b -> ((s %1 -> a) -> (b %1 -> t) -> r) -> r
withIso (Optical l) f = f fro to
  where
    Exchange fro to = l (Exchange id id)

withPrism :: Optic_ (Market a b) s t a b -> ((b %1 -> t) -> (s %1 -> Either t a) -> r) -> r
withPrism (Optical l) f = f b m
  where
    Market b m = l (Market id Right)

-- XXX: probably a direct implementation would be better
withLens ::
  Optic_ (Linear.Kleisli (Compose ((,) a) (FUN 'One b))) s t a b ->
  (forall c. (s %1 -> (c, a)) -> ((c, b) %1 -> t) -> r) ->
  r
withLens l k = k (Bifunctor.swap . (reifyLens l)) (uncurry ($))

reifyLens :: Optic_ (Linear.Kleisli (Compose ((,) a) (FUN 'One b))) s t a b -> s %1 -> (a, b %1 -> t)
reifyLens (Optical l) s = getCompose (Linear.runKleisli (l (Linear.Kleisli (\a -> Compose (a, id)))) s)

-- linear variant of getCompose
getCompose :: Compose f g a %1 -> f (g a)
getCompose (Compose x) = x