#
#   Imandra Inc.
#
#   help.py
#

from collections import defaultdict
from functools import partial
from pathlib import Path

from rich.style import Style
from rich.text import Text
from textual.containers import Horizontal, VerticalScroll
from textual.screen import Screen
from textual.widgets import Footer, Markdown, Tree

from codelogician.doc.md_utils import parse_frontmatter
from codelogician.doc.search import SKILL_DIR, collect_iml_files, collect_md_files
from codelogician.util import maybe

from ..common import MyHeader

TOGGLE_STYLE = Style.from_meta({'toggle': True})


class TOC(Tree):
    def __init__(self, md_files, skill_dir):
        Tree.__init__(self, 'CodeLogician Help')
        self.root.expand()

        # Group files by their parent directory relative to skill_dir
        groups = defaultdict(list)
        top_level = []
        for md_path, fm in md_files:
            rel = md_path.relative_to(skill_dir)
            if rel.parent == Path('.'):
                top_level.append((md_path, fm))
            else:
                groups[rel.parts[0]].append((md_path, fm))

        def label(md_path, fm):
            return fm.get('name', md_path.stem) if fm else md_path.stem

        # Add top-level files
        for md_path, fm in top_level:
            self.root.add_leaf(label(md_path, fm), data=md_path)

        # Add grouped directories
        for dir_name in sorted(groups):
            dir_node = self.root.add(dir_name)
            dir_node.expand()
            for md_path, fm in groups[dir_name]:
                dir_node.add_leaf(label(md_path, fm), data=md_path)

    def render_label(self, node, base_style, style):
        node_label = node._label.copy()
        node_label.stylize(style)

        def toggled_icon():
            icon = self.ICON_NODE_EXPANDED if node.is_expanded else self.ICON_NODE
            return icon, base_style + TOGGLE_STYLE

        prefix = toggled_icon() if node._allow_expand else ('', base_style)
        return Text.assemble(prefix, node_label)


def code_block(lang, code):
    return f'```{lang}\n{code}\n```\n'


def fix_code_blocks(s):
    return s.replace('```iml', '```ocaml')


content_wrappers = {
    '': fix_code_blocks,
    '.md': fix_code_blocks,
    '.iml': partial(code_block, 'ocaml'),
    '.py': partial(code_block, 'python'),
    '.ts': partial(code_block, 'typescript'),
}


class HelpScreen(Screen):
    DEFAULT_CSS = 'TOC { width: 25% } VerticalScroll { width: 75% }'

    def __init__(self):
        Screen.__init__(self)
        self.title = 'Help'

    def compose(self):
        files = collect_md_files() + collect_iml_files()

        yield MyHeader()
        with Horizontal():
            yield TOC(files, SKILL_DIR)
            with VerticalScroll():
                markdown = Markdown(
                    '# CodeLogician\nClick the tree to read the docs.', id='content'
                )
                yield markdown
        yield Footer()

    def on_tree_node_selected(self, event):
        def select_node(md_path):
            text = md_path.read_text()
            _, body = parse_frontmatter(text)
            content = fix_code_blocks(body)
            name = md_path.stem
            self.query_one('#content', Markdown).update(f'# {name}\n{content}')

        maybe(select_node, event.node.data)
