#!/usr/bin/env python3
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import gzip
import os
import random
import re
import shutil
import subprocess
import tempfile
import time
import urllib.request
from io import TextIOWrapper

DEBUG = False

TARGET_DOC_CHARS = 1024


def compress_with_seek_points(file_name_in: str, file_name_out: str, num_seek_points: int):
  bytes_per_chunk = os.path.getsize(file_name_in) / num_seek_points

  seek_points: list[int] = []

  if os.path.exists(file_name_out):
    os.remove(file_name_out)

  with open(file_name_in, "rb") as f_in:
    f_out = None

    bytes_in_chunk = 0

    chunk_count = 0

    while True:
      if f_out is None:
        if os.path.exists(file_name_out):
          seek_points.append(os.path.getsize(file_name_out))
          print("  create chunk %s at pos=%s" % (chunk_count, seek_points[-1]))
        else:
          print("  create chunk %s at pos=0" % chunk_count)
        f_out = gzip.open(file_name_out, "ab")
        chunk_count += 1

      line = f_in.readline()
      if len(line) == 0:
        break

      bytes_in_chunk += len(line)
      f_out.write(line)

      if bytes_in_chunk > bytes_per_chunk and chunk_count < num_seek_points:
        f_out.close()
        f_out = None
        bytes_in_chunk = 0

  with open(file_name_out[:-3] + ".seek", "w") as f_out:
    for seek_point in seek_points:
      f_out.write("%d\n" % seek_point)


re_tag = re.compile(r"<[^>]+?>")
re_newlines = re.compile(r"\n+")
re_space = re.compile(r"\s")

# used to find word break, for splitting docs into ~1 KB sized smaller docs:
re_next_non_word_character = re.compile(r"\W", re.UNICODE)

EUROPARL_V7_URL = "https://www.statmt.org/europarl/v7/europarl.tgz"


def split_docs(all_out: TextIOWrapper, title_string: str, date_string: str, body_string: str):
  """Splits docs into smallish (~1 KB) sized docs, repeating same title and date"""
  doc_count = 0
  while len(body_string) > 0:
    char_count = int(random.gauss(TARGET_DOC_CHARS, TARGET_DOC_CHARS / 4))
    if char_count < 64:
      # trimmed normal?
      continue

    m = re_next_non_word_character.search(body_string, char_count)
    if m is not None:
      char_count = m.start(0)
    else:
      char_count = len(body_string)

    body_string_fragment = body_string[:char_count].strip()

    # print('write title %d, body %d' % (len(title_string), len(body_string_fragment)))
    all_out.write("%s\t%s\t%s\n" % (title_string, date_string, body_string_fragment))
    body_string = body_string[char_count:]
    doc_count += 1

  return doc_count


