import Control.Concurrent
import Control.Concurrent.STM
import Control.Monad.State
import Data.Complex
import Data.List
import Data.Maybe
import Data.Function
import qualified Graphics.UI.GLUT as G

-- We need a notion of time:

type Time = Int

-- The basic resource type is ore, and it comes from mines:

type Ore = Int
data OreMine = OreMine { oreid :: UnitID, remaining :: Int, at :: Location }
	deriving Show
instance Eq OreMine where
  (==) = (==) `on` oreid

-- We will have three types of units:

data UnitType = Builder  Ore
		| Refinery Ore
		| Fighter
	deriving (Show, Eq)

-- Each type has its own initial health:

type Health = Int
defaultHealth :: UnitType -> Health
defaultHealth (Builder _)  = 55
defaultHealth (Refinery _) = 8
defaultHealth Fighter      = 13

{-
We need a notion of location.  I'm going to use complex numbers for everything
since they come with a lot of functionality.
-}

type Location = Complex Double
type Velocity = Complex Double
type Speed = Double

-- There are a few commands we can give our units:

data Command = Attack Unit       -- Attack a unit
	     | Mine OreMine Time -- Use time to control the rate of mining
	     | GiveOre Ore Unit  -- Transfer ore from a refinery to a builder
	     | Build UnitType Health -- Build a unit.  Track build progress.
	     | Go Location       -- Go to a location
	     | Idle              -- Do nothing
   deriving Eq
instance Show Command where
  show (Attack _)   = "Attack"
  show (Mine m t)   = "Mine " ++ show m ++ " (" ++ show t ++ ")"
  show (GiveOre o u) = "GiveOre " ++ show o ++ " to " ++ show u
  show (Build ut h) = "Build " ++ show ut ++
			" (" ++ show (100*h `div` defaultHealth ut) ++ "%)"
  show (Go l)       = "Go " ++ show l
  show Idle         = "Idle"

-- There are also a few "effects" units can have on the game:

data Effect = AddOre Ore Unit
	    | RmOre  Ore OreMine
	    | CreateUnit Player UnitType Location
	deriving Show

-- Units need to have some attributes:

type UnitID = Int
data Unit = Unit { ident   :: UnitID,
		   owner   :: Player,
		   utype   :: UnitType,
		   pos     :: Location,
		   health  :: Health,
		   cmd     :: Command }
  deriving Show
instance Eq Unit where
  u == v = ident u == ident v

{-
Not all units can travel at the same speed.  Here's a function that tells us how
fast units move:
-}

speed :: Unit -> Double
speed u = case (utype u) of
		Builder _  -> 1
		Refinery _ -> 2
		Fighter    -> 3

-- Not all units cause the same amount of damage:

damage :: Unit -> Health
damage Unit{utype = Builder _}  = 0
damage Unit{utype = Refinery _} = 0
damage Unit{utype = Fighter}    = 3

-- We need a notion of player.

data Color = Red | Blue | Green
	deriving (Eq, Show)

type Name = String
data Player = Player Name Color
	deriving (Eq, Show)

-- And a notion of game:

data Game = Game { units   :: [Unit],
		   mines   :: [OreMine],
		   players :: [Player],
		   idents  :: [UnitID],
		   msgs :: [ (Unit, Command) ] }
instance Show Game where
  show g = "Players: " ++ (show $ players g) ++
	   "\nUnits: " ++ (show $ units g) ++
	   "\nMines: " ++ (show $ mines g)
makeGame :: [Player] -> Game
makeGame ps = Game{ units = [],
		    mines = [],
		    players = ps,
		    idents  = [1 .. ],
		    msgs = [] }

{-
We're going to use the State monad for, well, managing the state of the game.
But let's give it a fancy name, just to keep things clear:
-}

type GameSt a = State Game a

-- We need a command for adding a unit to the game.

addUnit :: Player -> UnitType -> Location -> GameSt Unit
addUnit p ut l
     = do g@Game{units=us, idents=i:is} <- get
	  let u = Unit { ident = i, owner = p, utype = ut, pos = l,
			 health = defaultHealth ut, cmd = Idle }
	  put g{units=u:us, idents=is}
	  return u

-- We also need a command for adding a mine to the game.

addMine :: Int -> Location -> GameSt OreMine
addMine r loc
     = do g@Game{mines=ms, idents=i:is} <- get
	  let m = OreMine { oreid = i, remaining = r, at = loc }
	  put g{mines=m:ms, idents=is}
	  return m

-- We also need a command for telling units what to do.

command :: Unit -> Command -> GameSt ()
command u c = do g <- get
		 put g{ msgs = (u,c):(msgs g) }
		 return ()

