#
#   Imandra Inc.
#
#   codelogician/modeling/renderer.py
#

import json

from rich import box
from rich.console import Group
from rich.panel import Panel
from rich.pretty import Pretty
from rich.syntax import Syntax
from rich.table import Table
from rich.text import Text

from ..strategy.cl_agent_state import frm_status_to_rich
from ..util import intersperse, maybe_else
from .dep_context import ModelDepContext
from .metamodel import MetaModel
from .model import Model
from .task_creator import ModelTaskCreator

# ------------------------------------------------------------------------------------------
#   Model renderers
# ------------------------------------------------------------------------------------------


def model_to_json(m: Model):
    # This will save the reasons which we can only calculate if
    # we have "live" dependency models (not strings as is the case
    # post-deserialization)

    # Removed this from Model class
    # @field_serializer('static_frm_reasons')
    # def serialize_frm_reasons(self, _):
    #    return self.formalization_reasons()

    task_creator = ModelTaskCreator()
    m.static_frm_reasons = task_creator.formalization_reasons(m, 5)

    return m.model_dump_json()


def model_to_str_summary(m: Model) -> str:
    return str(ModelTaskCreator().gen_model_status(m))


def model_to_str(m: Model) -> str:
    agentStateStr = 'None' if m.agent_state is None else m.agent_state.status

    s = f'Model: {m.rel_path} \n'
    # s += f'{str(ModelTaskCreator().gen_model_status.status())}\n'
    s += f'Formalization state: {agentStateStr}\n'
    s += f'Depends on: {ModelDepContext.deps_paths(m)}\n'
    s += f'Source language: {m.src_language}\n'

    if m.src_code:
        s += f'Source code (condensed): \n {m.src_code[:100]}\n'
    else:
        s += 'Source code (condensed): \n None \n'

    s += f'Opaque funcs: {len(m.opaque_funcs())}\n'
    s += f'Decomps: {len(m.decomps())}\n'
    return s


def iml_code_to_print(m: Model) -> str:
    """
    Provide IML code (if available) along with relevant import statements
    """
    if m.iml_code() is None:
        return 'N\\A'

    import_statements = ModelDepContext().create_import_iml_statements(m)
    return f'{import_statements}\n\n{m.iml_code()}'


def render_model(model: Model):
    table = Table(
        title=Text(f'Model: {model.rel_path}', style='bold italic magenta'),
        box=box.MINIMAL_DOUBLE_HEAD,
        show_lines=True,
    )

    table.add_column('Name')
    table.add_column('Attribute')

    src_code_str = 'N/A' if model.src_code is None else model.src_code
    user_iml_code_str = (
        'N/A' if model.user_iml_edit is None else model.user_iml_edit.user_iml_entry
    )
    user_iml_code_last_change = (
        'N/A' if model.user_iml_edit is None else str(model.user_iml_edit.time_of_edit)
    )

    needs_frm_str = 'Yes' if model.static_frm_reasons == [] else 'No'
    frozen_str = 'Yes' if model.iml_code_frozen else 'No'

    iml_code_str = 'N/A' if not model.iml_code() else str(model.iml_code())

    eval_result_str = 'N/A' if not model.agent_state else model.agent_state.eval_res

    # fmt: off
    table.add_row('Relative path',              Text(model.rel_path, style='bold')                      )
    table.add_row('Formalization status',       frm_status_to_rich(model.formalization_status())        )
    table.add_row('Eval result',                Text(eval_result_str, style='bold')                    )
    table.add_row('Needs formalization?',       Text(needs_frm_str, style='bold')                      )
    table.add_row('Formalization reasons',      Text(str(model.static_frm_reasons))                     )
    table.add_row('Frozen?',                    Text(frozen_str)                                       )
    table.add_row('Src last changed',           Text(f'{model.src_code_last_changed}', style='bold')    )
    table.add_row('Src code',                   Syntax(src_code_str, 'python')                         )
    table.add_row('IML code',                   Syntax(iml_code_str, 'ocaml')                          )
    table.add_row('User IML code',              Syntax(user_iml_code_str, 'ocaml')                     )
    table.add_row('User IML code last changed', Text(f'{user_iml_code_last_change}', style='bold')     )
    table.add_row('Context',                    Text(f'{model.context}', style='bold')                  )
    table.add_row('Task ID',                    Text(f'{model.outstanding_task_ID}', style='bold')      )
    table.add_row('Deps',                       Text(f'{model.dependencies}', style='bold')             )
    table.add_row('Opaques',                    Text(f'{model.opaque_funcs()}', style='bold')           )
    table.add_row('Verification Goals',         Text(f'{model.verification_goals()}', style='bold')     )
    table.add_row('Decompositions',             Text(f'{model.decomps()}', style='bold')                )
    # fmt: on

    return table


