# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/00_core.ipynb.

# %% auto 0
__all__ = ['extract_comments_from_docx', 'extract_ruby_info', 'extract_paragraphs_with_style', 'merge_runs_with_same_style',
           'split_text', 'convert_json_to_tei', 'format_xml_string', 'convert_docx_to_tei']

# %% ../nbs/00_core.ipynb 3
from bs4 import BeautifulSoup
import docx
import zipfile
from typing import Dict, List, Any

# %% ../nbs/00_core.ipynb 4
def extract_comments_from_docx(input_path: str) -> Dict[str, List[str]]:
    """
    Extract comments from a .docx file and return a dictionary mapping comment IDs to their content.

    Args:
        input_path (str): The path to the .docx file from which to extract comments.

    Returns:
        Dict[str, List[str]]: A dictionary mapping comment IDs (str) to the content of the comments
                            as lists of paragraphs (List[str]).
    """
    import zipfile
    from bs4 import BeautifulSoup

    with zipfile.ZipFile(input_path, 'r') as zfp:
        with zfp.open('word/comments.xml') as fp:
            soup = BeautifulSoup(fp.read(), 'xml')

    comments = soup.find_all("w:comment")
    comment_map = {}

    for comment in comments:
        cid = comment.attrs["w:id"]
        texts = [p.text for p in comment.find_all("w:p")]
        comment_map[cid] = texts

    return comment_map

def extract_ruby_info(run) -> dict:
    """
    Extracts ruby information from a given run in a Word document.

    Args:
    - run: A Run object from python-docx.

    Returns:
    - A dictionary containing information about the run.
    """

    # Create a dictionary to store run data
    run_data = {
        "run_text": run.text,
        "run_style_name": run.style.name,
    }

    # Extract any ruby text
    rt_elements = run._r.xpath("w:ruby/w:rt/w:r/w:t")
    rb_elements = run._r.xpath("w:ruby/w:rubyBase/w:r/w:t")
    rb_rubyAlign = run._r.xpath("w:ruby/w:rubyPr/w:rubyAlign")

    if len(rb_elements) > 0:
        # If there is ruby text, extract it and add it to the run data
        run_data["rb"] = rb_elements[0].text

        # Extract any ruby text in <w:rt> tags and remove any parentheses
        rt_text = ""
        for rt in rt_elements:
            kigo_list = ["（", "）", "〔", "〕"]
            t = rt.text
            for kigo in kigo_list:
                t = t.replace(kigo, "")
            rt_text += t

        # Add the ruby text to the run data
        run_data["rt"] = rt_text

        if len(rb_rubyAlign) > 0:
            # Extract the ruby alignment, if it exists, and add it to the run data
            align = rb_rubyAlign[0].attrib['{http://schemas.openxmlformats.org/wordprocessingml/2006/main}val']
            run_data["align"] = align

    # Return the run data as a dictionary
    return run_data


def extract_paragraphs_with_style(input_path: str) -> Dict[int, Dict[str, Any]]:
    """
    Extract paragraphs with their style and associated runs from a .docx file, and return
    a dictionary mapping paragraph indices to their content and style.

    Args:
        input_path (str): The path to the .docx file from which to extract paragraphs and style.

    Returns:
        Dict[int, Dict[str, Any]]: A dictionary mapping paragraph indices (int) to their content
                                and style as a dictionary containing 'style' and 'values' keys.
    """
    doc = docx.Document(input_path)
    result = {}

    for index, para in enumerate(doc.paragraphs):
        runs = []

        for run in para.runs:
            run_data = extract_ruby_info(run)

            runs.append(run_data)

        result[index] = {
            "style": para.style.name,
            "values": runs
        }

    return result

