import os
import sys
import argparse
import datetime
import subprocess as sp
from post import (
    CompilationContext,
    ParserStage,
    DerivativeExpansionStage,
    DependencyResolutionStage,
    SympyOptimizationStage,
    FortranProgramWriter,
    ReportWriter,
    LatexWriter
)

parser = argparse.ArgumentParser()
parser.add_argument("term_file", help="name of file containing postprocessing term spec")
parser.add_argument("-b", "--build-info", action="store_true", help="print build_info instead of code")
parser.add_argument("-x", "--print-latex", action="store_true", help="print latex equation instead of code")
parser.add_argument("-p", "--print-report", action="store_true", help="print code generation report instead of code")
parser.add_argument("-m", "--make-build-module", action="store_true", help="print build_info Fortran module instead of code")
args = parser.parse_args()


class VersionInfo(object):
    """Encapsulates build and version metadata for the code generator execution.

    This class queries the current Git repository state to fetch the HEAD commit
    hash and checks if there are uncommitted changes, appending metadata to
    identify the exact codebase state used to generate the output files.
    """

    def __init__(self, input_raw):
        """Initializes VersionInfo with build date, git hash, and terms metadata.

        Args:
            input_raw (str): The raw input string containing DSL term definitions.
        """
        bd = 'build date: {}'
        bb = 'build base: {}'

        self.build_date = bd.format(datetime.datetime.today().ctime())

        try:
            self.build_base = bb.format(sp.check_output(["git", "rev-parse", "HEAD"]).strip())
        except sp.CalledProcessError:
            self.build_base = "None"

        try:
            sp.check_output("git diff --exit-code".split())
            sp.check_output("git diff --cached --exit-code".split())
            self.off_base = ""
        except sp.CalledProcessError:
            self.off_base = "built with changes not commited"

        tinfo = '''
================================================================================
************************************ terms *************************************
--------------------------------------------------------------------------------
{}
================================================================================
'''

        self.term_info = tinfo.format(input_raw)

    def fparams(self):
        """Formats the collected build parameters into a single string.

        Returns:
            str: Multi-line string summarizing build date, revision, and terms details.
        """
        return "\n".join([self.build_date, self.build_base, self.off_base, self.term_info])


def compile_terms(terms_raw):
    """Parses DSL terms spec and runs it through the four compiler stages.

    This function initializes a new `CompilationContext` and executes the compiler pipeline
    sequentially: parsing, derivative/fluctuation expansion, data dependency resolution,
    and SymPy optimization (including buffer pooling).

    Args:
        terms_raw (str): The raw string contents of the DSL terms specification input file.

    Returns:
        CompilationContext: The fully compiled and optimized compilation context.
    """
    ctx = CompilationContext()
    ParserStage().execute(terms_raw, ctx)
    DerivativeExpansionStage().execute(ctx)
    DependencyResolutionStage().execute(ctx)
    SympyOptimizationStage().execute(ctx)
    return ctx


def get_latex_equations_str(ctx):
    """Generates the LaTeX equations dictionary string from the context.

    Args:
        ctx (CompilationContext): The compiled context containing the averaged fields.

    Returns:
        str: A formatted string representing a Python dictionary mapping field names to LaTeX code.
    """
    latex_lines = ["{"]
    for avg in ctx.averaged.values():
        latex_lines.append('  "{}" : r"${}$",'.format(avg.name, avg.latex))
    latex_lines.append("}")
    return "\n".join(latex_lines)


def test(terms_raw, report=False, latex=False):
    """Compiles the post-processing term specifications from DSL into targets.

    This compiles the Lark AST through all 4 pipeline stages
    and outputs the resulting Fortran source code via FortranProgramWriter,
    a LaTeX equation dictionary via LatexWriter, or an IR report via ReportWriter.

    Args:
        terms_raw (str): The raw string contents of the DSL terms specification.
        report (bool, optional): If True, prints a JSON-based compilation IR report instead. Defaults to False.
        latex (bool, optional): If True, prints LaTeX equation dictionary instead. Defaults to False.
    """
    ctx = compile_terms(terms_raw)

    if report:
        ReportWriter().write(ctx)
    elif latex:
        LatexWriter().write(ctx)
    else:
        FortranProgramWriter().write(ctx)


def escape_fortran_string(s):
    """Escapes string content for standard Fortran source formatting.

    Converts internal double quotes into double-double quotes and splits the string lines,
    wrapping them in Fortran concatenation syntax `// char(10) // &` to respect line length limit.

    Args:
        s (str): The raw string to escape.

    Returns:
        str: The escaped and formatted Fortran string parameter literal.
    """
    escaped = s.replace('"', '""')
    lines = escaped.splitlines()
    if not lines:
        return '""'
    formatted_lines = [f'"{line}"' for line in lines]
    return ' // char(10) // &\n    '.join(formatted_lines)


def generate_build_info_module(terms_raw):
    """Generates a complete standard Fortran module containing build info and LaTeX equations.

    The generated module contains metadata about the compiler run, including the build base Git hash,
    and a string representations of the LaTeX equations.

    Args:
        terms_raw (str): The raw string contents of the DSL terms specification input file.
    """
    vinfo = VersionInfo(terms_raw)
    build_info_str = vinfo.fparams()

    ctx = compile_terms(terms_raw)
    latex_str = get_latex_equations_str(ctx)

    fortran_build_info = escape_fortran_string(build_info_str)
    fortran_latex = escape_fortran_string(latex_str)

    fortran_code = f"""module m_build_info
  implicit none
  character(len=*), parameter :: build_info_str = &
    {fortran_build_info}
  character(len=*), parameter :: latex_equations_str = &
    {fortran_latex}
end module m_build_info
"""
    print(fortran_code)


def build_info(terms_raw):
    """Generates and prints build information for the current code generation run.

    Args:
        terms_raw (str): The raw string contents of the DSL terms specification input file.
    """
    vinfo = VersionInfo(terms_raw)
    print(vinfo.fparams())



if __name__ == '__main__':

    with open(args.term_file) as inputfile:
        terms_raw = ((inputfile.read()))

    if args.make_build_module:
        generate_build_info_module(terms_raw)
    elif args.build_info:
        build_info(terms_raw)
    else:
        test(terms_raw, report=args.print_report, latex=args.print_latex)