module Example where

import Ast
import Typecheck
import Control.Exception (catch, evaluate, SomeException, displayException)
import Control.Exception.Base
import System.IO (stderr, hPutStrLn)
import Data.Maybe
import Data.List

green, red, yellow, blue, magenta, cyan, white :: String -> String
green str = "\x1b[32m" ++ str ++ "\x1b[0m"
red str = "\x1b[31m" ++ str ++ "\x1b[0m"
yellow str = "\x1b[33m" ++ str ++ "\x1b[0m"
blue str = "\x1b[34m" ++ str ++ "\x1b[0m"
magenta str = "\x1b[35m" ++ str ++ "\x1b[0m"
cyan str = "\x1b[36m" ++ str ++ "\x1b[0m"
white str = "\x1b[37m" ++ str ++ "\x1b[0m"

printSuccess :: String -> IO ()
printSuccess msg = putStrLn $ green "Success:" ++ white msg

handleError :: SomeException -> IO ()
handleError e = hPutStrLn stderr $ red ("Error: " ++ displayException e)

printResult :: Show a => String -> a -> IO ()
printResult title result = do
    putStrLn $ green title
    print result

sampleClasses :: [Class]
sampleClasses = [
    Class "Person" [
        MethodDeclaration "void" "setAge" [ParameterDeclaration "int" "newAge"]
        (Block [
            LocalVariableDeclaration (VariableDeclaration "int" "age" (Just (Reference "newAge")))
        ]),
        MethodDeclaration "int" "getAge" [] (Return (Just (Reference "age"))),
        MethodDeclaration "Person" "Person" [ParameterDeclaration "int" "initialAge"] (Block [])
    ] [
        VariableDeclaration "int" "age" (Just (IntegerLiteral 25))
    ]
 ]

initialSymtab :: [(DataType, Identifier)]
initialSymtab = []

exampleExpression :: Expression
exampleExpression = BinaryOperation NameResolution (Reference "bob") (Reference "age")

exampleAssignment :: Expression
exampleAssignment = StatementExpressionExpression (Assignment "a" (IntegerLiteral 30))

exampleMethodCall :: Statement
exampleMethodCall = StatementExpressionStatement (MethodCall (Reference "this") "setAge" [IntegerLiteral 30])

exampleConstructorCall :: Statement
exampleConstructorCall = LocalVariableDeclaration (VariableDeclaration "Person" "bob" (Just (StatementExpressionExpression  (ConstructorCall "Person" [IntegerLiteral 30]))))

exampleNameResolution :: Expression
exampleNameResolution = BinaryOperation NameResolution (Reference "b") (Reference "age")

exampleBlockResolution :: Statement
exampleBlockResolution = Block [
    LocalVariableDeclaration (VariableDeclaration "Person" "bob" (Just (StatementExpressionExpression  (ConstructorCall "Person" [IntegerLiteral 30])))),
    LocalVariableDeclaration (VariableDeclaration "int" "age" (Just (StatementExpressionExpression (MethodCall (Reference "bob") "getAge" [])))),
    StatementExpressionStatement (MethodCall (Reference "bob") "setAge" [IntegerLiteral 30])
    ]

exampleBlockResolutionFail :: Statement
exampleBlockResolutionFail = Block [
    LocalVariableDeclaration (VariableDeclaration "Person" "bob" (Just (StatementExpressionExpression  (ConstructorCall "Person" [IntegerLiteral 30])))),
    LocalVariableDeclaration (VariableDeclaration "bool" "age" (Just (StatementExpressionExpression (MethodCall (Reference "bob") "getAge" [])))),
    StatementExpressionStatement (MethodCall (Reference "bob") "setAge" [IntegerLiteral 30])
    ]

exampleMethodCallAndAssignment :: Statement
exampleMethodCallAndAssignment = Block [
    LocalVariableDeclaration (VariableDeclaration "Person" "bob" (Just (StatementExpressionExpression  (ConstructorCall "Person" [IntegerLiteral 30])))),
    LocalVariableDeclaration (VariableDeclaration "int" "age" (Just (StatementExpressionExpression (MethodCall (Reference "bob") "getAge" [])))),
    StatementExpressionStatement (MethodCall (Reference "bob") "setAge" [IntegerLiteral 30]),
    LocalVariableDeclaration (VariableDeclaration "int" "a" Nothing),
    StatementExpressionStatement (Assignment "a" (Reference "age"))
    ]


exampleMethodCallAndAssignmentFail :: Statement
exampleMethodCallAndAssignmentFail = Block [
    LocalVariableDeclaration (VariableDeclaration "Person" "bob" (Just (StatementExpressionExpression  (ConstructorCall "Person" [IntegerLiteral 30])))),
    LocalVariableDeclaration (VariableDeclaration "int" "age" (Just (StatementExpressionExpression (MethodCall (Reference "bob") "getAge" [])))),
    StatementExpressionStatement (MethodCall (Reference "bob") "setAge" [IntegerLiteral 30]),
    StatementExpressionStatement (Assignment "a" (Reference "age"))
    ]

