#!/usr/bin/env python3

# (c) 2018-2021 Vladimír Štill

import os
import os.path
import re
import shutil
import subprocess
import sys
import tempfile
import textwrap
import yaml
from typing import List, Pattern, TextIO, Optional, Dict, Any, TypeVar

from typing_extensions import Final

T = TypeVar("T")
FilePath = str

if len(sys.argv) < 3:
    print(f"usage: {sys.argv[0]} TEACHER_FILE STUDENT_FILE [-IDIR ...] "
          "[-pFD] [-cCONFIG] [--hint]", file=sys.stderr)
    sys.exit(1)


def except_hook(exctype, value, traceback):
    print("INTERNAL ERROR, please report this result;\n"
          f"terminating with error of type {exctype.__name__}")
    sys.__excepthook__(exctype, value, traceback)


sys.excepthook = except_hook

teacher: Final[str] = os.path.abspath(sys.argv[1])
student: Final[str] = os.path.abspath(sys.argv[2])
hint: bool = False
include: List[str] = []
journal_fd: int = 2
_config: Optional[Dict[str, Any]] = None


def _get_opt(val: str, *opts: str) -> Optional[str]:
    for opt in opts:
        if getattr(_get_opt, "next_opt", None) is None:
            if val.startswith(f"-{opt}"):
                opt_len = len(opt) + 1
                if len(val) > opt_len:
                    return val[opt_len:]
                else:
                    setattr(_get_opt, "next_opt", opt)
                    return None
        elif getattr(_get_opt, "next_opt", None) == opt:
            setattr(_get_opt, "next_opt", None)
            return val
    return None


for i in range(3, len(sys.argv)):
    v = sys.argv[i]
    if v == "--hint":
        hint = True
    if (val := _get_opt(v, "p")):
        journal_fd = int(val)
    elif (val := _get_opt(v, "i", "I")):
        include.append(val)
    elif (val := _get_opt(v, "c")):
        with open(val, "r") as h:
            _config = yaml.safe_load(h)


def from_source_file(path: str) -> Dict[str, Any]:
    pat = re.compile("-- @\\s+([^:]*):\\s*(.*)")
    collected = ""
    with open(path, "r") as h:
        for line in h:
            m = pat.match(line)
            if m is not None:
                collected += f"{m[1]}: {m[2]}\n"
    if not collected:
        return {}
    out = yaml.safe_load(collected)
    assert isinstance(out, dict)
    return out  # type: ignore


if _config is None:
    _config = from_source_file(teacher)


NO_DERIVING: Final[List[str]] = \
    ["NoDeriveDataTypeable", "NoDeriveFoldable",
     "NoDeriveFunctor", "NoDeriveGeneric", "NoDeriveLift",
     "NoDeriveTraversable", "NoEmptyDataDeriving",
     "NoGeneralizedNewtypeDeriving", "NoStandaloneDeriving"]

ALL_EXTS: Final[List[str]] = []
ST_EXTS: Final[List[str]] = ["Safe"] + ALL_EXTS + NO_DERIVING
T_EXTS: Final[List[str]] = ["TemplateHaskell", "Unsafe", "ExplicitForAll",
                            "ScopedTypeVariables", "DataKinds",
                            "StandaloneDeriving", "DeriveDataTypeable",
                            "TypeOperators", "KindSignatures",
                            "DerivingStrategies",
                            "GeneralizedNewtypeDeriving"] \
                            + ALL_EXTS
MAIN_EXTS: Final[List[str]] = T_EXTS

OPT_R: Final[Pattern[str]] = re.compile("-- @ +([^:]*)(?:: +(.*))?")
LIST_SPLIT_R: Final[Pattern[str]] = re.compile(" *, *| +")
ERR_HEADER: Final[Pattern[str]] = re.compile(r"^([^:]*)[.]hs:[0-9]*:[0-9]*:")
IGNORE: Final[Pattern[str]] = re.compile(r"^(-- .*)?\s*$")
AST_CHECK_DECL: Final[Pattern[str]] = re.compile(r"^([^\s:]*)\s*::\s*ASTCheck")

GHC_LINE_ERR: Final[Pattern[str]] = re.compile(r"^[0-9]*\s*\|")
COMPILE_LINE: Final[Pattern[str]] = re.compile(
    r"^(\[[0-9]+\s+of\s+[0-9]+\]|Linking .* ...)")
AST_CHECK_EXC: Final = re.compile(
    r"^\s+Exception when trying to run compile-time code:")
AST_CHECK_CALL_STACK: Final = re.compile(
    r"^(CallStack \(from HasCallStack\):"
    r"|\s+error, called at [.]/Teacher[.]hs:[0-9]+:[0-9]+ in main:Teacher"
    r"|\s+Code: runASTCheck)")