def merge_runs_with_same_style(paragraph_data: Dict[int, Dict[str, Any]]) -> Dict[int, Dict[str, Any]]:
    """
    Merge consecutive runs with the same style in the given paragraph data.
    
    Args:
        paragraph_data (Dict[int, Dict[str, Any]]): A dictionary mapping paragraph indices to their
                                                    content and style as a dictionary containing 'style'
                                                    and 'values' keys.

    Returns:
        Dict[int, Dict[str, Any]]: A dictionary with the same structure as the input, but with consecutive
                                runs with the same style merged together.
    """
    merged_data = {}

    for index, para in paragraph_data.items():
        merged_runs = []
        previous_run = None

        # Iterate through each run in the paragraph
        for run in para["values"]:
            # If this is the first run, or the run_style_name is different from the previous run's
            if previous_run is None or run["run_style_name"] != previous_run["run_style_name"] or run["run_style_name"] == "tei:ruby":
                merged_runs.append(run)
            else:
                # Concatenate the run_text of the current run with the previous run
                merged_runs[-1]["run_text"] += run["run_text"]

                # If the current run has rt (ruby text), update the previous run's rt
                if "rt" in run:
                    if "rt" not in merged_runs[-1]:
                        merged_runs[-1]["rt"] = ""
                    merged_runs[-1]["rt"] += run["rt"]

                # If the current run has rb (ruby base), update the previous run's rb
                if "rb" in run:
                    if "rb" not in merged_runs[-1]:
                        merged_runs[-1]["rb"] = ""
                    merged_runs[-1]["rb"] += run["rb"]

            # Update the previous_run for the next iteration
            previous_run = run

        # Store the merged runs in the merged_data dictionary
        merged_data[index] = {
            "style": para["style"],
            "values": merged_runs
        }

    return merged_data

# %% ../nbs/00_core.ipynb 5
def split_text(string):
    half_length = len(string) // 2
    first_half = string[:half_length]
    second_half = string[half_length:]
    return first_half, second_half

