
import re
import pysubs2
import os
from collections import Counter
from tqdm.notebook import tqdm
from pathlib import Path
from typing import Union, Dict, Iterable, List, Optional, Tuple, Any
import pickle
import logging
import chardet
import re
from bs4 import BeautifulSoup
import unicodedata
from langdetect import detect
from langdetect.lang_detect_exception import LangDetectException
from types import SimpleNamespace
import fitz
import zipfile
import shutil
import tarfile
from concurrent.futures import ProcessPoolExecutor, as_completed

from ankipan import PROJECT_ROOT
from ankipan.lemma_parsers import DictValidator, load_lemmatizer, supported_languages

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

SEGMENT_DELIMITER = '\n\n'

ASCII_PUNCT = r"!\"#\$%&'\(\)\*\+,\-./:;<=>\?@\[\]\\\^_`{\|\}~"
# General + Supplemental + CJK punct + Fullwidth punct (subset)
UNI_PUNCT_BLOCKS = "\u2000-\u206F\u2E00-\u2E7F\u3000-\u303F\uFF01-\uFF65"

ALLOWED = rf"\w\s{UNI_PUNCT_BLOCKS}{ASCII_PUNCT}"
pattern = re.compile(rf"[^{ALLOWED}]")

_WORKER_VALIDATOR: DictValidator | None = None

def _init_worker_validator(learning_lang: str, project_root: Path) -> None:
    global _WORKER_VALIDATOR
    _WORKER_VALIDATOR = load_lemmatizer(learning_lang, project_root=project_root)
    if _WORKER_VALIDATOR is None:
        raise RuntimeError(f"No dict validator found for language '{learning_lang}'")

def _process_file_in_worker(file, save_pos_data: bool, save_sentence_mapping: bool):
    if _WORKER_VALIDATOR is None:
        raise RuntimeError("Worker dict validator not initialized")
    file.extract_lemmas(
        validator=_WORKER_VALIDATOR,
        save_pos_data=save_pos_data,
        save_sentence_mapping=save_sentence_mapping,
    )
    return file