{-
Now we need a method for ``ticking'' the game.  First, we say that a player is
dead when they have no builders left.
-}

playerIsAlive :: Game -> Player -> Bool
playerIsAlive g p = foldl (\t -> \u -> t || isBuilder u) False (units g)
	where isBuilder u@Unit{utype = Builder _} = (owner u == p)
	      isBuilder _ = False

{-
A game tick consists of ticking each of the units, then leaving only the players
who still are alive.
-}

tick :: GameSt (Maybe Player)
tick =   do g@Game{ players = ps, units = us } <- get
	    let g' = foldl applyEffect
			   g{ units   = catMaybes mus',
			      players = (filter (playerIsAlive g) ps),
			      msgs = [] }
			   evnts
		(mus', evnts) = runState (sequence $ map (tickUnit g) us) []
	    put g'
	    if (length $ players g') == 1 then return (Just $ head $ players g')
					  else return Nothing

{-
We need a way to apply an effect to the game.  First, let's define the functions
that will carry out the effect:
-}

addOre :: Unit -> Ore -> GameSt ()
addOre u ore' = let updateOre u' = if u /= u' then u' else case utype u' of
			Refinery ore -> u'{utype = Refinery $ ore+ore'}
			Builder ore -> u'{utype = Builder $ ore+ore'}
			_ -> u'
		in modify (\g -> g{units = map updateOre (units g)})

removeOre :: OreMine -> Ore -> GameSt ()
removeOre m ore = let updateOre m' = if m /= m'
					then m'
					else m'{remaining = remaining m' - ore}
		  in modify (\g -> g{mines = map updateOre (mines g)})

applyEffect :: Game -> Effect -> Game
applyEffect g (AddOre ore' u)     = execState (addOre u ore') g
applyEffect g (RmOre ore m)       = execState (removeOre m ore) g
applyEffect g (CreateUnit p ut l) = execState (addUnit p ut l) g

{-
Unit ticking is the trickiest part.  Since the collection of units is a directed
cyclic graph (units can attack each other), it is the case that we need to ``tie
the knot.''

Here's what we do: when units need to interact with one another, we have them
point to the -previous- version of the unit.  If no such thing exists, then the
assocaited unit must be dead, so we revert to Idle.  If the unit does exist,
then we persist with our plans.  This keeps the pointers one-generation out of
date, but that's just how the rules work.

We need to do this to avoid infinite loops.  Consider the Attack sequence.
Before we attack, we check to make sure that the target is still alive (and thus
attackable).  But before we know if the unit is still alive, we need to know its
health, which is influenced by whether or not the unit is under attack!  And so
we get an infinite loop.

Units are able to alter their environment, either by killing each other, giving
each other ore, or building new units.  All of this happens in the tickUnit
function, so we need a way for tickUnit to report some effects.  By the rules of
our game, effects get applied after all units have been ticked.  This may lead
to some unusual gameplay -- for instance, an OreMine might have only 1 ore left,
but if 5 units are mining from it in the same game tick, then all 5 will get a
chance to mine the ore -before- the mine is depleted.
-}