def sample_europarl():
  # download europarl.tgz v7, if not already here (in cwd):
  file_name = "europarl.tgz"
  if not os.path.exists(file_name):
    print("Download %s to %s..." % (EUROPARL_V7_URL, file_name))
    urllib.request.urlretrieve(EUROPARL_V7_URL, file_name + ".tmp")
    os.rename(file_name + ".tmp", file_name)
  else:
    print("%s already here; skipping download..." % file_name)

  if not DEBUG:
    tmp_dir_path = tempfile.mkdtemp()
  else:
    tmp_dir_path = "/tmp/tmp31ekzg75"
  print('Using tmp dir "%s"...' % tmp_dir_path)
  try:
    if not DEBUG:
      cmd = "tar xzf %s -C %s" % (file_name, tmp_dir_path)
      print("Run: %s" % cmd)
      subprocess.run(cmd, shell=True, check=False)

    doc_count = 0
    skip_count = 0
    file_count = 0

    all_txt_file_name = "%s/all.txt" % tmp_dir_path

    print("Extract text...")

    start_time = time.time()
    next_print_time = start_time + 3
    # normalize text a bit and concatenate all lines into single file, counting total lines/bytes
    with open(all_txt_file_name, "w", encoding="utf-8") as all_out:
      for dir_path, _, file_names in os.walk("%s/txt" % tmp_dir_path):
        for file_name in file_names:
          if file_name.endswith(".txt"):
            file_count += 1

            year, month, day = (int(x) for x in file_name[3:-4].split("-")[:3])
            if year >= 50:
              year = 1900 + year
            else:
              year = 2000 + year

            date_string = "%04d-%02d-%02d" % (year, month, day)

            # unfortunately we need errors='ignore' since in Europarl v7, one file (pl/ep-09-10-22-009.txt) has invalid utf-8:
            chapter_count = 0
            with open("%s/%s" % (dir_path, file_name), encoding="utf-8", errors="ignore") as f_in:
              last_text: list[str] = []
              last_title = None
              while True:
                line = f_in.readline()
                if line == "":
                  break
                line = line.strip()
                if line.startswith("<CHAPTER "):
                  if last_title is not None:
                    s = " ".join(last_text)
                    s = re_tag.sub(" ", s)
                    s = re_newlines.sub(" ", s)
                    s = s.strip()
                    if len(s) > 0:
                      doc_count += split_docs(all_out, last_title, date_string, s)
                    else:
                      skip_count += 1

                    last_text = []
                    chapter_count += 1
                  while True:
                    last_title = f_in.readline()
                    if last_title == "":
                      last_title = None
                      break
                    last_title = re_tag.sub(" ", last_title).strip()
                    if len(last_title) > 0:
                      break
                  continue
                last_text.append(line)

              if last_title is not None:
                s = " ".join(last_text)
                s = re_tag.sub(" ", s)
                s = re_newlines.sub(" ", s)
                s = s.strip()
                if len(s) > 0:
                  doc_count += split_docs(all_out, last_title, date_string, s)
                else:
                  skip_count += 1
                chapter_count += 1
              else:
                skip_count += 1

              if chapter_count > 0:
                # print('%s/%s: %d chapters' % (dir_path, file_name, chapter_count))
                pass

            now = time.time()
            if now > next_print_time:
              print(
                "%4.1fs: keep %.2f K of %.2f K files (%.1f%%), %.2f M docs, %.2f GB..."
                % (now - start_time, (file_count - skip_count) / 1000, file_count / 1000, 100 * (file_count - skip_count) / file_count, doc_count / 1000000, all_out.tell() / 1024 / 1024 / 1024)
              )
              while next_print_time < now:
                next_print_time += 3

    total_mb = os.path.getsize(all_txt_file_name) / 1024 / 1024
    now = time.time()
    print(
      "%4.1fs (done): keep %.2f K of %.2f K files (%.1f%%), %.2f M docs, %.2f GB..."
      % (
        now - start_time,
        (file_count - skip_count) / 1000,
        file_count / 1000,
        100 * (file_count - skip_count) / file_count,
        doc_count / 1000000,
        os.path.getsize(all_txt_file_name) / 1024 / 1024 / 1024,
      )
    )

    print("Shuffle...")
    subprocess.run("shuf %s > %s.shuffled" % (all_txt_file_name, all_txt_file_name), shell=True, check=False)

    for mb in (20, 200, 2000):
      print("Sample %d MB file..." % mb)
      file_name_out = "%dmb.txt" % mb
      with open(file_name_out, "w", encoding="utf-8") as f_out:
        chance = mb / total_mb

        with open(all_txt_file_name + ".shuffled", encoding="utf-8") as f:
          while True:
            line = f.readline()
            if len(line) == 0:
              break
            if random.random() <= chance:
              f_out.write(line)

      print("  got %.2f MB" % (os.path.getsize(file_name_out) / 1024 / 1024))

      compress_with_seek_points(file_name_out, file_name_out + ".gz", mb)

  finally:
    print('Removing tmp dir "%s"...' % tmp_dir_path)
    if not DEBUG:
      shutil.rmtree(tmp_dir_path)

  print("\nWARNING: left ./europarl.tgz, which you should delete if you do not want it!\n")


if False:
  compress_with_seek_points("/x/tmp/europarl.lines.txt", "/x/tmp/foo.txt.gz", 16)
else:
  sample_europarl()
