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):

    def __init__ (self, input_raw):

        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):
        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 compiler stages 1-4.

    Returns:
        ctx (CompilationContext): The compiled 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."""
    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.
    """
    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."""
    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."""
    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)