#
#   Imandra Inc.
#
#   model.py
#

# ruff: noqa: UP031
import sys
from functools import partial
from typing import cast

from textual import on
from textual.containers import (
    Horizontal,
    HorizontalGroup,
    ScrollableContainer,
    VerticalGroup,
    VerticalScroll,
)
from textual.reactive import reactive
from textual.screen import Screen
from textual.widgets import (
    Button,
    Collapsible,
    Footer,
    Label,
    Static,
    TabbedContent,
    TabPane,
)

from codelogician.modeling.metamodel import MetaModel
from codelogician.modeling.model import Model
from codelogician.modeling.task_creator import ModelTaskCreator
from codelogician.strategy.cl_agent_state import CLAgentState
from codelogician.util import filter, fst, maybe, maybe_else, not_none, translate

from ..common import Border, MyHeader, bind, decomp_ui, opaques_rich, vg_ui
from ..tree_views import TreeViews
from .model_cmds import ModelCommandsView


# fmt: off
def _opaques_view(opaques): return [Static(opaques_rich(opaques))]
def _bordered_views(prefix, f, xs):
    return [Border(f'{prefix} #{i+1}', f(x)) for i, x in enumerate(xs)]

_decomps_view = partial(_bordered_views, 'Decomp', decomp_ui)
_vgs_view     = partial(_bordered_views, 'VG', vg_ui)

def rel_path(model): return model.rel_path
def src_code(model): return model.src_code

def fmt_count(n, sing, plu): return f'{n} {sing if n == 1 else plu}'

task_creator = ModelTaskCreator()

PathStatePairs = list[tuple[str, None | CLAgentState]]
# fmt: on


class Container(VerticalGroup):
    items = reactive([], recompose=True)

    def __init__(self, generator):
        super().__init__()
        self.generator = generator

    def compose(self):
        yield from self.generator(self.items)


class ModelStatus(VerticalGroup):
    model: reactive[None | Model] = reactive(None, recompose=True)

    def compose(self):
        if self.model is not None:
            yield Static(task_creator.gen_model_status(self.model), id='status')


class ModelName(VerticalGroup):
    model: reactive[None | Model] = reactive(None, recompose=True)

    def compose(self):
        label = maybe_else('<No model>', rel_path, self.model)
        yield Static(f'[$accent][b]{label}[/b][/]')


class CodeStatic(Static):
    DEFAULT_CSS = 'CodeStatic { width: auto; min-width: 100% }'

    code: reactive[str] = reactive('')

    def render_code(self, code):
        from rich.syntax import Syntax

        theme = 'ansi_dark' if 'pytest' in sys.modules else 'monokai'
        opts = {'line_numbers': True, 'indent_guides': True, 'theme': theme}
        return Syntax(code, self.lang, **opts)

    def __init__(self, lang):
        self.lang = lang
        super().__init__(self.render_code(''), id=lang, classes='code')

    def watch_code(self, code):
        self.update(self.render_code(code))


class ReactiveCollapsible(Collapsible):
    content = reactive('')

    def __init__(self, mk_title, generator, collapsed=True):
        super().__init__(collapsed=collapsed)
        self.generator = generator
        self.mk_title = mk_title

    def watch_content(self, content):
        self.styles.display = 'block' if content else 'none'
        self.title = self.mk_title(content)

    def compose(self):
        # this is a bit cheaty
        yield self._title
        with self.Contents():
            yield from self.generator(ReactiveCollapsible.content)


class CodeView(ReactiveCollapsible):
    DEFAULT_CSS = 'ScrollableContainer { height: auto; max-height: 16 }'

    def __init__(self, title, lang):
        def mk_title(code):
            return f'{title} ({fmt_count(len(code.split("\n")), "line", "lines")})'

        def generator(upstream):
            with ScrollableContainer():
                yield CodeStatic(lang).data_bind(code=upstream)

        super().__init__(mk_title, generator, collapsed=False)


