import glob
import os
from pathlib import Path
import re

import yaml

from scripts.docs.lit_to_doc import lit_to_doc

SRC_BASE = "bazel-bin/lib"
DEST_BASE = "docs/content/en/docs"
PASSES_FILE = f"{DEST_BASE}/passes.md"
FRONTMATTER_PATTERN = r"^---\n(.*?)---\n"


def cleanup(content):
  content = content.replace(
      "<!-- Autogenerated by mlir-tblgen; don't manually edit -->", ""
  )
  return content


def extract_frontmatter(content):
  match = re.search(FRONTMATTER_PATTERN, content, re.DOTALL)
  if match:
    frontmatter = match.group(1)
    return yaml.safe_load(frontmatter), content[match.end() :]
  return None, content


def split_sections(content):
  # Assumes only level ### headers delineate sections, (#### is used for pass options)
  sections = re.split(r"(^### .*$)", content, flags=re.MULTILINE)
  return sections


def sort_sections(sections):
  headers_and_content = []

  for i in range(1, len(sections), 2):
    header = sections[i].strip()
    body = sections[i + 1]
    headers_and_content.append((header, body))

  sorted_sections = sorted(headers_and_content, key=lambda x: x[0])
  return sorted_sections


def rebuild_content(frontmatter, sorted_sections):
  frontmatter_str = (
      "---\n" + yaml.dump(frontmatter) + "---\n\n" if frontmatter else ""
  )
  content_str = "".join([
      f"{header.strip()}\n\n{body.strip()}\n\n"
      for header, body in sorted_sections
  ])
  return frontmatter_str + content_str


def sort_markdown_file_by_header(path):
  with open(path, "r") as f:
    content = f.read()

  frontmatter, content_without_frontmatter = extract_frontmatter(content)
  sections = split_sections(content_without_frontmatter)
  sorted_sections = sort_sections(sections)
  sorted_content = rebuild_content(frontmatter, sorted_sections)

  with open(path, "w") as f:
    f.write(sorted_content)


def generate_doctest_examples():
  print("Generating doctests")
  pattern = re.compile(r"^ *(\(\* example filepath=)(.*?)(\s*\*\))")

  def replacer(match):
    captured_filepath = match.group(2)
    print(f"Running doctest for {captured_filepath}")
    return lit_to_doc(lit_test_file=captured_filepath)

  new_lines = []
  with open(PASSES_FILE, "r") as passes_file:
    for line in passes_file:
      modified_line = pattern.sub(replacer, line)
      new_lines.append(modified_line)

  with open(PASSES_FILE, "w") as passes_file:
    passes_file.write("".join(new_lines))


if __name__ == "__main__":
  # Create passes.md file with the front matter
  with open(PASSES_FILE, "w") as f:
    f.write("""---
    title: Passes
    weight: 70
    ---\n""")

  print("Processing Passes")
  passes_files = glob.glob(f"{SRC_BASE}/**/*Passes.md", recursive=True)
  conversion_files = glob.glob(
      f"{SRC_BASE}/Dialect/**/Conversions/**/*.md", recursive=True
  )
  for src_path in set(passes_files + conversion_files):
    print(src_path)
    with open(src_path, "r") as src_file:
      with open(PASSES_FILE, "a") as dest_file:
        dest_file.write(cleanup(src_file.read()))
        dest_file.write("\n")

  sort_markdown_file_by_header(PASSES_FILE)
  generate_doctest_examples()

  print("Processing Dialects")
  Path(f"{DEST_BASE}/Dialects/").mkdir(parents=True, exist_ok=True)
  for dialect_dir in glob.glob(f"{SRC_BASE}/Dialect/*"):
    if not os.path.isdir(dialect_dir):
      print(f"Skipping non-directory file {dialect_dir}")
      continue

    dialect_name = os.path.basename(dialect_dir)
    print(f"Processing {dialect_name}")
    filename = f"{dialect_name}.md"
    dest_path = f"{DEST_BASE}/Dialects/{filename}"

    markdown_files = {}
    for src_path in glob.glob(f"{dialect_dir}/IR/**/*.md", recursive=True):
      markdown_files[src_path] = True
      print(f"Adding {src_path} to the queue")

    if not markdown_files:
      print(f"Skipping {dialect_name} as no markdown files found")
      continue

    # Write front matter for the dialect markdown
    with open(dest_path, "w") as f:
      f.write(f"""---
title: {dialect_name}
github_url: https://github.com/google/heir/edit/main/lib/Dialect/{dialect_name}/IR
---\n""")

    # Process files in a specific order
    groups = ["Dialect", "Attributes", "Types", "Ops"]

    for index, group in enumerate(groups):
      search_pattern = f"{dialect_dir}/IR/*{group}.md"
      for src_path in glob.glob(search_pattern):
        if not os.path.isfile(src_path):
          continue

        with open(dest_path, "a") as dest_file:
          if index == 0:
            # Special care for the first group
            with open(src_path, "r") as src_file:
              content = (
                  src_file.read().replace("# Dialect", "").replace("[TOC]", "")
              )
              dest_file.write(content)
          else:
            dest_file.write(f"## {dialect_name} {group.lower()}\n")
            with open(src_path, "r") as src_file:
              dest_file.write(src_file.read())

        markdown_files.pop(src_path, None)

    # Include additional files not processed in the groups
    if markdown_files:
      with open(dest_path, "a") as dest_file:
        dest_file.write(f"## {dialect_name} additional definitions\n")
        for src_path in list(markdown_files.keys()):
          if os.path.isfile(src_path):
            with open(src_path, "r") as src_file:
              dest_file.write(src_file.read())
            markdown_files.pop(src_path)
