import os
import hashlib
import pickle
import re
from pathlib import Path
from lxml import etree
import datetime

from justai.agent.agent import Agent
from justai.tools.log import Log
from justai.tools.prompts import get_prompt, set_prompt_file
from justai.translator.languages import LANGUAGES


def log(message, param):
    def nonl(text):
        return str(text).replace('\n', '<br/>')

    with open('justai.txt', 'a') as f:
        timestamp = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        f.write(f'<p><b>{timestamp} - {message.upper()}</b><br/>')
        if isinstance(param, list):
            for item in param:
                f.write('-' + nonl(item) + '<br/>')
            f.write(f'{len(param)}<br/>')
        elif isinstance(param, dict):
            for key, value in param.items():
                f.write(f'{key}: {nonl(value)}<br/>')
        else:
            f.write(nonl(param))
        f.write('</p>')


class Translator(Agent):

    def __init__(self, model=None):
        if not model:
            model = os.environ.get('MODEL', 'claude-3-opus-20240229')
        super().__init__(model, temperature=0, max_tokens=4096)
        set_prompt_file(Path(__file__).parent / 'prompts.toml')
        self.system_message = get_prompt('SYSTEM')
        self.xml = ''
        self.version = ''
        self.logger = Log()

    def load(self, input_file: str | Path):
        with open(input_file, 'r') as f:
            self.read(f.read())

    def read(self, input_string: str):
        # Input bestaat uit <transunit> elementen. Die hebben een datatype property.
        # Binnen elke <transunit> zit een <source> element en komt (na vertaling) een <target> element.
        # ALs datatype == "plaintext" dan zit de te vertalen tekst direct in de <source>
        # Als datatype == "x-DocumentState" dan zit er in de <source> een <g> element met daarin de te vertalen tekst.

        # In 2.0:
        # Input bestaat uit <unit> elementen. Die hebben een Id.
        # Binnen elke <unit> zit een <segment> en daarin een <source>
        # In de source zit ofwel direct tekst, ofwel een <pc> element
        # met daarin nog een <pc> element met daarin de te vertalen tekst
        self.xml = input_string
        self.messages = []
        try:
            self.version = self.xml.split('xliff:document:')[1].split('"')[0].split("'")[0]
        except IndexError:
            raise ValueError('No XLIFF version found in input')
        if self.version not in ['1.2', '2.0']:
            raise ValueError(f'Unsupported XLIFF version: {self.version}')

    def translate(self, language: str, string_cached: bool = False) -> str:
        if self.version == '1.2':
            return self.translate1_2(language, string_cached=string_cached)
        elif self.version == '2.0':
            return self.translate2_01(language, string_cached=string_cached)

    def translate1_2(self, language, string_cached: bool = False):
        # XML-data laden met lxml
        parser = etree.XMLParser(ns_clean=True)
        root = etree.fromstring(self.xml.encode('utf-8'), parser=parser)
        namespaces = {'ns': 'urn:oasis:names:tc:xliff:document:1.2'}

        # Verzamel alle te vertalen teksten en hun paden
        texts_to_translate = []

        # Start het verzamelproces vanuit <source> elementen en vertaal de teksten
        for trans_unit in root.xpath('.//ns:trans-unit', namespaces=namespaces):
            source = trans_unit.xpath('.//ns:source', namespaces=namespaces)[0]
            texts_to_translate.extend(collect_texts_from_element(source))

        # Vertaal met AI
        translated_texts = self.do_translate(texts_to_translate, language, string_cached=string_cached)

        # Plaats vertaalde teksten terug in nieuwe <target> elementen met behoud van structuur
        counter = [0]
        for trans_unit in root.xpath('.//ns:trans-unit', namespaces=namespaces):
            source = trans_unit.xpath('.//ns:source', namespaces=namespaces)[0]
            target = etree.Element('{urn:oasis:names:tc:xliff:document:1.2}target')
            copy_structure_with_texts(source, target, translated_texts, counter)
            trans_unit.append(target)

        # De bijgewerkte XLIFF-structuur omzetten naar een string en afdrukken
        updated_xml = etree.tostring(root, pretty_print=True, xml_declaration=True, encoding='UTF-8').decode('utf-8')
        return updated_xml

    def translate2_0(self, language, string_cached: bool = False):
        # return self.experiment_with_translating_xml_source_blocks(language, string_cached)

        # XML-data laden met lxml
        parser = etree.XMLParser(ns_clean=True)
        root = etree.fromstring(self.xml.encode('utf-8'), parser=parser)
        namespaces = {'ns': 'urn:oasis:names:tc:xliff:document:2.0'}

        # Speciaal voor xliff 2.0: voeg de target language toe aan het root element
        language_code = LANGUAGES.get(language)
        root.attrib['trgLang'] = language_code

        # Verzamel alle te vertalen teksten en hun paden
        texts_to_translate = []

        # Start het verzamelproces vanuit <source> elementen en vertaal de teksten
        for source in root.xpath('.//ns:source', namespaces=namespaces):
            texts_to_translate.extend(collect_texts_from_element(source))

        # Vertaal met AI
        translated_texts = self.do_translate(texts_to_translate, language, string_cached=string_cached)

        # Plaats vertaalde teksten terug in nieuwe <target> elementen met behoud van structuur
        counter = [0]
        for segment in root.xpath('.//ns:segment', namespaces=namespaces):
            source = segment.xpath('.//ns:source', namespaces=namespaces)[0]
            target = etree.SubElement(segment, '{urn:oasis:names:tc:xliff:document:2.0}target')
            copy_structure_with_texts(source, target, translated_texts, counter)

        updated_xml = etree.tostring(root, pretty_print=True, xml_declaration=True, encoding='UTF-8').decode('utf-8')
        return updated_xml

    def translate2_01(self, language, string_cached: bool = False):
        log = self.logger
        # Versie van 2_0 die teksten uit een <source> samen aan de AI aanbiedt ter vertaling.
        parser = etree.XMLParser(ns_clean=True)
        root = etree.fromstring(self.xml.encode('utf-8'), parser=parser)
        namespaces = {'ns': 'urn:oasis:names:tc:xliff:document:2.0'}

        language_code = LANGUAGES.get(language)
        root.attrib['trgLang'] = language_code

        # Verzamel teksten als lijst van samengevoegde strings per <source>
        all_texts = []
        translatable_texts = []
        for source in root.xpath('.//ns:source', namespaces=namespaces):
            texts = collect_texts_from_element(source)
            all_texts.append(texts)
            # Verzamel alleen de teksten die ook echt vertaald moeten worden
            translatable_text = "||".join(text for text in texts if is_translatable(text))
            if translatable_text:
                translatable_texts.append(translatable_text)
        log.info('all_texts', all_texts)
        log.info('translatable_texts', translatable_texts)

        # Vertaal de lijst van samengevoegde strings
        translated_texts = self.do_translate(translatable_texts, language, string_cached=string_cached)
        original_translated_texts_list = translated_texts[:]
        log.info('translated_texts', translated_texts)

        # Zet nu de delen die niet vertaald hoefden te worden terug in de lijst
        for i1, line in enumerate(all_texts):
            if any(is_translatable(text) for text in line):
                if len([l for l in line if is_translatable(l)]) != len(translated_texts[0].split("||")):
                    log.error('mismatch', f'Number of translated texts ({len(translated_texts[0].split("||"))}) '
                          f'does not match number of source texts ({len([l for l in line if is_translatable(l)])})')
                    log.info('line', line)
                    log.info('original:', [l for l in line if is_translatable(l)])
                    log.info('Translated:', translated_texts[0].split("||"))

                assert len([l for l in line if is_translatable(l)]) == len(translated_texts[0].split("||"))
                translated = translated_texts.pop(0).split("||")
                for index, text in enumerate(line):
                    if is_translatable(text):
                        line[index] = translated.pop(0)
        translated_texts = [item for sublist in all_texts for item in sublist]  # And flatten

        # Plaats vertaalde teksten terug in nieuwe <target> elementen met behoud van structuur
        counter = [0]
        for segment in root.xpath('.//ns:segment', namespaces=namespaces):
            source = segment.xpath('.//ns:source', namespaces=namespaces)[0]
            target = etree.SubElement(segment, '{urn:oasis:names:tc:xliff:document:2.0}target')
            copy_structure_with_texts(source, target, translated_texts, counter)

        updated_xml = etree.tostring(root, pretty_print=True, xml_declaration=True, encoding='UTF-8').decode('utf-8')
        return updated_xml

    def do_translate(self, texts, language: str, string_cached=False):
        # source_list = list(set([text for text in texts if is_translatable(text)]))  # Filter out doubles
        assert all(is_translatable(text) for text in texts)  # !! Tijdelijk
        source_list = list(set(texts))
        cache = StringCache(language) if string_cached else {}
        source_list = [text for text in source_list if text not in cache]

        if source_list:
            vars = []
            source_str_no_vars = ''
            for index, text in enumerate(source_list):
                text_no_vars, v = replace_variables_with_hash(text)
                vars.extend(v)
                source_str_no_vars += f'{index + 1} [[{text_no_vars}]]\n'
            prompt = get_prompt('TRANSLATE_MULTIPLE', language=language, translate_str=source_str_no_vars,
                                count=len(source_list))
            log.prompt('prompt', prompt)
            target_str_no_vars = self.chat(prompt, return_json=False)
            log.response('target_str_no_vars', target_str_no_vars)
            target_str = replace_hash_with_variables(target_str_no_vars, vars)
            log.info('target_str', target_str)

            target_list = [t.split(']]')[0] for t in target_str.split('[[')[1:]]
            translation_dict = dict(zip(source_list, target_list))
            log.info('translation_dict', translation_dict)

            count = 1
            for source, translation in translation_dict.items():

                source_parts = source.split('||')
                translation_parts = translation.split('||')
                log.info('source_parts', source_parts)
                log.info('translation_parts', translation_parts)

                # Check op het aantal onderdelen tussen ||. Dit moet gelijk zijn in de bron en de vertaling
                sc = len(source_parts)
                tc = len(translation_parts)
                if tc != sc:
                    log.error('do_translate', f'Number of translated texts ({tc}) does not match number of source texts ({sc}).')
                    if tc < sc:
                        translation_parts += [' '] * (sc - tc)
                    else:
                        # Dit is een hack! Het model geeft kennelijk meer ||'s terug in dan in het origineel.
                        # Dat breekt de code verderop. Daarom voegen we de laatste onderdelen samen.
                        # Maar dat levert wel een onjuiste vertaling op.
                        translation_parts = translation_parts[:sc - 1] + [' '.join(translation_parts[sc - 1:])]
                assert len(source_parts) == len(translation_parts)

                # Zorg dat de vertaling dezelfde whitespace aan het begin en eind heeft heeft als de bron
                for i, (source_part, translation_part) in enumerate(zip(source_parts, translation_parts)):
                    if source_part.strip() and (source_part[0] == ' ' or source_part[-1] == ' '):
                        start_spaces = (len(source_part) - len(source_part.lstrip(' '))) * ' '
                        end_spaces = (len(source_part) - len(source_part.rstrip(' '))) * ' '
                        translation_parts[i] = start_spaces + translation_part.strip() + end_spaces
                translation = '||'.join(translation_parts)

                log.info('', f'{count}. [{source}] -> [{translation}]')
                count += 1
                ratio = len(source) / len(translation)
                if ratio >= 1.5 or ratio <= 0.7:
                    log.warning('', f'Vertaling van {source} naar {translation} is onverwacht lang of kort')

            cache.update(translation_dict)
            if string_cached:
                cache.save()

        translations = [cache.get(text, text) for text in texts]

        return translations


