{-# LANGUAGE OverloadedStrings    #-}
{-# LANGUAGE UndecidableInstances #-}

module MetaBrush.Document
  ( AABB(..), mkAABB
  , Document(..), DocumentContent(..)
  , emptyDocument
  , Stroke(..), StrokeHierarchy(..), visibleStrokes
  , StrokeSpline, _strokeSpline, overStrokeSpline
  , PointData(..), BrushPointData(..), DiffPointData(..)
  , FocusState(..), Hoverable(..), HoverContext(..)
  , Guide(..), Ruler(..)
  , _selection, _coords, coords
  , addGuide, selectedGuide
  )
  where

-- base
import Control.Monad.ST
  ( RealWorld )
import Data.Coerce
  ( coerce )
import Data.Functor.Identity
  ( Identity(..) )
import Data.Semigroup
  ( Arg(..), Min(..), ArgMin )
import Data.Typeable
  ( Typeable )
import GHC.Generics
  ( Generic, Generic1 )
import GHC.TypeLits
  ( Symbol )

-- acts
import Data.Act
  ( Act(..), Torsor(..) )

-- containers
import Data.Map.Strict
  ( Map )
import qualified Data.Map.Strict as Map
  ( empty, insert )
import Data.Sequence
  ( Seq(..) )
import qualified Data.Sequence as Seq
  ( empty, singleton )

-- deepseq
import Control.DeepSeq
  ( NFData(..), NFData1, deepseq )

-- generic-lens
import Data.Generics.Product.Fields
  ( field' )

-- groups
import Data.Group
  ( Group(..) )

-- lens
import Control.Lens
  ( Lens'
  , set, view, over
  )

-- stm
import Control.Concurrent.STM
  ( STM )

-- text
import Data.Text
  ( Text )

-- transformers
import Control.Monad.Trans.Reader
  ( ReaderT, runReaderT )

-- MetaBrush
import Math.Bezier.Spline
  ( Spline(..), KnownSplineType )
import Math.Bezier.Stroke
  ( CachedStroke )
import Math.Module
  ( Module
    ( origin, (^+^), (^-^), (*^) )
  , Inner((^.^))
  , squaredNorm, quadrance
  )
import Math.Linear
  ( ℝ(..), T(..) )
import MetaBrush.Brush
  ( Brush, PointFields )
import MetaBrush.Records
import MetaBrush.Unique
  ( UniqueSupply, Unique, freshUnique )

--------------------------------------------------------------------------------

data AABB
  = AABB
  { topLeft, botRight :: !( ℝ 2 ) }
  deriving stock    ( Show, Generic )
  deriving anyclass NFData

mkAABB :: ℝ 2 -> ℝ 2 -> AABB
mkAABB ( ℝ2 x1 y1 ) ( ℝ2 x2 y2 ) = AABB ( ℝ2 xmin ymin ) ( ℝ2 xmax ymax )
  where
    ( xmin, xmax )
      | x1 > x2   = ( x2, x1 )
      | otherwise = ( x1, x2 )
    ( ymin, ymax )
      | y1 > y2   = ( y2, y1 )
      | otherwise = ( y1, y2 )

-- | Document, together with some extra metadata.
data Document
  = Document
  { displayName     :: !Text
  , mbFilePath      :: !( Maybe FilePath )
  , viewportCenter  :: !( ℝ 2 )
  , zoomFactor      :: !Double
  , documentUnique  :: Unique
  , documentContent :: !DocumentContent
  }
  deriving stock    ( Show, Generic )
  deriving anyclass NFData

-- | Main content of document (data which we kept track of throughout history).
data DocumentContent
  = Content
  { unsavedChanges :: !Bool
  , latestChange   :: !Text
  , guides         :: !( Map Unique Guide )
  , strokes        :: !( Seq StrokeHierarchy )
  }
  deriving stock    ( Show, Generic )
  deriving anyclass NFData

-- | Hierarchy for groups of strokes.
data StrokeHierarchy
  = StrokeGroup
    { groupName     :: !Text
    , groupVisible  :: !Bool
    , groupContents :: !( Seq StrokeHierarchy )
    }
  | StrokeLeaf
    { strokeLeaf :: !Stroke }
  deriving stock    ( Show, Generic )
  deriving anyclass NFData

visibleStrokes :: StrokeHierarchy -> Seq Stroke
visibleStrokes ( StrokeGroup { groupVisible, groupContents } )
  | groupVisible
  = foldMap visibleStrokes groupContents
  | otherwise
  = Empty
visibleStrokes ( StrokeLeaf { strokeLeaf } )
  | strokeVisible strokeLeaf
  = Seq.singleton strokeLeaf
  | otherwise
  = Empty

type StrokeSpline clo brushParams =
  Spline clo ( CachedStroke RealWorld ) ( PointData brushParams )

data Stroke where
  Stroke
    :: forall clo pointParams ( pointFields :: [ Symbol ] ) ( brushFields :: [ Symbol ] )
    .  ( KnownSplineType clo
       , pointParams ~ Record pointFields
       , PointFields pointFields, Typeable pointFields
       )
    =>
      { strokeName     :: !Text
      , strokeVisible  :: !Bool
      , strokeUnique   :: Unique
      , strokeBrush    :: !( Maybe ( Brush brushFields ) )
      , strokeSpline   :: !( StrokeSpline clo pointParams )
      }
    -> Stroke
deriving stock instance Show Stroke
instance NFData Stroke where
  rnf ( Stroke { strokeName, strokeVisible, strokeUnique, strokeBrush, strokeSpline } )
    = deepseq strokeSpline
    . deepseq strokeBrush
    . deepseq strokeUnique
    . deepseq strokeVisible
    $ rnf strokeName

_strokeSpline
  :: forall f
  .  Functor f
  => (  forall clo pointParams ( pointFields :: [ Symbol ] )
     .  ( KnownSplineType clo
        , pointParams ~ Record pointFields
        , PointFields pointFields
        )
     => StrokeSpline clo pointParams
     -> f ( StrokeSpline clo pointParams )
     )
  -> Stroke -> f Stroke
_strokeSpline f ( Stroke { strokeSpline = oldStrokeSpline, .. } )
  = ( \ newSpline -> Stroke { strokeSpline = newSpline, .. } ) <$> f oldStrokeSpline

overStrokeSpline
  :: (  forall clo pointParams ( pointFields :: [ Symbol ] )
     .  ( KnownSplineType clo
        , pointParams ~ Record pointFields
        , PointFields pointFields
        )
     => StrokeSpline clo pointParams
     -> StrokeSpline clo pointParams
     )
  -> Stroke -> Stroke
overStrokeSpline f = coerce ( _strokeSpline @Identity ( coerce . f ) )


data PointData params
  = PointData
  { pointCoords :: !( ℝ 2 )
  , pointState  :: FocusState
  , brushParams :: !params
  }
  deriving stock ( Show, Generic )
  deriving anyclass NFData

instance Act (T ( ℝ 2 )) (PointData params) where
  v • ( dat@( PointData { pointCoords = p } ) ) =
    dat { pointCoords = v • p }

data BrushPointData
  = BrushPointData
  { brushPointState :: FocusState }
  deriving stock ( Show, Generic )
  deriving anyclass NFData

data FocusState
  = Normal
  | Hover
  | Selected
  deriving stock ( Show, Eq, Generic )
  deriving anyclass NFData

instance Semigroup FocusState where
  Selected <> _        = Selected
  Normal   <> s        = s
  _        <> Selected = Selected
  s        <> Normal   = s
  _        <> _        = Hover
instance Monoid FocusState where
  mempty = Normal

emptyDocument :: Text -> Unique -> Document
emptyDocument docName unique =
  Document
    { displayName     = docName
    , mbFilePath      = Nothing
    , viewportCenter  = ℝ2 0 0
    , zoomFactor      = 1
    , documentUnique  = unique
    , documentContent =
      Content
      { unsavedChanges = False
      , latestChange   = "New document"
      , strokes        = Seq.empty
      , guides         = Map.empty
      }
    }

--------------------------------------------------------------------------------

data HoverContext
  = MouseHover     !( ℝ 2 )
  | RectangleHover !AABB
  deriving stock    ( Show, Generic )
  deriving anyclass NFData

instance Act ( T ( ℝ 2 ) ) HoverContext where
  v • MouseHover p = MouseHover ( v • p )
  v • RectangleHover ( AABB p1 p2 ) = RectangleHover ( AABB ( v • p1 ) ( v • p2 ) )

instance Act ( T ( ℝ 2 ) ) ( Maybe HoverContext ) where
  (•) v = fmap ( v • )

class Hoverable a where
  hovered :: Maybe HoverContext -> Double -> a -> FocusState

instance Hoverable ( ℝ 2 ) where
  hovered Nothing _ _ = Normal
  hovered ( Just ( MouseHover p ) ) zoom q
    | quadrance @( T ( ℝ 2 ) ) p q * zoom ^ ( 2 :: Int ) < 16
    = Hover
    | otherwise
    = Normal
  hovered ( Just ( RectangleHover ( AABB ( ℝ2 x1 y1 ) ( ℝ2 x2 y2 ) ) ) ) _ ( ℝ2 x y )
    | x >= x1 && x <= x2 && y >= y1 && y <= y2
    = Hover
    | otherwise
    = Normal

class HasSelection pt where
  _selection :: Lens' pt FocusState
instance HasSelection ( PointData brushParams ) where
  _selection = field' @"pointState"
instance HasSelection BrushPointData where
  _selection = field' @"brushPointState"

_coords :: Lens' ( PointData brushParams ) ( ℝ 2 )
_coords = field' @"pointCoords"

coords :: PointData brushParams -> ℝ 2
coords = view _coords

data FocusDifference
  = DifferentFocus
  | SameFocus
  deriving stock    ( Show, Generic )
  deriving anyclass NFData

instance Semigroup FocusDifference where
  SameFocus <> SameFocus = SameFocus
  _ <> _ = DifferentFocus

instance Monoid FocusDifference where
  mempty = SameFocus

instance Group FocusDifference where
  invert = id

data DiffPointData diffBrushParams
  = DiffPointData
  { diffVector :: !( T ( ℝ 2 ) )
  , diffParams :: !diffBrushParams
  , diffState  :: !FocusDifference
  }
  deriving stock    ( Show, Generic, Generic1, Functor, Foldable, Traversable )
  deriving anyclass ( NFData, NFData1 )

instance Module Double diffBrushParams => Semigroup ( DiffPointData diffBrushParams ) where
  DiffPointData v1 p1 s1 <> DiffPointData v2 p2 s2 =
    DiffPointData ( v1 <> v2 ) ( p1 ^+^ p2 ) ( s1 <> s2 )
instance Module Double diffBrushParams => Monoid ( DiffPointData diffBrushParams ) where
  mempty = DiffPointData mempty origin mempty
instance Module Double diffBrushParams => Group ( DiffPointData diffBrushParams ) where
  invert ( DiffPointData v1 p1 s1 ) =
    DiffPointData ( invert v1 ) ( -1 *^ p1 ) ( invert s1 )

instance ( Module Double diffBrushParams, Act diffBrushParams brushParams )
      => Act    ( DiffPointData diffBrushParams ) ( PointData brushParams ) where
  (•) ( DiffPointData { diffVector = dp, diffParams = db, diffState = focusDiff } )
    = over _coords ( dp • )
    . over ( field' @"brushParams" ) ( db • )
    . ( case focusDiff of { SameFocus -> id; DifferentFocus -> set ( field' @"pointState" ) Normal } )
instance ( Module Double diffBrushParams, Torsor diffBrushParams brushParams )
      => Torsor ( DiffPointData diffBrushParams ) ( PointData brushParams ) where
  ( PointData { pointCoords = p1, brushParams = b1, pointState = s1 } ) <-- ( PointData { pointCoords = p2, brushParams = b2, pointState = s2 } ) =
    DiffPointData
      { diffVector = p1 <-- p2
      , diffParams = b1 <-- b2
      , diffState  = if s1 == s2 then SameFocus else DifferentFocus
      }
instance Module Double brushParams => Module Double ( DiffPointData brushParams ) where
  origin = mempty
  (^+^) = (<>)
  x ^-^ y = x <> invert y
  d *^ DiffPointData v1 p1 s1 = DiffPointData ( d *^ v1 ) ( d *^ p1 ) s1


--------------------------------------------------------------------------------
-- Guides.

data Guide
  = Guide
  { guidePoint  :: !( ℝ 2 )       -- ^ point on the guide line
  , guideNormal :: !( T ( ℝ 2 ) ) -- ^ /normalised/ normal vector of the guide
  , guideFocus  :: !FocusState
  , guideUnique :: Unique
  }
  deriving stock    ( Show, Generic )
  deriving anyclass NFData

data Ruler
  = RulerCorner
  | LeftRuler
  | TopRuler
  deriving stock Show

-- | Try to select a guide at the given document coordinates.
selectedGuide :: ℝ 2 -> Document -> Maybe Guide
selectedGuide c ( Document { zoomFactor, documentContent = Content { guides } } ) =
  \case { Min ( Arg _ g ) -> g } <$> foldMap ( selectGuide_maybe c zoomFactor ) guides

selectGuide_maybe :: ℝ 2 -> Double -> Guide -> Maybe ( ArgMin Double Guide )
selectGuide_maybe c zoom guide@( Guide { guidePoint = p, guideNormal = n } )
  | sqDist * zoom ^ ( 2 :: Int ) < 4
  = Just ( Min ( Arg sqDist guide ) )
  | otherwise
  = Nothing
  where
    t :: Double
    t = ( c --> p ) ^.^ n
    sqDist :: Double
    sqDist = t ^ ( 2 :: Int ) / squaredNorm n

-- | Add new guide after a mouse drag from a ruler area.
addGuide :: UniqueSupply -> Ruler -> ℝ 2 -> Document -> STM Document
addGuide uniqueSupply ruler p doc = ( `runReaderT` uniqueSupply ) $ ( field' @"documentContent" . field' @"guides" ) insertNewGuides doc
  where
    insertNewGuides :: Map Unique Guide -> ReaderT UniqueSupply STM ( Map Unique Guide )
    insertNewGuides gs = case ruler of
      RulerCorner
        -> do
          uniq1 <- freshUnique
          uniq2 <- freshUnique
          let
            guide1, guide2 :: Guide
            guide1 = Guide { guidePoint = p, guideNormal = V2 0 1, guideFocus = Normal, guideUnique = uniq1 }
            guide2 = Guide { guidePoint = p, guideNormal = V2 1 0, guideFocus = Normal, guideUnique = uniq2 }
          pure ( Map.insert uniq2 guide2 . Map.insert uniq1 guide1 $ gs )
      TopRuler
        -> do
          uniq1 <- freshUnique
          let
            guide1 :: Guide
            guide1 = Guide { guidePoint = p, guideNormal = V2 0 1, guideFocus = Normal, guideUnique = uniq1 }
          pure ( Map.insert uniq1 guide1 gs )
      LeftRuler
        -> do
          uniq2 <- freshUnique
          let
            guide2 :: Guide
            guide2 = Guide { guidePoint = p, guideNormal = V2 1 0, guideFocus = Normal, guideUnique = uniq2 }
          pure ( Map.insert uniq2 guide2 gs )

instance Hoverable Guide where
  hovered ( Just ( MouseHover c ) ) zoom guide
    | Just _ <- selectGuide_maybe c zoom guide
    = Hover
    | otherwise
    = Normal
  hovered _ _ _ = Normal