class ListView(ReactiveCollapsible):
    def __init__(self, kind, inner_generator):
        def mk_title(items):
            return f'{kind} ({fmt_count(len(items), "item", "items")})'

        def generator(upstream):
            yield Container(inner_generator).data_bind(items=upstream)

        super().__init__(mk_title, generator, collapsed=True)


def _errors_view(errors):
    from regex import compile

    parse_kind = compile(r'\s*{\s*Kind.name\s*=\s*"(.*)"\s*}\s*').match
    kinds = {
        'SyntaxErr': 'Syntax error',
        'TypeErr': 'Type error',
        'TacticEvalErr': 'Tactic evaluation error',
    }

    def fmt_kind(kind):
        return translate(
            kinds, maybe_else(kind, lambda m: m.group(1), parse_kind(kind))
        )

    def error_ui(e, contexts):
        with VerticalGroup():
            label, msg = (
                Label(f'[bold]{fmt_kind(e.kind)}[/bold]: '),
                Static(f'[red]{e.msg.msg}[/]'),
            )
            yield from (
                [label, msg] if '\n' in e.msg.msg else [HorizontalGroup(label, msg)]
            )

            for loc, c in zip(e.msg.locs, contexts):
                start_l, start_c = c.start_line, c.start_col
                stop_l, stop_c = c.stop_line, c.stop_col
                file = c.file.replace('.py', '.iml')  # TODO: prefer not to hack this
                yield Static(f'in {file} {start_l}:{start_c} -- {stop_l}:{stop_c}')
                yield Static('\n'.join(c.lines))

    for i, (e, contexts) in enumerate(errors):
        with Border('#%d' % (i + 1)):
            yield from error_ui(e, contexts)


class ModelStateView(VerticalGroup):
    model: reactive[None | Model] = reactive(None)
    iml_code: reactive[str] = reactive('')
    opaque_funcs = reactive([])
    region_decomps = reactive([])
    errors = reactive([])
    vgs = reactive([])

    async def watch_model(self, old, new):
        def class_(m):
            return 0 if m is None else 1 if m.agent_state is None else 2

        if class_(new) == 2:
            fstate = new.agent_state
            self.iml_code = fstate.iml_code
            self.opaque_funcs = fstate.opaque_funcs
            self.region_decomps = fstate.region_decomps
            self.vgs = fstate.vgs
            self.errors = list(zip(fstate.errors, fstate.error_contexts))

        if class_(old) != class_(new):
            await self.recompose()

    def compose(self):
        def if_nonempty(f, xs):
            if len(xs) > 0:
                yield from f(xs)

        fstate = self.model.agent_state if self.model else None
        if self.model is None:
            pass
        elif fstate is None:
            yield Static('Not formalised')
        else:
            # fmt: off
            yield CodeView('IML code', 'ocaml').data_bind(content=ModelStateView.iml_code)
            yield ListView('Opaque functions', _opaques_view).data_bind(content=ModelStateView.opaque_funcs) 
            yield ListView('Verification goals', _vgs_view).data_bind(content=ModelStateView.vgs)
            yield ListView('Region decompositions', _decomps_view).data_bind(content=ModelStateView.region_decomps)
            yield ListView('Formalization errors', _errors_view).data_bind(content=ModelStateView.errors)
            # fmt: in


class ModelView(VerticalGroup):
    DEFAULT_CSS = 'ModelName { margin-bottom: 1 }'

    model: reactive[None | Model] = reactive(None)
    src_code: reactive[str] = reactive('')

    def watch_model(self, model: Model):
        self.src_code = '' if model is None else (model.src_code or '')

    def compose(self):
        yield ModelName().data_bind(ModelView.model)
        with VerticalScroll():
            src_code = ModelView.src_code
            yield ModelStatus().data_bind(ModelView.model)
            yield CodeView('Source code', 'python').data_bind(content=src_code)
            yield ModelStateView().data_bind(ModelView.model)


