module Main where

import Control.Exception (SomeException, try)
import Data.List (foldl', intercalate, isSuffixOf)
import Data.Maybe (fromMaybe, mapMaybe)
import qualified Data.Map.Strict as Map
import Data.Time.Clock (diffUTCTime, getCurrentTime)
import System.Directory (doesDirectoryExist, listDirectory)
import System.Environment (getArgs, getProgName)
import System.Exit (exitFailure, exitSuccess)
import System.IO (hPutStrLn, stderr)

import Tessera.Analyze
import Tessera.Check (checkProgram)
import Tessera.Env (Env, envConstructorTypes)
import Tessera.Eval
import Tessera.Graph (renderDot)
import Tessera.Parser (parseExpression, parseProgram)
import Tessera.Serialize hiding (renderDiagnostic)
import Tessera.Syntax
import Tessera.Tree

main :: IO ()
main = do
  arguments <- getArgs
  case arguments of
    ("check" : file : rest) -> commandCheck file rest
    ("explain" : file : rest) -> commandExplain file rest
    ("compile" : file : rest) -> commandCompile file rest
    ("run" : file : rest) -> commandRun file rest
    ("graph" : file : rest) -> commandGraph file rest
    ("bench" : directory : rest) -> commandBench directory rest
    _ -> usage

usage :: IO ()
usage = do
  name <- getProgName
  hPutStrLn stderr ("usage: " ++ name ++ " <check|explain|compile|run|graph|bench> ...")
  exitFailure

optionValue :: String -> [String] -> Maybe String
optionValue name arguments = case dropWhile (/= name) arguments of
  (_ : value : _) -> Just value
  _ -> Nothing

strategyFrom :: [String] -> Strategy
strategyFrom arguments = case optionValue "--strategy" arguments of
  Just "heuristic" -> Heuristic
  _ -> Leftmost

loadProgram :: FilePath -> IO Program
loadProgram file = do
  source <- readFile file
  case parseProgram source of
    Left diagnostic -> do
      hPutStrLn stderr (renderDiagnostic diagnostic)
      exitFailure
    Right program -> return program

loadEnv :: Program -> IO Env
loadEnv program = case checkProgram program of
  Left diagnostic -> do
    hPutStrLn stderr (renderDiagnostic diagnostic)
    exitFailure
  Right env -> return env

renderDiagnostic :: Diagnostic -> String
renderDiagnostic (Diagnostic span' message) =
  "line " ++ show (spanLine span') ++ " column " ++ show (spanColumn span') ++ ": " ++ message

renderType :: Type -> String
renderType ty = case ty of
  TVar name -> name
  TName name [] -> name
  TName name arguments -> name ++ " " ++ unwords (map renderType arguments)
  TTuple elements -> "(" ++ intercalate ", " (map renderType elements) ++ ")"

renderMatchReport :: MatchReport -> [String]
renderMatchReport report =
  [ "match " ++ reportMatchName report
  , if reportInconclusive report
      then "  coverage: inconclusive within the configured budget"
      else if reportExhaustive report
        then "  coverage: exhaustive"
        else "  coverage: missing " ++ maybe "" (intercalate ", " . map renderValue) (reportMissingCase report)
  ]
  ++ concatMap renderClause (reportClauseReports report)
  where
    renderClause clause =
      let status = case reportClauseStatus clause of
            ClauseUseful -> "useful"
            ClauseRedundant -> "redundant"
          example = case reportClauseExample clause of
            Just values -> " example " ++ intercalate ", " (map renderValue values)
            Nothing -> ""
      in ["  clause " ++ show (reportClauseIndex clause) ++ ": " ++ status ++ example]

commandCheck :: FilePath -> [String] -> IO ()
commandCheck file _ = do
  program <- loadProgram file
  env <- loadEnv program
  let reports = mapMaybe (analyzeMatch env) program
      failures = [report | report <- reports, not (reportExhaustive report) || hasRedundant report]
      output = concatMap renderMatchReport reports
  mapM_ putStrLn output
  if null failures then exitSuccess else exitFailure
  where
    hasRedundant report = any ((== ClauseRedundant) . reportClauseStatus) (reportClauseReports report)

commandExplain :: FilePath -> [String] -> IO ()
commandExplain file arguments = do
  program <- loadProgram file
  env <- loadEnv program
  case optionValue "--function" arguments of
    Nothing -> do
      hPutStrLn stderr "explain requires --function"
      exitFailure
    Just name -> case mapMaybe (analyzeMatch env) program of
      reports -> case filter ((== name) . reportMatchName) reports of
        (report : _) -> do
          mapM_ putStrLn (renderMatchReport report)
          putStrLn "inhabited types"
          mapM_ (\(typeName, flag) -> putStrLn ("  " ++ typeName ++ ": " ++ if flag then "inhabited" else "uninhabited")) (Map.toList (reportInhabitedTypes report))
          exitSuccess
        [] -> do
          hPutStrLn stderr ("unknown function " ++ name)
          exitFailure

commandCompile :: FilePath -> [String] -> IO ()
commandCompile file arguments = do
  source <- readFile file
  program <- case parseProgram source of
    Left diagnostic -> hPutStrLn stderr (renderDiagnostic diagnostic) >> exitFailure
    Right p -> return p
  env <- loadEnv program
  case optionValue "--function" arguments of
    Nothing -> hPutStrLn stderr "compile requires --function" >> exitFailure
    Just name -> case findMatch name program of
      Nothing -> hPutStrLn stderr ("unknown function " ++ name) >> exitFailure
      Just match' -> do
        let strategy = strategyFrom arguments
            tree = compileMatchFor env strategy match'
            output = fromMaybe (name ++ ".tree") (optionValue "-o" arguments)
        writeFile output (encodeArtifact name source strategy tree)
        putStrLn ("wrote " ++ output ++ " with " ++ strategyName strategy ++ " strategy, " ++ show (treeSize tree) ++ " nodes")

commandRun :: FilePath -> [String] -> IO ()
commandRun file arguments = do
  content <- readFile file
  artifact <- case decodeArtifact content of
    Left message -> hPutStrLn stderr message >> exitFailure
    Right a -> return a
  case optionValue "--function" arguments of
    Nothing -> hPutStrLn stderr "run requires --function" >> exitFailure
    Just name -> case findMatch name (artifactProgram artifact) of
      Nothing -> hPutStrLn stderr ("unknown function " ++ name) >> exitFailure
      Just match' -> case optionValue "--input" arguments of
        Nothing -> hPutStrLn stderr "run requires --input" >> exitFailure
        Just inputFile -> do
          inputText <- readFile inputFile
          let arities = Map.toList (Map.map (length . snd) (envConstructorTypes (artifactEnv artifact)))
          values <- case parseInputs (artifactEnv artifact) arities inputText of
            Left message -> hPutStrLn stderr message >> exitFailure
            Right v -> return v
          results <- mapM (runOne (artifactEnv artifact) match' (artifactTree artifact)) values
          if and results then exitSuccess else exitFailure
  where
    runOne env match' tree values = do
      let compiled = evaluateCompiled env (matchClauses match') tree values
          sourceResult = evaluateSource env (matchClauses match') values
      case (compiled, sourceResult) of
        (Right a, Right b) -> do
          putStrLn (intercalate ", " (map renderValue values) ++ " => " ++ renderValue a ++ (if a == b then "" else " (source disagrees: " ++ renderValue b ++ ")"))
          return (a == b)
        (Left message, _) -> do
          putStrLn (intercalate ", " (map renderValue values) ++ " => error " ++ message)
          return False
        (_, Left message) -> do
          putStrLn (intercalate ", " (map renderValue values) ++ " => source error " ++ message)
          return False

commandGraph :: FilePath -> [String] -> IO ()
commandGraph file arguments = do
  content <- readFile file
  artifact <- case decodeArtifact content of
    Left message -> hPutStrLn stderr message >> exitFailure
    Right a -> return a
  let output = fromMaybe "tree.dot" (optionValue "-o" arguments)
  writeFile output (renderDot (artifactTree artifact))
  putStrLn ("wrote " ++ output)

commandBench :: FilePath -> [String] -> IO ()
commandBench directory arguments = do
  exists <- doesDirectoryExist directory
  if not exists
    then hPutStrLn stderr ("not a directory: " ++ directory) >> exitFailure
    else do
      entries <- listDirectory directory
      let files = [directory ++ "/" ++ entry | entry <- entries, ".tess" `isSuffixOf` entry]
      mapM_ benchFile files
  where
    benchFile file = do
      program <- loadProgram file
      env <- loadEnv program
      let matches = [match' | match' <- program, isMatch match']
      mapM_ (benchMatch env) matches

    benchMatch env match' = do
      leftmostStart <- getCurrentTime
      let leftmostTree = compileMatchFor env Leftmost match'
      leftmostEnd <- getCurrentTime
      heuristicStart <- getCurrentTime
      let heuristicTree = compileMatchFor env Heuristic match'
      heuristicEnd <- getCurrentTime
      putStrLn (matchNameOf match')
      putStrLn ("  leftmost nodes " ++ show (treeSize leftmostTree) ++ " depth " ++ show (treeDepth leftmostTree) ++ " time " ++ show (diffUTCTime leftmostEnd leftmostStart))
      putStrLn ("  heuristic nodes " ++ show (treeSize heuristicTree) ++ " depth " ++ show (treeDepth heuristicTree) ++ " time " ++ show (diffUTCTime heuristicEnd heuristicStart))

isMatch :: Declaration -> Bool
isMatch (MatchDeclaration _ _ _ _) = True
isMatch _ = False

matchNameOf :: Declaration -> String
matchNameOf (MatchDeclaration name _ _ _) = name
matchNameOf _ = ""

findMatch :: String -> Program -> Maybe Declaration
findMatch name program = case [match' | match' <- program, matchNameOf match' == name, isMatch match'] of
  (match' : _) -> Just match'
  [] -> Nothing

compileMatchFor :: Env -> Strategy -> Declaration -> Tree
compileMatchFor env strategy (MatchDeclaration _ arguments _ clauses) =
  let types = map snd arguments
      rows = zip (map clausePatterns clauses) [0 .. length clauses - 1]
  in compileMatch env strategy types rows (initialOccurrences (length arguments))
compileMatchFor _ _ _ = Failure

parseInputs :: Env -> [(String, Int)] -> String -> Either String [[Value]]
parseInputs env arities text = go (filter (not . null) (lines text)) []
  where
    go [] acc = Right (reverse acc)
    go (line : rest) acc = case parseInputLine env arities line of
      Left message -> Left message
      Right values -> go rest (values : acc)

parseInputLine :: Env -> [(String, Int)] -> String -> Either String [Value]
parseInputLine env arities line =
  let pieces = map trim (splitSemicolon line)
  in mapM parsePiece pieces
  where
    parsePiece piece = case parseExpression arities piece of
      Left diagnostic -> Left ("input parse error: " ++ renderDiagnostic diagnostic)
      Right expression -> case evalExpr env Map.empty expression of
        Left message -> Left message
        Right value -> Right value

trim :: String -> String
trim = reverse . dropWhile (== ' ') . reverse . dropWhile (== ' ')

splitSemicolon :: String -> [String]
splitSemicolon [] = [""]
splitSemicolon text = case break (== ';') text of
  (piece, []) -> [piece]
  (piece, _ : rest) -> piece : splitSemicolon rest