tickUnit :: Game -> Unit -> State [Effect] (Maybe Unit)
tickUnit g u = do u' <- procCmd u
		  return $ (takeDamage $ deliverMsgs u') >>= checkPlayer
   where procCmd u@Unit{cmd=Attack v} =
	    let f v' = (goNear (pos v') u){ cmd = Attack v' }
	    in return $ maybe u{cmd=Idle} f (find (v==) $ units g)
	 procCmd u@Unit{cmd=Mine om t} =
	    if remaining om <= 0
		then return u{cmd=Idle}
		else let f om' = if pos u == at om
				  then if t == 0
					then do modify $ ([AddOre 1 u,
							   RmOre 1 om]++)
						return u{cmd=Mine om' mineDelay}
					else return u{cmd=Mine om' (t-1)}
				  else return $ (goTo (at om) u){cmd=Mine om' t}
		     in maybe (return u{cmd=Idle}) f (find (om==) $ mines g)
	 procCmd u@Unit{cmd=GiveOre o v} =
	    case utype u of
	     Refinery ore -> if ore < o
			     then return u{cmd=Idle}
			     else if pos u /= pos v
				   then return $ goTo (pos v) u
				   else do modify $ (AddOre o v:)
					   return u{ cmd = Idle,
						     utype = Refinery $ ore-o }
	     _ -> return u{cmd=Idle}
	 procCmd u@Unit{cmd=Build ut h} =
	    case utype u of
	     Builder o -> if h == defaultHealth ut
			   then do modify (CreateUnit (owner u) ut (pos u):)
				   return u{cmd=Idle}
			   else if o > 0 then return u{cmd=Build ut (h+1),
						       utype=Builder (o-1)}
					 else return u
	     _ -> return u{cmd=Idle}
	 procCmd u@Unit{cmd=Go l} =
	    return $ if pos u == l then u{cmd=Idle} else goTo l u
	 procCmd u@Unit{cmd=Idle} = return u
	 mineDelay = 5
	 goTo l u = let vec = (l - (pos u))
			vel = vec / (magnitude vec / speed u :+ 0)
		    in if (magnitude vec) > (magnitude vel)
			 then u{ pos = (pos u) + vel }
			 else u{ pos = l }
	 goNear l u = let dist = magnitude $ (pos u) - l
		      in if dist > attackRange then (goTo l u) else u
	 takeDamage u = let h' = foldl (checkAttack) (health u) $ units g
			in if h' <= 0 then Nothing else Just u{ health = h' }
	 checkAttack h v = let dist = magnitude $ (pos u) - (pos v)
			       d    = if isAttacking then (-1) else 0
			       isAttacking = (dist <= attackRange) &&
					     (cmd v == Attack u)
			   in (d * damage v)+h
	 attackRange = 10
	 deliverMsgs u = case (filter ((==) u . fst) (msgs g)) of
				((_,c):_) -> u{ cmd = c }
				[]	  -> u
	 checkPlayer u = if (playerIsAlive g $ owner u) then Just u else Nothing

{-
Here's a function that, given a game, plays it for n turns.  If there's a winner
by then, the winner gets returned.  If not, it returns Nothing.
-}

play :: Int -> GameSt (Maybe Player)
play n = if n == 0 then return Nothing
		   else do mp <- tick
			   case mp of
				Nothing -> play $ max (-1) (n-1)
				Just p  -> return $ Just p

{-
Once we start the game, it just kinda does its thing.  It'd be much better if
we could feed commands in during the simulation.  To do this, we need to
introduce some concurrency.  We will have one thread that plays the game, and
another to feed commands into it.  We will use Software Transactional Memory to
manage the shared game object.
-}

makeSharedGame :: [Player] -> IO (TVar Game)
makeSharedGame = atomically . newTVar . makeGame

playSharedGame :: (TVar Game) -> IO Player
playSharedGame tvg = do w <- atomically $
				do g <- readTVar tvg
				   let (w,g') = runState tick g
				   writeTVar tvg g'
				   return w
			case w of
				Nothing -> do threadDelay 50000
					      playSharedGame tvg
				Just p  -> return p

{-
We've got a bunch of commands for modifying game state.  Here's a function that
applies them to the shared game state.
-}

applyCommand :: (TVar Game) -> (GameSt a) -> IO a
applyCommand tvg c = atomically $
			do g <- readTVar tvg
			   let (a,g') = runState c g
			   writeTVar tvg g'
			   return a

{-
We might want to occasionally take a look at the state of the game.  Here's an
easy way to do that:
-}

printGame :: (TVar Game) -> IO ()
printGame tvg = do g <- atomically $ readTVar tvg
		   print g

-- Heck, we might even want a way to render the game.

ui :: (TVar Game) -> IO ()
ui tvg =
     do	G.getArgsAndInitialize
	G.initialDisplayMode G.$=
		[ G.DoubleBuffered, G.RGBMode, G.WithAlphaComponent,
		  G.WithDepthBuffer, G.Multisampling ]
	G.idleCallback G.$= Nothing
	G.initialWindowSize G.$= G.Size winW winH
	win <- G.createWindow "armada"
	G.currentWindow G.$= Just win
	G.depthFunc G.$= Just G.Less
	useView
	G.displayCallback G.$= drawGame win tvg
	G.reshapeCallback G.$= Nothing
	G.passiveMotionCallback G.$= Nothing
	G.keyboardMouseCallback G.$= Nothing
	G.idleCallback G.$= Just (G.postRedisplay $ Just win)
	forkIO $ G.mainLoop
	return ()

useView =
     do G.viewport G.$= (G.Position 0 0, G.Size winW winH)
	G.matrixMode G.$= G.Projection
	G.loadIdentity
	G.perspective 35 (fromIntegral winW / fromIntegral winH) 1 10000
	G.matrixMode G.$= G.Modelview 0
	G.loadIdentity

winW = 800 :: G.GLsizei
winH = 600 :: G.GLsizei
zDepth = (-200) :: G.GLdouble

drawGame :: G.Window -> (TVar Game) -> IO ()
drawGame win tvg =
     do G.clearColor G.$= G.Color4 0 0 0 0
	G.clear [ G.ColorBuffer, G.DepthBuffer ]
	G.loadIdentity
	g <- atomically $ readTVar tvg
	G.color (G.Color4 0.8 0.8 0.8 1.0 :: G.Color4 G.GLfloat)
	G.renderPrimitive G.Points $
		  mapM_ (\(x,y) -> G.vertex $ G.Vertex3 x y zDepth)
			[(x*5,y*5) | x <- [-10 .. 10], y <- [-10 .. 10]]
	sequence $ map drawUnit $ units g
	sequence $ map drawMine $ mines g
	G.swapBuffers

drawUnit :: Unit -> IO ()
drawUnit u =
     do G.loadIdentity
	let x :+ y = pos u
	    Player _ color = owner u
	G.color $ case color of
			Red   -> G.Color4 1.0 0.6 0.6 1.0 :: G.Color4 G.GLfloat
			Green -> G.Color4 0.8 1.0 0.4 1.0
			Blue  -> G.Color4 0.7 0.8 1.0 1.0
	G.translate $ G.Vector3 x y zDepth
	G.renderPrimitive
		G.Polygon
		(mapM_ (\(x,y) -> G.vertex $ G.Vertex3 x y 0) $ model u)
	G.translate $ G.Vector3 (x+2) (y-2) zDepth
	G.scale 0.09 0.09 (0.09::G.GLfloat)
	G.renderString G.Roman $ case utype u of
		Builder ore -> "H: " ++ show (health u) ++ "  O: " ++ show ore
		Refinery ore ->  "H: " ++ show (health u) ++ "  O: " ++ show ore
		Fighter -> "H: " ++ show (health u)
model :: Unit -> [(G.GLdouble, G.GLdouble)]
model u@Unit{utype = Builder ore} = [(-1, -1), (-1,1), (0,2), (1,1), (1,-1)]
model u@Unit{utype = Refinery ore} = [(-1, -1), (-1,1), (1,1), (1,-1)]
model u@Unit{utype = Fighter} = [(-1,0), (0,1), (1,0)]

drawMine :: OreMine -> IO ()
drawMine om =
     do G.loadIdentity
	G.color $ (G.Color4 0.2 0.2 0.2 1.0 :: G.Color4 G.GLdouble)
	let x :+ y = at om
	G.translate $ G.Vector3 x y zDepth
	G.renderObject G.Wireframe G.Octahedron
	G.translate $ G.Vector3 (x+2) (y-10) zDepth
	G.scale 0.09 0.09 (0.09::G.GLfloat)
	G.renderString G.Roman $ show (remaining om)

scrollTo (x,y) = G.lookAt (G.Vertex3 x y 50)
			  (G.Vertex3 x y 0)
			  (G.Vector3 0 1 0)

{-
Here's an example of how we use this from ghci.  No special extensions needed:
-}

buildDemo =
     do let tim = Player "Tim" Red
	    cyn = Player "Cyndi" Blue
	    chr = Player "Chris" Green
	shared <- makeSharedGame [tim, chr, cyn]
	mine1 <- applyCommand shared $ addMine 5000 ((-23) :+ (-14))
	tim_builder <- applyCommand shared $ addUnit tim (Builder 8) (0 :+ 0)
	chr_builder <- applyCommand shared $ addUnit chr (Builder 0) (9 :+ (-9))
	cyn_builder <- applyCommand shared $ addUnit cyn (Builder 0) ((-20) :+ 12)
	forkIO $ (playSharedGame shared) >>= (putStrLn . ("Winner: " ++) . show)
	tim_f1 <- applyCommand shared $ addUnit tim Fighter (0 :+ 0)
	tim_f2 <- applyCommand shared $ addUnit tim Fighter (0 :+ 0)
	tim_r1 <- applyCommand shared $ addUnit tim (Refinery 0) (0 :+ 0)
	applyCommand shared $ command tim_f1 (Attack chr_builder)
	applyCommand shared $ command tim_r1 (Mine mine1 0)
	ui shared
	scrollTo (0, 0)
	return (shared, tim_r1, tim_builder)

{-
Now try these commands out from ghci:

(shared, tim_r1, tim_builder) <- buildDemo
let p = printGame shared
applyCommand shared $ command tim_r1 (GiveOre 5 tim_builder)
applyCommand shared $ command tim_builder (Build Fighter 0)

-}

main = do print "Armada demo"
	  (shared, tim_r1, tim_builder) <- buildDemo
	  threadDelay 10000000
	  applyCommand shared $ command tim_r1 (GiveOre 5 tim_builder)
	  threadDelay 2000000
	  applyCommand shared $ command tim_builder (Build Fighter 0)
	  threadDelay 2000000
	  applyCommand shared $ command tim_builder (Go $ (-20) :+ (-10))
	  applyCommand shared $ command tim_r1 (Go $ (-15) :+ (10))