class Reader():
    SUPPORTED_LEARNING_LANGUAGES = supported_languages()
    ZIP_ALLOWED_SUFFIXES = {
        '.ass', '.srt', '.pdf', '.html', '.htm',
        '.txt', '.md', '.csv', '.tsv', '.log',
    }

    def __init__(self, learning_lang):
        if learning_lang not in self.SUPPORTED_LEARNING_LANGUAGES:
            raise RuntimeError(
                f"Unsupported learning_lang '{learning_lang}'. Supported languages: {sorted(self.SUPPORTED_LEARNING_LANGUAGES)}"
            )
        self.learning_lang = learning_lang
        self.validator: DictValidator | None = load_lemmatizer(learning_lang, project_root=PROJECT_ROOT)
        self.corpus_occurrences_ = None
        if self.validator is None:
            raise RuntimeError(f"No dict validator found for language '{learning_lang}'")

    @property
    def corpus_occurrences(self):
        if self.corpus_occurrences_ is None:
            try:
                # TODO: Dependency to parent folder not good
                with open(f'{PROJECT_ROOT}/corpus_resources/{self.learning_lang}.pkl', "rb") as f:
                    self.corpus_occurrences_ = pickle.load(f)
            except:
                self.corpus_occurrences_ = Counter()
        return self.corpus_occurrences_

    def collect_file_paths(self, path, exclude_paths = None, file_match_pattern = None, dir_match_pattern = None):
        compiled_file_pattern = re.compile(file_match_pattern) if file_match_pattern else None
        compiled_dir_pattern = re.compile(dir_match_pattern) if dir_match_pattern else None
        if not exclude_paths:
            exclude_paths = set()
        else:
            exclude_paths = {os.path.normcase(Path(path).as_posix()) for path in exclude_paths}
        def allowed(p: Path) -> bool:
            path_normed = os.path.normcase(p.as_posix())
            if path_normed in exclude_paths:
                logger.info(f'Skipping path {p}, listed in exclude paths')
                return False
            if compiled_file_pattern and not compiled_file_pattern.match(p.name):
                logger.info(f'Skipping path {p}, matches compiled_file_pattern')
                return False
            return True

        file_paths = set()
        path = Path(path)

        if path.is_file():
            cleaned_path = Path(path).relative_to(path.parent)
            if allowed(cleaned_path):
                file_paths.add(path)

        elif path.is_dir():
            for root, dirs, files in os.walk(path):
                if compiled_dir_pattern:
                    dirs[:] = [d for d in dirs if compiled_dir_pattern.fullmatch(d)]
                    if not compiled_dir_pattern.fullmatch(Path(root).name):
                        continue
                for filename in files:
                    if compiled_file_pattern and not compiled_file_pattern.fullmatch(filename):
                        continue
                    cleaned_path = Path(root, filename).relative_to(path.parent)
                    if allowed(cleaned_path):
                        file_paths.add(Path(root, filename))
        else:
            raise RuntimeError(f'Path does not exist: "{path}"')
        return file_paths

    def detect_languages(self, file_contents, n_chunks = 10):
        """
        Currently used to assert integrity of files (only one language used)

        """
        part_length = len(file_contents) // n_chunks
        parts = [file_contents[i * part_length:(i + 1) * part_length] for i in range(n_chunks)]

        detected_languages = set()
        for part in parts:
            try:
                lang = detect(part)
                detected_languages.add(lang)
            except LangDetectException:
                continue
        return detected_languages

    def clean_string(self, string, chars_to_filter=None):
        ESCAPE_PATTERN = re.compile(r"\\[A-Za-z0-9_]+")
        # https://www.unicode.org/reports/tr44/tr44-34.html#General_Category_Values
        REMOVE_CATEGORIES = {
            "Cc", "Cf", "Cs", "Co", "Cn",
            "Zl", "Zp",
        }
        string = string.replace("\\N", "\n").replace("\\n", "\n")
        string = ESCAPE_PATTERN.sub("", string)
        result_chars = []
        for c in string:
            cat = unicodedata.category(c)
            if c == "\n":
                result_chars.append(c)
            elif cat not in REMOVE_CATEGORIES:
                result_chars.append(c)

        string = "".join(result_chars)
        if chars_to_filter:
            for char_to_filter in chars_to_filter:
                replacement = " "
                if isinstance(char_to_filter, tuple):
                    char_to_filter, replacement = char_to_filter
                string = re.sub(char_to_filter, replacement, string)
        return string

    @staticmethod
    def prepare_stanza_input(raw_text, index_separators, secondary_separators, lookahead=24):
        text = raw_text
        for sep in index_separators:
            text = text.replace(sep, f"{sep}{SEGMENT_DELIMITER}")
        stanzas = text.split(SEGMENT_DELIMITER)

        def split_soft(s: str, L: int, look: int = 24) -> list[str]:
            s = s.strip()
            if L is None or L <= 0 or len(s) <= L:
                return [s] if s else []
            out, i, n = [], 0, len(s)
            while n - i > L:
                window = s[i:i+L+1]
                left_ws = [m.start() for m in re.finditer(r"\s", window)]
                if left_ws:
                    cut = i + left_ws[-1]
                else:
                    m = re.search(r"\s", s[i+L:i+L+look+1])
                    cut = (i + L + m.start()) if m else (i + L)
                head = s[i:cut].rstrip()
                if head:
                    out.append(head)
                i = cut
                while i < n and s[i].isspace():
                    i += 1
            tail = s[i:].strip()
            if tail:
                out.append(tail)
            return out

        final = []
        for s in stanzas:
            if not s:
                continue
            tokens = re.split(f"({'|'.join(map(re.escape, secondary_separators))})", s) if secondary_separators else [s]
            out, buf = [], ""
            for t in tokens:
                buf += t
            if buf.strip():
                out.append(buf.strip())
            for piece in out:
                final.extend(split_soft(piece, len(piece), lookahead))
        return SEGMENT_DELIMITER.join(final)

    def open_files(self,
                   file_paths: List[Union[str, Path]],
                   replace_chars: List[str] = None,
                   source_name: str = None,
                   assert_coherence: bool = False,
                   index_separators: Iterable[str] = None,
                   secondary_separators: Iterable[str] = None):

        """
        Extract raw text from list of files.

        Valid formats:
        - Any raw text (.txt etc.)
        - Subtitles (.ass, .srt)
        - Websites (.html)
        - Zip archives containing supported files (.zip)
        - Tar archives containing supported files (.tar, .tar.gz, .tgz)

        Parameters:
        -----------

        file_paths: list of files to process
        replace_chars: List of characters to be removed when processing
        source_name: Allows for special parcing of content on specific websites in parse_html
            Currently only have rule for ['wikipedia']
        assert_coherence: check if files have sentences in more than one language, skip if that is the case.
        index_separators: Custom delimiter to text into small segments.
            Only relevant if we want to create example sentences from this in the db.
            Subtitles are already segmented in small chunks by default.
            If left empty, a custom segmentation of newlines, dots and commas will be used.

        """

        index_separators_ = index_separators if index_separators else ['\n', '.', '。', '!', '?']
        secondary_separators_ = secondary_separators if secondary_separators else [',', ';', '、']
        files = []
        processed_words = {} # caching past results

        pending_paths = [Path(file_path) for file_path in file_paths]
        seen_paths = set()

        while pending_paths:
            file_path = Path(pending_paths.pop(0))
            if file_path in seen_paths:
                continue
            seen_paths.add(file_path)
            logger.debug(f'Opening file {file_path}')
            suffix = file_path.suffix.lower()
            lower_name = file_path.name.lower()

            if lower_name.endswith('.zip'):
                try:
                    extracted_paths = self._ensure_zip_extracted(file_path)
                except Exception as e:
                    logger.error(f'Error: "{file_path}": {type(e).__name__} - {e}')
                    continue
                for extracted_path in extracted_paths:
                    if self._zip_suffix_allowed(extracted_path):
                        pending_paths.append(extracted_path)
                    else:
                        logger.info(f'Skipping extracted file {extracted_path}, unsupported file ending')
                continue
            if lower_name.endswith('.tar') or lower_name.endswith('.tar.gz') or lower_name.endswith('.tgz'):
                try:
                    extracted_paths = self._ensure_tar_extracted(file_path)
                except Exception as e:
                    logger.error(f'Error: "{file_path}": {type(e).__name__} - {e}')
                    continue
                for extracted_path in extracted_paths:
                    if self._zip_suffix_allowed(extracted_path):
                        pending_paths.append(extracted_path)
                    else:
                        logger.info(f'Skipping extracted file {extracted_path}, unsupported file ending')
                continue

            with open(file_path, 'rb') as f:
                raw_data = f.read()
            encoding_result = chardet.detect(raw_data)
            encoding = encoding_result['encoding']
            if not encoding:
                encoding = 'utf-8'
            try:
                raw_file_contents = raw_data.decode(encoding, errors='replace')
            except LookupError as e:
                raise RuntimeError(f'Decode error: {e}')
            file_contents = None
            sub_timestamps = None

            logger.debug(f'Reading file {file_path}')
            if suffix in ['.ass', '.srt']:
                if index_separators:
                    raise RuntimeError('index_separators not used when parsing subs')
                try:
                    file_contents, sub_timestamps = self.parse_subs(raw_file_contents, replace_chars=replace_chars)
                except Exception as e:
                    logger.error(f'Error: "{file_path}": {type(e).__name__} - {e}')
            elif suffix in ['.pdf']:
                try:
                    file_contents = self.parse_pdf(file_path)
                except Exception as e:
                    logger.warning(f'Reading pdf file "{file_path}" failed, skipping...')
            else:
                if suffix in ['.html', '.htm']:
                    raw_file_contents = self.parse_html(raw_file_contents, source_name = source_name)
                file_contents = self.prepare_stanza_input(self.clean_string(raw_file_contents,
                                                                            replace_chars),
                                                                            index_separators_,
                                                                            secondary_separators_)

            if file_contents:
                if assert_coherence and len(file_contents) > 500: # TODO: Find more efficient way to ignore parts of the file that have a different language than the learning one
                    langs = self.detect_languages(file_contents)
                    if len(langs)>1:
                        logger.error(f"Skipping file {file_path}, Multiple languages detected: {langs}")
                        continue
                files.append(File(self.learning_lang,
                                  file_contents,
                                  processed_words=processed_words,
                                  path=file_path,
                                  sub_timestamps=sub_timestamps))
        return sorted(files, key=lambda x: x.path.name)

    def _zip_suffix_allowed(self, path: Path) -> bool:
        return path.suffix.lower() in self.ZIP_ALLOWED_SUFFIXES

    def _ensure_zip_extracted(self, zip_path: Path) -> list[Path]:
        zip_path = Path(zip_path)
        base_dir = zip_path.parent.resolve()
        extracted_paths: list[Path] = []
        missing_entries: list[tuple[zipfile.ZipInfo, Path]] = []

        with zipfile.ZipFile(zip_path) as zf:
            for info in zf.infolist():
                if info.is_dir():
                    continue
                member_path = Path(info.filename)
                if member_path.is_absolute():
                    logger.warning(f'Skipping absolute zip member {member_path} in {zip_path}')
                    continue
                dest_path = (base_dir / member_path).resolve()
                if dest_path != base_dir and base_dir not in dest_path.parents:
                    logger.warning(f'Skipping unsafe zip member {member_path} in {zip_path}')
                    continue
                extracted_paths.append(dest_path)
                if not dest_path.exists():
                    missing_entries.append((info, dest_path))

            if missing_entries:
                logger.info(f'Extracting {len(missing_entries)} files from {zip_path} into {base_dir}')
                for info, dest_path in missing_entries:
                    dest_path.parent.mkdir(parents=True, exist_ok=True)
                    with zf.open(info) as src, open(dest_path, 'wb') as dst:
                        shutil.copyfileobj(src, dst)
                try:
                    zip_path.unlink()
                except OSError as e:
                    logger.warning(f'Failed to delete archive {zip_path}: {e}')

        return extracted_paths

    def _ensure_tar_extracted(self, tar_path: Path) -> list[Path]:
        tar_path = Path(tar_path)
        base_dir = tar_path.parent.resolve()
        extracted_paths: list[Path] = []
        missing_members: list[tuple[tarfile.TarInfo, Path]] = []

        with tarfile.open(tar_path) as tf:
            for member in tf.getmembers():
                if not member.isfile():
                    continue
                member_path = Path(member.name)
                if member_path.is_absolute():
                    logger.warning(f'Skipping absolute tar member {member_path} in {tar_path}')
                    continue
                dest_path = (base_dir / member_path).resolve()
                if dest_path != base_dir and base_dir not in dest_path.parents:
                    logger.warning(f'Skipping unsafe tar member {member_path} in {tar_path}')
                    continue
                extracted_paths.append(dest_path)
                if not dest_path.exists():
                    missing_members.append((member, dest_path))

            if missing_members:
                logger.info(f'Extracting {len(missing_members)} files from {tar_path} into {base_dir}')
                for member, dest_path in missing_members:
                    dest_path.parent.mkdir(parents=True, exist_ok=True)
                    with tf.extractfile(member) as src, open(dest_path, 'wb') as dst:
                        if src is None:
                            continue
                        shutil.copyfileobj(src, dst)
                try:
                    tar_path.unlink()
                except OSError as e:
                    logger.warning(f'Failed to delete archive {tar_path}: {e}')

        return extracted_paths

    def process_files(self, files, save_pos_data=False, n_threads=5, save_sentence_mapping=False, silent=False):
        if not self.validator:
            raise RuntimeError(f"No dict validator configured for language '{self.learning_lang}'")
        if n_threads <= 1 or len(files) <= 1:
            iterator = files if silent else tqdm(files, total=len(files), desc="Extracting lemmas")
            for f in iterator:
                f.extract_lemmas(
                    validator=self.validator,
                    save_pos_data=save_pos_data,
                    save_sentence_mapping=save_sentence_mapping,
                )
            return

        results = [None] * len(files)
        progress = None if silent else tqdm(total=len(files), desc="Extracting lemmas")
        with ProcessPoolExecutor(
            max_workers=n_threads,
            initializer=_init_worker_validator,
            initargs=(self.learning_lang, PROJECT_ROOT),
        ) as executor:
            future_by_idx = {
                executor.submit(
                    _process_file_in_worker,
                    file,
                    save_pos_data,
                    save_sentence_mapping,
                ): idx
                for idx, file in enumerate(files)
            }
            for future in as_completed(future_by_idx):
                idx = future_by_idx[future]
                results[idx] = future.result()
                if progress is not None:
                    progress.update(1)
        if progress is not None:
            progress.close()
        for idx, processed in enumerate(results):
            files[idx] = processed

    def parse_html(self, file_contents: str, source_name = None) -> str:
        """
        Clean html files to extract text (relevant for scraped pages)

        """
        soup = BeautifulSoup(file_contents, 'lxml')
        if source_name == 'wikipedia':
            if soup.find('meta', attrs={'http-equiv': 'Refresh'}):
                return ""
            body_content_div = soup.find('div', id='bodyContent')
            text = ""
            if body_content_div:
                paragraphs = body_content_div.find_all('p')
                for paragraph in paragraphs:
                    text += paragraph.get_text() + "\n"
            clean_text = re.sub(r'\[.*?\]', '', text)
            clean_text = re.sub(r'[^\x00-\x7F]+', '', clean_text)
        else:
            clean_text = soup.get_text()
        return clean_text.strip()

    sub_filter_expressions = {
        'jp': [r'《', r'》', r'（.*?）', r'\[.*\]', r'\{.*?\}', r'\(.*?\)', '\u3000', '\n', '“','”','「','」',
            r'\\N', r'\n', '…', '→', '〉', '〈', '‥', '⁉', r'[\u2460-\u2473]', '‪', '‬', r'\r', r'  '],
        'de': [r'\{.*?\}', r'\\N'],
    }

    def parse_subs(self, file_contents: str, replace_chars: list[str] | None = None):
        try:
            subs = pysubs2.SSAFile.from_string(file_contents)
        except Exception:
            subs = pysubs2.SSAFile.from_string(file_contents, format="srt")

        def is_jp_dialogue(text: str) -> bool:
            non_dialogue = [r'---', r'翻译:', r'www\.', '字幕組', '校对', '后期', '广告', '制作成员']
            if any(re.search(p, text) for p in non_dialogue):
                return False
            # if re.search(r'[\u3000\u3040-\u30ff\u4e00-\u9fff]+', text):
            #     return True
            return True

        segment_texts: list[str] = []
        sub_timestamps: list[tuple[int, int]] = []

        for line in subs:
            raw = line.text
            cleaned = self.clean_string(raw, self.sub_filter_expressions.get(self.learning_lang, []) + (replace_chars or []))
            if self.learning_lang == 'jp' and not is_jp_dialogue(cleaned):
                continue
            cleaned = cleaned.replace("\r\n", " ").replace("\n", " ").strip()
            cleaned = re.sub(r"\s+", " ", cleaned)
            if not cleaned:
                continue
            start_sec = round(line.start / 1000)
            end_sec   = round(line.end  / 1000)
            segment_texts.append(cleaned)
            sub_timestamps.append((start_sec, end_sec))
        final_text = SEGMENT_DELIMITER.join(segment_texts).strip()
        return final_text, sub_timestamps

    def parse_pdf(self, pdf_path):
        parts = []
        with fitz.open(pdf_path) as doc:
            for page in doc:
                txt = page.get_text("text", sort=True)
                txt = re.sub(r"[ \t]+", " ", txt)
                txt = re.sub(r"\n", "", txt).strip()
                parts.append(txt)
        return "".join(parts)