def convert_json_to_tei(json_data: Dict[int, Dict[str, Any]], base_path) -> str:
    """
    Convert the given JSON data to TEI XML.

    Args:
        json_data (Dict[int, Dict[str, Any]]): A dictionary mapping paragraph indices to their
                                            content and style as a dictionary containing 'style'
                                            and 'values' keys.

    Returns:
        str: The TEI XML representation of the given JSON data.
    """

    with open(base_path, "r") as file:
        content = file.readlines()
        content = "".join(content)
        bs_content = BeautifulSoup(content, "xml")

    parent = bs_content.find("body")

    p = bs_content.new_tag("p")
    parent.append(p)

    target = p

    for index, para in json_data.items():

        para_style = para["style"]

        if para_style == "tei:head":

            values = para["values"]

            if len(values) == 0:
                continue

            div = bs_content.new_tag("div")
            div["type"] = "section"
            parent.append(div)

            head = bs_content.new_tag("head")

            head.append(values[0]["run_text"])
            div.append(head)

            p = bs_content.new_tag("p")
            div.append(p)

            target = p

            continue

        else:

            lb = bs_content.new_tag("lb")
            # lb["n"] = str(index + 1)
            target.append(lb)

            # target = ab

        # values
        for run in para["values"]:
            run_style_name = run["run_style_name"]
            run_text = run["run_text"]
            if run_style_name == "tei:pb":
                pb = bs_content.new_tag("pb")
                pb["n"] = run["run_text"] # .replace("p", "")
                pb["corresp"] = f"#{pb['n']}"
                target.append(pb)

            elif run_style_name == "tei:del":
                e = bs_content.new_tag("del")

                run_text = run_text.replace("（抹消）", "")
                e.append(run_text)
                target.append(e)

            elif run_style_name == "tei:warichu":
                e = bs_content.new_tag("note")
                target.append(e)
                
                e["type"] = "割書"

                f, l = split_text(run["run_text"])

                e.append(f)

                milestone = bs_content.new_tag("milestone")
                milestone["unit"] = "wbr"
                e.append(milestone)

                e.append(l)

            elif run_style_name == "tei:add":
                e = bs_content.new_tag("add")
                target.append(e)
                e.append(run_text)

            elif run["run_style_name"] in ["tei:persName", "tei:placeName", "tei:name", "tei:date"]:

                run_style_name = run["run_style_name"]
                name = ""
                if run_style_name == "tei:persName":
                    name = "persName"
                elif run_style_name == "tei:placeName":
                    name = "placeName"
                elif run_style_name == "tei:name":
                    name = "name"
                elif run_style_name == "tei:date":
                    name = "date"

                # print(run)

                e = bs_content.new_tag(name)
                target.append(e)

                if "rb" in run:

                    rb = run["rb"]
                    rt = run["rt"]
                    
                    e.append(run["rb"])

                    if run_style_name == "tei:date":
                        e["when-custom"] = run["rt"]
                    else:
                        e["corresp"] = "#" + run["rt"]
                    pass
                else:
                    e.append(run["run_text"])
                    pass

            elif run["run_style_name"] == "tei:rubyOther":

                ruby = bs_content.new_tag("ruby")
                target.append(ruby)

                rb = bs_content.new_tag("rb")
                rb.string = run["rb"]
                ruby.append(rb)

                rt = bs_content.new_tag("rt")
                rt.string = run["rt"]
                ruby.append(rt)

            elif run["run_style_name"] == "tei:okurikana":
                e = bs_content.new_tag("seg")
                e.string = run["run_text"]
                e["type"] = "okurikana"
                target.append(e)

            elif run["run_style_name"] == "tei:kaeriten":

                metamark = bs_content.new_tag("metamark")
                metamark["function"] = "kaeriten"
                metamark.append(run_text)
                target.append(metamark)

            elif run["run_style_name"] == "tei:choice":

                choice = bs_content.new_tag("choice")
                target.append(choice)

                sic = bs_content.new_tag("sic")
                choice.append(sic)
                sic.string = run["rb"]

                corr = bs_content.new_tag("corr")
                choice.append(corr)
                corr.string = run["rt"]


            elif run["run_style_name"] == "tei:small":

                
                note = bs_content.new_tag("seg")
                note["rend"] = "small"
                note.append(run_text)
                target.append(note)

            elif run["run_style_name"] == "tei:ruby":
                if "rb" not in run:
                    continue

                ruby = bs_content.new_tag("ruby")
                target.append(ruby)

                rb = bs_content.new_tag("rb")
                rb.string = run["rb"]
                ruby.append(rb)

                align = run["align"]

                

                if align == "distributeSpace":
                    spl = run["rt"].split("　")
                    if len(spl) == 2:
                        yomi, okuri = spl
                        rt = bs_content.new_tag("rt")
                        rt.string = yomi
                        ruby.append(rt)

                        rt = bs_content.new_tag("rt")
                        rt.string = okuri
                        rt["type"] = "okurigana"
                        ruby.append(rt)
                    else:
                        rt = bs_content.new_tag("rt")
                        rt.string = run["rt"]
                        ruby.append(rt)
                else:
                    rt = bs_content.new_tag("rt")

                    rt.string = run["rt"]

                    if align == "left":
                        # rt["type"] = "yomikana"
                        pass
                    elif align == "right":
                        rt["type"] = "okurigana"

                    ruby.append(rt)

                
            else:
                target.append(run["run_text"])

    return bs_content

# %% ../nbs/00_core.ipynb 6
from xml.dom.minidom import parseString
from xml.dom import minidom
import re
def format_xml_string(xml_string):
    """xml文字列を整形する

    Parameters:  
    * `xml_string` - xml文字列 (str)

    Returns:  
    * 整形されたxml文字列 (str)
    """
    
    xml_string = xml_string.replace("\n", "")

    reparsed = minidom.parseString(xml_string)
    pretty = re.sub(r"[\t ]+\n", "", reparsed.toprettyxml(indent="\t"))
    pretty = pretty.replace(">\n\n\t<",">\n\t<")
    return pretty

# %% ../nbs/00_core.ipynb 7
def convert_docx_to_tei(input_path, output_path, template_path):
    """docxファイルをTEIに変換する

    Parameters:  
    * `docx_file_path` - docxファイルのパス (str)

    Returns:  
    * TEIのxml文字列 (str)
    """
    tmp = extract_paragraphs_with_style(input_path)
    tmp2 = merge_runs_with_same_style(tmp)

    '''
    import json
    with open("tmp.json", "w") as f:
        json.dump(tmp2, f, ensure_ascii=False, indent=4)
    '''

    xml = convert_json_to_tei(tmp2, template_path)
    xml_str = format_xml_string(str(xml))
    with open(output_path, "w") as f:
        f.write(xml_str)