def replace_variables_with_hash(text):
    # Vindt alle variabelen in de tekst
    variables = re.findall(r'%[^%]+%', text)
    # Vervang alle variabelen in de tekst met ###
    # Het model heeft moeite met newlines. Daarom vervangen we ze door @@ en na vertaling weer terug.
    modified_text = re.sub(r'%[^%]+%', '###', text).replace('\n', '@@')
    return modified_text, variables


def replace_hash_with_variables(text, variables):
    for variable in variables:
        text = text.replace('###', variable, 1)
    # en zet de newlines terug
    text = text.replace('@@', '\n')
    return text


def collect_texts_from_element(element):
    texts = []
    # if element.text and element.text.strip():
    #    texts.append(element.text.strip())
    if element.text:
        texts.append(element.text)
    for child in element:
        texts.extend(collect_texts_from_element(child))
    return texts


def copy_structure_with_texts(source, target, translated_texts, counter=[0]):
    """ Kopieer de structuur van <source> naar <target> en behoud de teksten """
    if source.text:  # and source.text.strip():
        try:
            target.text = translated_texts[counter[0]]
            counter[0] += 1
        except IndexError:
            log.error('IndexError', 'IndexError in copy_structure_with_texts')
    for child in source:
        child_copy = etree.SubElement(target, child.tag, attrib=child.attrib)
        copy_structure_with_texts(child, child_copy, translated_texts, counter)