class File:
    def __init__(
        self,
        learning_lang: str,
        content: str,
        processed_words: Dict[Tuple[str, Optional[str]], Dict[str, Any]] | None = None,
        path=None,
        sub_timestamps: List[Tuple[int, int]] | None = None,
        *,
        segment_delimiter: str = SEGMENT_DELIMITER,
    ):
        self.learning_lang = learning_lang
        self.path = path
        self.content = content
        self.processed_words = processed_words if processed_words else {}
        self.sub_timestamps = sub_timestamps  # text indices -> time on screen

        self.segment_delimiter = segment_delimiter

        self.segments: List[str] = []
        self.segment_records: List[SimpleNamespace] | None = None
        self.lemma_word_index_mapping: Dict[str, Dict[str, List[int]]] = {}
        self.text_segment_components: List[List[Dict[str, Any]]] = []
        self.lemma_counts: Counter = Counter()

    def split_into_segments(self) -> List[str]:
        return [s for s in self.content.split(self.segment_delimiter) if s]

    def extract_lemmas(
        self,
        validator: DictValidator,
        *,
        save_pos_data: bool,
        save_sentence_mapping: bool,
    ) -> None:
        """
        Extract dictionary keys (lemmas) using a DictValidator that yields:
            (surface, lemma_lower, pos)

        Produces:
          - lemma_counts: Counter of lemma_lower
          - processed_words: (surface, pos) -> {'lemma': lemma_lower, 'xpos': pos}
          - lemma_word_index_mapping: lemma_lower -> surface -> [segment_idx] (optional)
          - text_segment_components: per-segment list of dicts with word/lemma/pos (optional)
        """
        self.segments = self.split_into_segments()
        self.segment_records = [SimpleNamespace(text=s) for s in self.segments]

        for segment_idx, segment_text in enumerate(self.segments):
            metadata: List[Dict[str, Any]] = []

            for surface, lemma_lower, pos in validator.iter_lemmas(segment_text):
                if not lemma_lower:
                    continue
                # your original filter — keep if you want it here rather than in validators
                if re.search(r"\d", lemma_lower):
                    continue

                self.lemma_counts[lemma_lower] += 1

                key = (surface, pos)
                if key not in self.processed_words:
                    self.processed_words[key] = {"lemma": lemma_lower, "xpos": pos}

                if save_sentence_mapping:
                    self.lemma_word_index_mapping.setdefault(lemma_lower, {}).setdefault(surface, []).append(segment_idx)

                if save_pos_data:
                    metadata.append({"word": surface, "lemma": lemma_lower, "pos": pos, "xpos": pos})

            if save_pos_data:
                self.text_segment_components.append(metadata)