HINT_TYPE_MISSING_LEAD: Final = re.compile(
    r"^ *Top-level binding with no type signature:")
HINT_TYPE_MISSING_TYPE: Final = re.compile(
    r"^(?P<fun> *[^ ]* ::) (?P<type>.*)$")


class TestConf:
    def __init__(self, config: Dict[str, object]) -> None:
        self.inject: str = ""
        self.typecmp: bool = False
        self.typecmpex: bool = False
        self.typecmp_num_types: int = 0
        self.ast_checks: List[str] = []
        self.allow_hint = True
        self.raw: bool = False
        self.hlint = True
        self.config = config

    def reflection(self) -> bool:
        return self.get_typed("reflection", False, bool)

    def compiled(self) -> bool:
        return self.get_typed("compiled", False, bool)

    def exts(self) -> List[str]:
        return self.get_typed("exts", [], list)

    def student_exts(self) -> List[str]:
        return ST_EXTS + self.exts() + self.get_typed("student exts", [], list)

    def teacher_exts(self) -> List[str]:
        return T_EXTS + self.exts() + self.get_typed("teacher exts", [], list)

    def main_exts(self) -> List[str]:
        return MAIN_EXTS + self.exts() + self.get_typed("main exts", [], list)

    def imports(self) -> List[str]:
        return self.get_typed("imports", [], list)

    def require_type(self) -> bool:
        return self.get_typed("require type", False, bool)

    def allow_imports(self) -> bool:
        return self.get_typed("allow imports", True, bool)

    def warnings(self) -> List[str]:
        return self.get_typed("warnings", [], list)

    def question_path(self) -> str:
        return self.get_typed("question_path", teacher, str)

    def _get(self, key: str, default: object) -> object:
        return self.config.get(key, default)

    def get_typed(self, key: str, default: T, ty: type) -> T:
        val = self._get(key, default)
        if not isinstance(val, ty):
            msg = f"Invalid config type of {key}, got {type(val)} expected " \
                  f"(subtype of) {ty}"
            print("Error:", msg)
            raise TypeError(msg)
        return val  # type: ignore


def warn(msg: str) -> None:
    print("W: " + msg, file=sys.stderr)


def write_type(dst: TextIO, t: str, i: int, prefix: str = "") -> None:
    dst.write(f"{prefix}t{i} :: {t}\n{prefix}t{i} = undefined\n")


def val2bool(val: str) -> bool:
    return val.lower() == "yes" or val.lower() == "true"


def mkTeacherFile(srcFile: FilePath, dstFile: FilePath) -> TestConf:
    with open(dstFile, "w") as dst:
        with open(srcFile, "r") as src:
            assert _config is not None
            conf = TestConf(_config)

            dst.write("{-# LANGUAGE " + ", ".join(conf.teacher_exts())
                                      + " #-}\n\n")
            dst.write("module Teacher where\n\n")

            inject: bool = False
            inject_end: bool = False

            for ln in src:
                match = OPT_R.fullmatch(ln.rstrip())
                if match:
                    key: str = match[1]

                    if key == "INJECT BEGIN":
                        assert not inject, "inject inside inject"
                        inject = True
                    elif key == "INJECT END":
                        assert inject, "inject end without start"
                        inject = False
                        inject_end = True
                    elif key.startswith("typecmp"):
                        conf.typecmp = True
                        if key == "typecmpex":
                            conf.typecmpex = True
                            conf.typecmp_num_types = 1
                        if key == "typecmp num types":
                            conf.typecmp_num_types = int(match[2])
                    elif key == "allow hint":
                        conf.allow_hint = match[2].strip().lower() == "true"
                    elif key == "hlint":
                        conf.hlint = match[2].strip().lower() == "true"

                if conf.reflection():
                    match = AST_CHECK_DECL.search(ln)
                    if match:
                        conf.ast_checks.append(match[1])

                if inject:
                    conf.inject += ln
                if inject_end:
                    conf.inject += ln
                    inject_end = False

                if not conf.typecmp or conf.typecmpex:
                    dst.write(ln)
                else:
                    if not IGNORE.match(ln):
                        write_type(dst, ln, conf.typecmp_num_types)
                        conf.typecmp_num_types += 1

            return conf


def write_student(src: TextIO, dst: TextIO, conf: TestConf, prefix: str = "") \
        -> None:
    type_idx = 0
    if not conf.typecmp:
        for ln in src:
            dst.write(prefix + ln)
    else:
        for ln in src:
            if not IGNORE.match(ln):
                write_type(dst, ln, type_idx, prefix)
                type_idx += 1