def is_translatable(text) -> bool:
    """ Returns True if the unit should be translated """
    return text and re.search('[a-zA-Z]{2}', text) and text[0] not in ('%', '<')


def split_list_in_sublists(source_list, max_chunk_len):
    chunks = []
    for text in source_list:
        if not chunks or chunks[-1] and len(chunks[-1]) + len(text) > max_chunk_len:
            chunks.append([text])
        else:
            chunks[-1].append(text)
    return chunks


class StringCache:
    def __init__(self, language: str):
        self.language = language
        self.cache = {}
        self.file = Path(__file__).parent / (self.language + '.pickle')
        try:
            with open(self.file, 'rb') as f:
                self.cache = pickle.load(f)
        except (FileNotFoundError, EOFError):
            self.cache = {}

    def get(self, source, default=None):
        key = self.get_key(source)
        return self.cache.get(key, default)

    def set(self, source, translation):
        key = self.get_key(source)
        self.cache[key] = translation

    def __contains__(self, source):
        key = self.get_key(source)
        return key in self.cache

    def update(self, translation_dict):
        for source, translation in translation_dict.items():
            self.set(source, translation)

    def save(self):
        with open(self.file, 'wb') as f:
            pickle.dump(self.cache, f)

    def clear(self):
        self.cache = {}
        self.save()

    @classmethod
    def get_key(cls, source):
        return hashlib.md5(source.encode('utf-8')).hexdigest()


