import pathlib


def split_doc_path(filepath):
  """Return paths for docs root prefix directory and the relative path to file.

  Given a full path to notebook file, standalone or within an established
  documentation directory layout, split the provided path into two:
  1. a path reprsenting the prefix directory to the docs root (if it exists),
  2. the relative path to the file from the docs root directory.
  If in an unknown docs directory layout, return an empty prefix path and the
  full path of the original argument.

  For example:
  "site/en/notebook.ipynb" => ("site/en", "notebook.ipynb")
  "tensorflow/docs/notebook.ipynb" => ("docs", "notebook.ipynb")
  "unknown/path/notebook.ipynb" => ("", "unknown/path/notebook.ipynb")

  Args:
    filepath: `pathlib.Path` to a documentation notebook.

  Returns:
    pathlib.Path: The path of the doc root prefix directory., if applicable.
    pathlib.Path: The relative path to notebook from the prefix directory.
  """
  fp_full = filepath.resolve()  # Check full path for sub-elements.

  def split_path_on_dir(fp, dirname, offset=1):
    parts = fp.parts
    idx = parts.index(dirname)
    docs_dir = pathlib.Path(*parts[idx:idx + offset])
    rel_path = fp.relative_to(*parts[:idx + offset])
    return docs_dir, rel_path

  if "site" in fp_full.parts:
    return split_path_on_dir(fp_full, "site", offset=2)  # site/<lang>/
  elif "docs" in fp_full.parts:
    return split_path_on_dir(fp_full, "docs")
  elif "g3doc" in fp_full.parts:
    idx = fp_full.parts.index("g3doc")
    if fp_full.parts[idx + 1] == "en":
      offset = 2
    else:
      offset = 1
    return split_path_on_dir(fp_full, "g3doc", offset=offset)
  else:
    # Unknown setup. Return empty root and unsplit path.
    return pathlib.Path(), filepath