def mkStudentFile(srcFile: FilePath, dstFile: FilePath, conf: TestConf) \
        -> None:
    with open(dstFile, "w") as dst:
        # one skipping BOM if present
        with open(srcFile, "r", encoding='utf-8-sig') as src:
            exts_string = ", ".join(conf.student_exts())
            dst.write("{-# LANGUAGE " + exts_string + " #-}\n\n")
            if conf.require_type():
                dst.write("{-# OPTIONS_GHC -Werror=missing-signatures #-}\n\n")
            if conf.typecmp:
                dst.write("{-# OPTIONS_GHC -w #-}")
            elif conf.warnings():
                ws = " ".join(f"-W{w}" for w in conf.warnings())
                dst.write("{-# OPTIONS_GHC " + ws + " #-}")
            dst.write("module Student where\n\n")

            if conf.inject:
                dst.write(conf.inject)

            if not conf.allow_imports():
                dst.write("noImportsForStudentAllowed :: ()\n"
                          "noImportsForStudentAllowed = ()\n\n")

            dst.write('{-# LINE 1 "IS.hs" #-}\n')
            write_student(src, dst, conf)


def mkStudentReflectorFile(srcFile: FilePath, dstDir: FilePath,
                           conf: TestConf) -> None:
    with open(os.path.join(dstDir, "StudentReflector.hs"), "w") as dst:
        with open(srcFile, "r", encoding='utf-8-sig') as src:
            dst.write(textwrap.dedent("""\
                {-# LANGUAGE TemplateHaskell #-}

                module StudentReflector (
                    studentAST, runASTCheck,
                    ) where

                import Test.Expr.Reflection ( Q, Dec, ASTCheck )
                import Control.Applicative ( pure )
                """))
            dst.write(format(conf.inject or ""))
            dst.write(textwrap.dedent("""\

                studentAST :: Q [Dec]
                studentAST = [d|
                """))

            write_student(src, dst, conf, "    ")
            dst.write("\n    |]\n\n")
            dst.write(textwrap.dedent("""\
                runASTCheck :: ASTCheck -> Q [Dec]
                runASTCheck check = do
                    ast <- studentAST
                    check ast
                    pure []
                """))
    with open(os.path.join(dstDir, "StudentASTChecks.hs"), "w") as dst:
        dst.write(textwrap.dedent("""\
            {-# LANGUAGE TemplateHaskell #-}

            module StudentASTChecks () where

            import StudentReflector ( runASTCheck )
            import qualified Teacher

            """))

        for c in conf.ast_checks:
            dst.write(f"$(runASTCheck Teacher.{c})\n")


def mkMainFile(filename: FilePath, conf: TestConf) -> None:
    reflect = ""
    if conf.reflection():
        reflect = "import qualified StudentASTChecks"
    with open(filename, "w") as dst:
        dst.write("{-# LANGUAGE " + ", ".join(conf.main_exts()) + "#-}\n")
        dst.write(f"""
module Main ( main ) where

import qualified Teacher
import safe qualified Student
{reflect}
import Test.Expr
import Test.Expr.Config
import Test.Expr.Types ( TypeOrder ( TEqual ) )
import Data.Maybe ( catMaybes )
import Data.Tuple.TH.Curry
import Language.Haskell.TH ( Exp ( ConE ) )
import Test.QuickCheck.Convertible ( convert )
import qualified Test.QuickCheck.Modifiers
import qualified Test.QuickCheck.Range
import qualified Test.QuickCheck.Literal
import qualified Test.QuickCheck.Union
import System.Exit
import Control.Monad
""")
        for imp in conf.imports():
            dst.write("import qualified {}\n".format(imp))

        if not conf.typecmp:
            dst.write(f"""
$(testMain (TestConfig
      ( ConfigEntry TypeOrd $(extractOptionDef "typeOrder" (ConE 'TEqual))
      : ConfigEntry JournalFd {journal_fd}
      : catMaybes [ConfigEntry TestPattern <$> $(extractOptionMaybe "pattern"),
                   ConfigEntry Expression <$> $(extractOptionMaybe "expr"),
                   ConfigEntry DegenType <$> $(extractOptionMaybe "degenType")
                  ])))
""")
        else:
            dst.write("main :: IO ()\nmain = do\n")
            for i in range(conf.typecmp_num_types):
                dst.write(f'''\
    $(testType $(extractOptionDef "typeOrder" (ConE 'TEqual)) "t{i}")
''')