def parse_xliff_with_unit_clusters(xliff_content, max_chunk_size):
    # Functie alleen gebruikt voor tests. Deze is gelijk aan de parseXLIFFWithUnitClusters uit javascript
    # en is bedoeld om de real life situatie met het splitsen van de XLIFF in clusters te simuleren.

    # Bepaal de versie van xliffContent
    version_match = re.search(r'<xliff[^>]*\s+version="([0-9.]+)"', xliff_content)
    version = version_match.group(1) if version_match else None

    # Splits xliff in header, clusters van maxChunkSize, and footer
    header_re = r'^(.*?)<unit ' if version == '2.0' else r'^(.*?)<trans-unit '
    header_match = re.match(header_re, xliff_content, re.DOTALL)
    header = header_match.group(1) if header_match else None

    # Extract units and cluster them
    units = []
    cluster = ''
    unit_re = r'<unit .*?</unit>' if version == '2.0' else r'<trans-unit .*?</trans-unit>'
    matches = re.findall(unit_re, xliff_content, re.DOTALL)

    for match in matches:
        unit = match
        if len(unit) > max_chunk_size:
            # If current unit is larger than maxChunkSize, push current cluster (if not empty) and then this unit
            if cluster:
                units.append(cluster)
                cluster = ''
            units.append(unit)
        elif len(cluster) + len(unit) > max_chunk_size:
            # If adding this unit exceeds the limit, push current cluster and start a new one
            units.append(cluster)
            cluster = unit
        else:
            # Add this unit to current cluster
            cluster += unit

    # Don't forget to add the last cluster if it exists
    if cluster:
        units.append(cluster)

    # Extract footer
    footer_re = r'</unit>' if version == '2.0' else r'</trans-unit>'
    footer = xliff_content.split(footer_re)[-1]

    return {"header": header, "units": units, "footer": footer, "version": version}