class SubtreeViewer(VerticalGroup):
    subtree: reactive[tuple[str, PathStatePairs]] = reactive(('', []), recompose=True)

    def __init__(self, pad: int, mk_view):
        super().__init__()
        self.mk_view = mk_view
        self.pad = pad

    def compose(self):
        descriptor, path_state_pairs = self.subtree
        yield Label(descriptor, classes='subtree-descriptor')
        with VerticalScroll():
            for p, state in path_state_pairs:
                yield Label('[$primary][b]%s[/b][/]' % p)
                with VerticalGroup() as v:
                    v.styles.padding = (0, 0, 0, self.pad)
                    if state is None:
                        yield Static(' [$foreground-muted]Not formalised[/]')
                    else:
                        yield from self.mk_view(state)


class ModelScreen(Screen):
    DEFAULT_CSS = """
        #panes { width: 65% }
        #panes TabPane { padding: 0 1 0 1; }
        Label.subtree-descriptor { color: $accent; margin-bottom: 1; width: 1fr }
        """

    mmodel: reactive[None | MetaModel] = reactive(None)
    model: reactive[None | Model] = reactive(None)
    selected: reactive[None | str] = reactive(None)
    subtree: reactive[tuple[str, PathStatePairs]] = reactive(('<None>', []))

    def on_mount(self):
        self.title = 'Model View'
        bind(self, 'mmodel')

    def empty_container(self, id):
        container = self.query_one(id)
        container.remove_children()
        return container

    def watch_mmodel(self, old_mm: MetaModel, new_mm: MetaModel):
        if self.selected:
            self.select_path(self.selected)

    def compose(self):
        def subtree(pad, mk_view):
            return SubtreeViewer(pad, mk_view).data_bind(ModelScreen.subtree)

        yield MyHeader()
        with Horizontal():
            yield TreeViews().data_bind(ModelScreen.mmodel)
            with TabbedContent(id='panes'):
                with TabPane('Model state', id='model_tab'):
                    yield ModelView().data_bind(ModelScreen.model)
                with TabPane('Command entry', id='commands_tab'):
                    yield ModelCommandsView().data_bind(ModelScreen.model)
                with TabPane('Opaque functions', id='opaques_tab'):
                    yield subtree(1, lambda st: _opaques_view(st.opaque_funcs))
                with TabPane('Decomposition requests', id='decomps_tab'):
                    yield subtree(2, lambda st: _decomps_view(st.region_decomps))
                with TabPane('Verification goals', id='vgs_tab'):
                    yield subtree(2, lambda st: _vgs_view(st.vgs))
        yield Footer()

    # fmt: off
    def paths_under(self, rel_path):
        tree = self.query_one(TreeViews).active_tree()
        file_path = None if rel_path == '' else tree.path_of_node(rel_path)
        children = tree.get_descendants(rel_path)
        paths = filter(not_none, [file_path, *map(tree.path_of_node, children)])
        descriptor = ('All modules' if rel_path == '' else
                      { 'src-tree':   (f'Just the module [b]{rel_path}[/b]' if file_path else
                                       f'Modules in directory [b]{rel_path}[/b]:')
                      , 'rev-deps':    f'Modules [b]{rel_path}[/b] and its dependents:'
                      , 'module-deps': f'Module [b]{rel_path}[/b] and its dependencies:'
                      }[cast(str, tree.id)])
        return descriptor, paths
    # fmt: on

    def watch_selected(self, _, selected):
        maybe(self.select_path, selected)

    def select_path(self, rel_path):
        descriptor, paths = self.paths_under(rel_path)

        def mk_subtree(mmodel):
            def state_of_path(p):
                return cast(MetaModel, mmodel).models[p].agent_state

            return descriptor, [(p, state_of_path(p)) for p in paths]

        self.subtree = maybe_else(('<Nothing>', []), mk_subtree, self.mmodel)
        self.model = cast(MetaModel, self.mmodel).models.get(rel_path)
        # print(f'---| got model: {maybe(lambda m: m.rel_path, self.model)}')

    def on_tree_node_selected(self, event):
        self.selected = maybe_else('', fst, event.node.data)

    @on(Button.Pressed, '#send_commands')
    def add_model_commands(self, event):
        """Add commands to the model"""
        pass

    def show_model(self, rel_path):
        self.selected = rel_path
        self.query_one(TreeViews).select_path(rel_path)