with tempfile.TemporaryDirectory(prefix="hsExprTest.") as wd:
    conf = mkTeacherFile(teacher, os.path.join(wd, "Teacher.hs"))

    if hint and not conf.allow_hint:
        print("syntax check not allowed for this question")
        sys.exit(1)

    mkStudentFile(student, os.path.join(wd, "Student.hs"), conf)
    if conf.reflection():
        mkStudentReflectorFile(student, wd, conf)
    os.chdir(wd)

    mkMainFile("Main.hs", conf)

    include.append(os.path.dirname(conf.question_path()))
    invocation: List[str] = [f"-i{x}" for x in include] + ["-dynamic"]
    execution: List[str] = []

    if conf.typecmp:
        # just compile, does not need to run
        invocation.insert(0, "ghc")
        invocation.extend(["-e", "return ()"])
        if hint:
            invocation.append("Student.hs")  # just parse student's file
            if conf.reflection():
                invocation.append("StudentASTChecks.hs")
        else:
            invocation.append("Main.hs")  # acutally do the comparison
    else:  # expression comparison
        if hint:
            invocation.insert(0, "ghc")
            invocation.extend(["-e", "return ()"])
        else:
            if conf.compiled():
                invocation[:0] = ["ghc", "-o", "test", "-odir", ".",
                                  "-hidir", ".", "-O2"]
                execution = ["./test"]
            else:
                invocation.insert(0, "runghc")

        if hint and not conf.get_typed("hint teacher", True, bool):
            invocation.append("Student.hs")  # just parse student's file
            if conf.reflection():
                invocation.append("StudentASTChecks.hs")
        else:
            invocation.append("Main.hs")
    pass_fds = [journal_fd] if journal_fd != 2 else []
    r = subprocess.run(invocation, stdin=subprocess.DEVNULL,
                       stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                       pass_fds=pass_fds, encoding="utf-8")

    if conf.compiled():
        for ln in r.stdout.splitlines():
            if COMPILE_LINE.search(ln) is None:
                print(ln, flush=True)
            else:
                print(ln, flush=True, file=sys.stderr)
    else:
        print(r.stdout, end='', flush=True)
    hide = False
    ast_check = False
    type_missing = False
    hidden = 0
    for ln in r.stderr.splitlines():
        m = ERR_HEADER.match(ln)
        if m:
            hide = False
            ast_check = False
            if hint and (m[1] != "IS" and m[1] != "StudentASTChecks"):
                hide = True
                hidden += 1
            elif m[1] == "StudentASTChecks":
                ast_check = True
        if hide:
            print(ln, file=sys.stderr, flush=True)
        elif ast_check and \
                (GHC_LINE_ERR.search(ln) is not None or
                 AST_CHECK_EXC.search(ln) is not None or
                 AST_CHECK_CALL_STACK.search(ln) is not None):
            print(ln, file=sys.stderr, flush=True)
        # don't show students the type inferred by GHC if they did not write it
        elif hint and type_missing:
            type_missing = False
            if (m := HINT_TYPE_MISSING_TYPE.match(ln)) is not None:
                print(m["fun"], "(type missing)", flush=True)
                print(ln, file=sys.stderr, flush=True)
        else:
            if hint and HINT_TYPE_MISSING_LEAD.match(ln) is not None:
                type_missing = True
            print(ln, flush=True)
    if hidden > 0:
        print(f"+ {hidden} more errors hidden in syntax check mode, check "
              "that your function is correctly typed and named", flush=True)

    is_hint = "hint" if hint else "no hint"
    if not hint and len(execution) != 0:
        if r.returncode != 0:
            print(f"failed (compilation): {r.returncode} ({is_hint})",
                  file=sys.stderr, flush=True)
            sys.exit(2)
        r = subprocess.run(execution, stdin=subprocess.DEVNULL,
                           stdout=sys.stdout, stderr=sys.stdout,
                           pass_fds=pass_fds, encoding="utf-8")

    lint = shutil.which("fi-hlint")
    if not conf.typecmp and conf.hlint and lint is not None:
        lp = subprocess.run([lint, "Student.hs"],
                            stdin=subprocess.DEVNULL,
                            stdout=subprocess.PIPE,
                            stderr=subprocess.STDOUT,
                            encoding="utf-8")
        if lp.stdout:
            print("\nSuggestions for code quality improvements:")
            print(lp.stdout)

    if r.returncode != 0:
        if r.returncode == -14:
            print("\nTIMEOUT -- total time limit exceeded", flush=True)
        print(f"failed: {r.returncode} ({is_hint})",
              file=sys.stderr, flush=True)
        sys.exit(1)
    sys.exit(0)

# vim: colorcolumn=80 expandtab sw=4 ts=4
