{-# LANGUAGE OverloadedStrings   #-}
{-# LANGUAGE ScopedTypeVariables #-}

module MetaBrush.Draw
  ( DrawAnchor(..), anchorsAreComplementary
  , getOrCreateDrawAnchor, addToAnchor
  , anchorPos
  )
  where

-- base
import Control.Monad
  ( when )
import Data.Foldable
  ( for_ )
import Data.Functor.Identity
  ( Identity(..) )
import GHC.Generics
  ( Generic )
import GHC.TypeLits
  ( Symbol )

-- containers
import Data.Sequence
  ( Seq(..) )
import qualified Data.Map.Strict as Map
import qualified Data.Set as Set

-- deepseq
import Control.DeepSeq
  ( NFData )

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

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

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

-- transformers
import Control.Monad.Trans.Except
  ( Except )
import qualified Control.Monad.Trans.Except as Except
import Control.Monad.Trans.Reader
  ( runReaderT )

-- brush-strokes
import Math.Bezier.Spline
import Math.Bezier.Stroke
  ( invalidateCache )
import Math.Linear
  ( ℝ(..) )

-- MetaBrush
import MetaBrush.Assert
  ( assert )
import MetaBrush.Brush
  ( NamedBrush(..) )
import MetaBrush.Document
import MetaBrush.Hover
  ( inPointClickRange )
import MetaBrush.Layer
import MetaBrush.Records
import MetaBrush.Stroke
import MetaBrush.Unique
  ( Unique, UniqueSupply, freshUnique )

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

-- | A draw anchor, to continue drawing from one end of an existing stroke.
data DrawAnchor
  = DrawAnchor
    { anchorIsNew   :: !Bool
    , anchorStroke  :: !Unique
    , anchorIsAtEnd :: !Bool
    }
  deriving stock ( Show, Eq, Generic )
  deriving anyclass NFData

-- | Computes whether two anchors are the two ends of the same stroke.
anchorsAreComplementary :: DrawAnchor -> DrawAnchor -> Bool
anchorsAreComplementary
  ( DrawAnchor { anchorStroke = uniq1, anchorIsAtEnd = end1 } )
  ( DrawAnchor { anchorStroke = uniq2, anchorIsAtEnd = end2 } )
  = uniq1 == uniq2 && end1 /= end2

-- | Compute a draw anchor at the given position, e.g. to continue
-- drawing a stroke or to start a new one.
getOrCreateDrawAnchor
  :: UniqueSupply
  -> ℝ 2
  -> Document
  -> STM ( Document, DrawAnchor )
getOrCreateDrawAnchor uniqueSupply c doc@( Document { documentContent = oldContent, documentMetadata } ) =

  -- Deselect all points, and try to find a valid anchor for drawing
  -- (a path start/end point at mouse click point).

  case
    Except.runExcept $
      forStrokeHierarchy
        ( layerMetadata documentMetadata )
        ( strokeHierarchy oldContent )
        findAnchor
  of
    -- Anchor found: use it.
    Left anchor ->
      let newMeta :: DocumentMetadata
          newMeta = documentMetadata { selectedPoints = mempty }
      in pure ( doc { documentMetadata = newMeta }
              , anchor )
    -- No anchor found: start a new stroke (on a new stroke layer).
    Right {} -> do
      newStrokeUnique <- runReaderT freshUnique uniqueSupply
      let
        newSpline :: StrokeSpline Open ( Record ( '[] :: [ Symbol ] ) )
        newSpline =
          Spline { splineStart  = PointData c ( MkR ℝ0 )
                 , splineCurves = OpenCurves Empty
                 }
        newStroke :: Stroke
        newStroke =
          Stroke
            { strokeSpline  = newSpline
            , strokeBrush   = Nothing :: Maybe ( NamedBrush ( '[] :: [ Symbol ] ) )
            }
        newSel = StrokePoints $ Map.singleton newStrokeUnique ( Set.singleton FirstPoint )
        newMeta :: DocumentMetadata
        newMeta =
            set ( field' @"selectedPoints" ) newSel
          . over ( field' @"layerMetadata" . field' @"layerNames" ) ( Map.insert newStrokeUnique "Stroke" )
          $ documentMetadata
        newHierarchy :: StrokeHierarchy
        newHierarchy =
          ( strokeHierarchy oldContent )
            { topLevel = newStrokeUnique : topLevel ( strokeHierarchy oldContent )
            , content = Map.insert newStrokeUnique newStroke ( content ( strokeHierarchy oldContent ) )
            }
        newContent :: DocumentContent
        newContent = oldContent { strokeHierarchy = newHierarchy }
        newDoc' :: Document
        newDoc' =
          doc { documentMetadata = newMeta
              , documentContent  = newContent
              }
        anchor =
          DrawAnchor
            { anchorIsNew   = True
            , anchorStroke  = newStrokeUnique
            , anchorIsAtEnd = True
            }
      pure ( newDoc', anchor )
  where
    zoom = documentZoom documentMetadata

    findAnchor :: Unique -> Stroke -> StrokeMetadata -> Except DrawAnchor UpdateStroke
    findAnchor strokeUnique ( Stroke { strokeSpline }) ( StrokeMetadata { strokeVisible, strokeLocked } ) = do
      when ( strokeVisible && not strokeLocked ) $
        findAnchorSpline strokeSpline
      return PreserveStroke

      where
        findAnchorSpline
          :: forall clo brushParams
          .  SplineTypeI clo
          => StrokeSpline clo brushParams
          -> Except DrawAnchor ()
        findAnchorSpline spline =
          for_ ( endpointAnchor strokeUnique spline ) Except.throwE
          where
            -- See if we can anchor a drawing operation on a given (visible) stroke.
            endpointAnchor :: Unique -> StrokeSpline clo brushParams -> Maybe DrawAnchor
            endpointAnchor uniq ( Spline { splineStart, splineCurves } ) = case ssplineType @clo of
              SOpen
                | let
                    p0 = coords splineStart
                , inPointClickRange zoom c p0
                , let anchor =
                       DrawAnchor
                         { anchorIsNew   = False
                         , anchorStroke  = uniq
                         , anchorIsAtEnd = False
                         }
                -> Just anchor
                | OpenCurves ( _ :|> lastCurve ) <- splineCurves
                , let
                    pn :: ℝ 2
                    pn = coords ( openCurveEnd lastCurve )
                , inPointClickRange zoom c pn
                , let anchor =
                       DrawAnchor
                         { anchorIsNew   = False
                         , anchorStroke  = uniq
                         , anchorIsAtEnd = True
                         }
                -> Just anchor
              _ -> Nothing

anchorPos :: DrawAnchor -> Document -> ℝ 2
anchorPos ( DrawAnchor { anchorStroke, anchorIsAtEnd } )
  ( Document { documentContent = Content { strokeHierarchy = Hierarchy { content = strokes } } } ) =
    case Map.lookup anchorStroke strokes of
      Nothing -> error "draw anchor: stroke not present in hierarchy"
      Just ( Stroke { strokeSpline } ) ->
        coords $
          if anchorIsAtEnd
          then splineEnd   ( adjustSplineType @Open strokeSpline )
          else splineStart ( adjustSplineType @Open strokeSpline )

addToAnchor :: DrawAnchor -> Spline Open () ( PointData () ) -> Document -> Document
addToAnchor anchor newSpline doc@( Document { documentMetadata, documentContent = oldContent } ) =
  let
    strokes' =
      runIdentity $
        forStrokeHierarchy
          ( layerMetadata documentMetadata )
          ( strokeHierarchy oldContent )
          ( \ u s _ -> Identity $ updateStroke u s )
  in doc { documentContent = oldContent { strokeHierarchy = strokes' } }
  where

    updateStroke :: Unique -> Stroke -> UpdateStroke
    updateStroke strokeUnique stroke
      | strokeUnique == anchorStroke anchor
      , let
          updateSpline
            :: forall clo brushData
            .  SplineTypeI clo
            => StrokeSpline clo brushData -> StrokeSpline clo brushData
          updateSpline prevSpline
            | SOpen <- ssplineType @clo
            = if anchorIsAtEnd anchor
              then
                let
                  i0 = case splineCurves prevSpline of
                    OpenCurves ( _ :|> lastCurve ) ->
                      curveIndex ( curveData lastCurve ) + 1
                    _ -> 0
                  setBrushData :: PointData () -> PointData brushData
                  setBrushData = set ( field @"brushParams" ) ( brushParams ( splineEnd prevSpline ) )
                  setData = bimapSpline ( \ _ -> bimapCurve ( \ s -> CurveData i0 ( invalidateCache $ undefined s ) ) ( \ _ -> setBrushData ) ) setBrushData
                in prevSpline <> setData newSpline
              else
                let
                  i0 = case splineCurves prevSpline of
                    OpenCurves ( firstCurve :<| _ ) ->
                      curveIndex ( curveData firstCurve ) - 1
                    _ -> 0
                  setBrushData :: PointData () -> PointData brushData
                  setBrushData = set ( field @"brushParams" ) ( brushParams ( splineStart prevSpline ) )
                  setData = bimapSpline ( \ _ -> bimapCurve ( \ s -> CurveData i0 ( invalidateCache $ undefined s ) ) ( \ _ -> setBrushData ) ) setBrushData
                in setData ( reverseSpline newSpline ) <> prevSpline
            | otherwise
            = assert False ( "addToAnchor: trying to add to closed spline " <> show strokeUnique )
                prevSpline -- should never add to a closed spline
      = UpdateStrokeTo $ overStrokeSpline updateSpline stroke
      | otherwise
      = PreserveStroke