testClasses :: [Class]
testClasses = [
    Class "Person" [
        MethodDeclaration "Person" "Person" [ParameterDeclaration "int" "initialAge"] 
        (Block [
            Return (Just (Reference "this"))
        ]),
        MethodDeclaration "void" "setAge" [ParameterDeclaration "int" "newAge"] 
        (Block [
            LocalVariableDeclaration (VariableDeclaration "int" "age" (Just (Reference "newAge")))
        ]),
        MethodDeclaration "int" "getAge" [] 
        (Return (Just (Reference "age")))
    ] [
        VariableDeclaration "int" "age" Nothing  -- initially unassigned
    ],
    Class "Main" [
        MethodDeclaration "int" "main" [] 
        (Block [
            LocalVariableDeclaration (VariableDeclaration "Person" "bob" (Just (StatementExpressionExpression (ConstructorCall "Person" [IntegerLiteral 25])))),
            StatementExpressionStatement (MethodCall (Reference "bob") "setAge" [IntegerLiteral 30]),
            LocalVariableDeclaration (VariableDeclaration "int" "bobAge" (Just (StatementExpressionExpression (MethodCall (Reference "bob") "getAge" [])))),
            Return (Just (Reference "bobAge"))
        ])
    ] []
 ]

runTypeCheck :: IO ()
runTypeCheck = do
    catch (do
        print "====================================================================================="
        evaluatedExpression <- evaluate (typeCheckExpression exampleExpression [("bob", "Person")] sampleClasses)
        printSuccess "Type checking of expression completed successfully"
        printResult "Result Expression:" evaluatedExpression
      ) handleError

    catch (do
        print "====================================================================================="
        evaluatedAssignment <- evaluate (typeCheckExpression exampleAssignment [("a", "int")] sampleClasses)
        printSuccess "Type checking of assignment completed successfully"
        printResult "Result Assignment:" evaluatedAssignment
      ) handleError

    catch (do
        print "====================================================================================="
        evaluatedMethodCall <- evaluate (typeCheckStatement exampleMethodCall [("this", "Person"), ("setAge", "Person"), ("getAge", "Person")] sampleClasses)
        printSuccess "Type checking of method call this completed successfully"
        printResult "Result MethodCall:" evaluatedMethodCall
      ) handleError

    catch (do
        print "====================================================================================="
        evaluatedConstructorCall <- evaluate (typeCheckStatement exampleConstructorCall [] sampleClasses)
        printSuccess "Type checking of constructor call completed successfully"
        printResult "Result Constructor Call:" evaluatedConstructorCall
      ) handleError

    catch (do
        print "====================================================================================="
        evaluatedNameResolution <- evaluate (typeCheckExpression exampleNameResolution [("b", "Person")] sampleClasses)
        printSuccess "Type checking of name resolution completed successfully"
        printResult "Result Name Resolution:" evaluatedNameResolution
      ) handleError

    catch (do
        print "====================================================================================="
        evaluatedBlockResolution <- evaluate (typeCheckStatement exampleBlockResolution [] sampleClasses)
        printSuccess "Type checking of block resolution completed successfully"
        printResult "Result Block Resolution:" evaluatedBlockResolution
      ) handleError

    catch (do
        print "====================================================================================="
        evaluatedBlockResolutionFail <- evaluate (typeCheckStatement exampleBlockResolutionFail [] sampleClasses)
        printSuccess "Type checking of block resolution failed"
        printResult "Result Block Resolution:" evaluatedBlockResolutionFail
      ) handleError

    catch (do
        print "====================================================================================="
        evaluatedMethodCallAndAssignment <- evaluate (typeCheckStatement exampleMethodCallAndAssignment [] sampleClasses)
        printSuccess "Type checking of method call and assignment completed successfully"
        printResult "Result Method Call and Assignment:" evaluatedMethodCallAndAssignment
      ) handleError

    catch (do
        print "====================================================================================="
        evaluatedMethodCallAndAssignmentFail <- evaluate (typeCheckStatement exampleMethodCallAndAssignmentFail [] sampleClasses)
        printSuccess "Type checking of method call and assignment failed"
        printResult "Result Method Call and Assignment:" evaluatedMethodCallAndAssignmentFail
      ) handleError

    catch (do
        print "====================================================================================="
        let mainClass = fromJust $ find (\(Class className _ _) -> className == "Main") testClasses
        case mainClass of
            Class _ [mainMethod] _ -> do
                let result = typeCheckMethodDeclaration mainMethod [] testClasses
                printSuccess "Full program type checking completed successfully."
                printResult "Main method result:" result
      ) handleError

    catch (do
        print "====================================================================================="
        let typedProgram = typeCheckCompilationUnit testClasses
        printSuccess "Type checking of Program completed successfully"
        printResult "Typed Program:" typedProgram
        ) handleError