from TexSoup import TexNode
from typing import Optional

from texmemo.types import SectionConfig, ParsedEnv
from texmemo.parser import get_environments


class Filter():
    def __init__(self) -> None:
        self._soup: TexNode

    def filter(self,
               soup: TexNode,
               title: Optional[str],
               author: Optional[str],
               sections: list[SectionConfig]
               ):
        self._soup = soup
        self._process_preamble(title, author)
        return self._filter_doc(sections)

    def _get_document(self) -> TexNode:
        docclass = self._soup.find("documentclass")
        if docclass is None:
            self._soup.insert(0, r"\documentclass{article}")
        else:
            docclass.string = 'article'

        doc = self._soup.find("document")
        if doc is None:
            raise ValueError("Could not find document within .tex file.")

        return doc

    def _pick_sections(self, doc: TexNode, sections: list[SectionConfig]) -> list[ParsedEnv]:
        envs = get_environments(doc)
        section_names = [config.name for config in sections]
        child_names = {config.name: config.child for config in sections}
        env_length = len(envs)
        parsed_envs: list[ParsedEnv] = []

        for i, env in enumerate(envs):
            parsed_env = self._get_parsed_env(
                child_names, env, envs, i, section_names, env_length)
            if parsed_env is not None:
                parsed_envs.append(parsed_env)
        return parsed_envs

    def _get_parsed_env(self, child_names, env, envs, env_idx, section_names, env_length) -> ParsedEnv | None:
        child_name = child_names.get(env.name)
        child = None
        if env.name in section_names:
            if env_idx < env_length - 1 and child_name is not None:
                next_env = envs[env_idx+1]
                if next_env.name == child_name:
                    child = next_env
            parsed_env = ParsedEnv(env=env, child=child)
            return parsed_env

    def _filter_doc(self, sections: list[SectionConfig]) -> list[ParsedEnv]:
        doc = self._get_document()
        parsed_sections = self._pick_sections(doc, sections)

        # Clear
        doc.contents = []

        # Rebuild
        doc.append("\n\\maketitle\n\n")
        for sec in parsed_sections:
            doc.append(sec.env)
            if sec.child:
                doc.append("\n")
                doc.append(sec.child)
            doc.append("\n\n")

        return parsed_sections

    def _process_preamble(self, title: Optional[str], author: Optional[str]):
        if author is not None:
            self._replace_author(author)
        if title is not None:
            self._replace_title(title)

    def _replace_author(self, author: str):
        if self._soup.author:
            self._soup.author.string = author
        else:
            self._soup.insert(1, f"\n\\author{{{author}}}")

    def _replace_title(self, title):
        if self._soup.title:
            self._soup.title.string = title
        else:
            self._soup.insert(1, f"\n\\title{{{title}}}")