def render_model_list(models: list[Model]):
    table = Table(
        title=Text('Models list', style='bold italic magenta'),
        box=box.MINIMAL_DOUBLE_HEAD,
        show_lines=True,
    )

    # fmt: off
    table.add_column('ID'                       , justify='right', no_wrap=True )
    table.add_column('Frm Status'               , justify='right', no_wrap=True )
    table.add_column('Frozen'                   , justify='right', no_wrap=True )
    table.add_column('Needs frm?'               , justify='right', no_wrap=True )
    table.add_column('Frm Needs'                , justify='right', no_wrap=False)
    table.add_column('Rel path'                 , justify='right', no_wrap=True )
    table.add_column('Src Last Changed'         , justify='right', no_wrap=True )
    table.add_column('User IML Code'            , justify='right', no_wrap=True )
    table.add_column('User IML Code Last Chged' , justify='right', no_wrap=True )
    table.add_column('Src Code Embds'           , justify='right', no_wrap=True )
    table.add_column('IML Code Embds'           , justify='right', no_wrap=True )
    table.add_column('Context'                  , justify='right' )
    table.add_column('Task ID'                  , justify='right' )
    table.add_column('Dependencies'             , justify='right', no_wrap=False)
    table.add_column('Num opaques'              , justify='right' )
    table.add_column('Num VGs'                  , justify='right' )
    table.add_column('Num Decomps'              , justify='right' )

    def yes_no(b): return 'Yes' if b else 'No'
    def edit_time(e): return str(e.time_of_edit)
    # fmt: on

    for m_id, m in enumerate(models):
        table.add_row(
            Text(f'{m_id}', style='bold'),
            frm_status_to_rich(m.formalization_status()),
            Text(f'{yes_no(m.iml_code_frozen)}', style='bold'),
            Text(f'{yes_no(m.static_frm_reasons)}', style='bold'),
            Text(f'{map(str, m.static_frm_reasons)}', style='bold'),
            Text(f'{m.rel_path}', style='bold'),
            Text(f'{m.src_code_last_changed}', style='bold'),
            Text(f'{m.user_iml_edit is not None}', style='bold'),
            Text(f'{maybe_else("N/A", edit_time, m.user_iml_edit)}', style='bold'),
            Text(f'{m.context}', style='bold'),
            Text(f'{m.outstanding_task_ID}', style='bold'),
            Text(f'{m.dependencies}', style='bold'),
            Text(f'{len(m.opaque_funcs())}', style='bold'),
            Text(f'{len(m.verification_goals())}', style='bold'),
            Text(f'{len(m.decomps())}', style='bold'),
        )

    return table


# ------------------------------------------------------------------------------------------
#   MetaModel renderers
# ------------------------------------------------------------------------------------------


def metamodel_rich_summary(mmodel: MetaModel):
    """Return rich summary of the metamodel state"""
    summary = mmodel.summary()

    def frm_status(status_count):
        status, count = status_count
        return Text.assemble(status.__rich__(), ' ', (str(count), 'bold'))

    def fmt_countable(label_key):
        label, key = label_key
        return Text.assemble(f'Num. {label}: ', (str(summary[key]), 'bold'))

    countables = [
        ('Models', 'num_models'),
        ('Verification goals', 'num_vgs'),
        ('Opaque functions', 'num_opaque_funcs'),
        ('Decompositions', 'num_decomps'),
    ]
    return Group(
        Text.assemble(*intersperse(' | ', map(fmt_countable, countables))),
        Text.assemble(
            *intersperse(' | ', map(frm_status, summary['frm_statuses'].items()))
        ),
    )


def mmodel_to_str(mmodel: MetaModel):
    """Return a nice representation"""

    metaSummary = mmodel.summary()
    statuses = {s.name: n for s, n in metaSummary['frm_statuses'].items()}

    s = f'Source file directory: {mmodel.src_dir_abs_path}\n'
    s += f'Number of models: {metaSummary["num_models"]}\n'
    s += f'Number of VGs: {metaSummary["num_vgs"]}\n'
    s += f'Number of decomps: {metaSummary["num_decomps"]}\n'
    s += f'Status breakdown: \n{json.dumps(statuses, indent=4)}'
    s += '\n'
    s += 'Model details:\n'

    for path in mmodel.models:
        s += f'Model [path={path}; {model_to_str_summary(mmodel.models[path])}]\n'

    return s


def __rich__(self):
    """Return RICH values"""
    pretty = Pretty(self.toJSON(), indent_size=2)
    return Panel(pretty, title='Current MetaModel', border_style='green')
