import re
from datetime import datetime, timezone
from lxml import etree


def _esc_xml(s):
    return (str(s)
            .replace('&', '&amp;')
            .replace('<', '&lt;')
            .replace('>', '&gt;')
            .replace('"', '&quot;'))


def _esc_xml_content(s):
    return (str(s)
            .replace('&', '&amp;')
            .replace('<', '&lt;')
            .replace('>', '&gt;'))


def _build_target_map(segments):
    return {str(s['mid']): s['target']
            for s in segments if s.get('target')}


# ── SDLXLIFF — raw string surgery (no XML parser) ────────────────────────────


def _update_sdl_seg_conf(xml, mids):
    """Set conf="Draft" on every <sdl:seg id="N"> whose id is in *mids*."""
    if not mids:
        return xml
    seg_re = re.compile(r'<sdl:seg\b([^>]*)(/?)>')

    def _replace(m):
        attrs, self_close = m.group(1), m.group(2)
        id_m = re.search(r'\bid\s*=\s*"([^"]*)"', attrs)
        if not id_m or id_m.group(1) not in mids:
            return m.group(0)
        if re.search(r'\bconf\s*=\s*"[^"]*"', attrs):
            attrs = re.sub(r'\bconf\s*=\s*"[^"]*"', 'conf="Draft"', attrs)
        else:
            attrs += ' conf="Draft"'
        return '<sdl:seg' + attrs + self_close + '>'

    return seg_re.sub(_replace, xml)


def generate_sdlxliff(raw_xml, segments):
    target_map = _build_target_map(segments)
    if not target_map:
        return raw_xml

    # Locate all <target>...</target> ranges
    target_ranges = []
    target_open_re = re.compile(r'<target[\s>]', re.IGNORECASE)
    for m in target_open_re.finditer(raw_xml):
        open_start = m.start()
        tag_end = raw_xml.find('>', open_start)
        if tag_end != -1 and raw_xml[tag_end - 1] == '/':
            continue
        close_idx = raw_xml.find('</target>', open_start)
        if close_idx == -1:
            continue
        target_ranges.append((open_start, close_idx + len('</target>')))

    mrk_open_close_re = re.compile(
        r'<mrk\b[^>]*\bmtype\s*=\s*"seg"[^>]*(?<!/)>([\s\S]*?)</mrk>', re.DOTALL)
    mrk_self_close_re = re.compile(
        r'<mrk\b[^>]*\bmtype\s*=\s*"seg"[^>]*/>')
    mid_re = re.compile(r'\bmid\s*=\s*"([^"]*)"')

    replacements = []
    translated_mids = set()

    for t_start, _t_end in target_ranges:
        target_block = raw_xml[t_start:_t_end]

        # Open-close <mrk>content</mrk>
        for mrk_match in mrk_open_close_re.finditer(target_block):
            mid_match = mid_re.search(mrk_match.group(0))
            if not mid_match:
                continue
            mid = mid_match.group(1)
            if mid not in target_map or not target_map[mid]:
                continue

            old_content = mrk_match.group(1)
            new_content = _esc_xml_content(target_map[mid])
            if old_content == new_content:
                continue

            tag_open_end = mrk_match.group(0).index('>') + 1
            content_abs_start = t_start + mrk_match.start() + tag_open_end
            content_abs_end = t_start + mrk_match.start() + mrk_match.group(0).rfind('</mrk>')
            replacements.append((content_abs_start, content_abs_end, new_content))
            translated_mids.add(mid)

        # Self-closing <mrk ... />
        for mrk_match in mrk_self_close_re.finditer(target_block):
            mid_match = mid_re.search(mrk_match.group(0))
            if not mid_match:
                continue
            mid = mid_match.group(1)
            if mid not in target_map or not target_map[mid]:
                continue

            full = mrk_match.group(0)
            attrs = full[4:-2].rstrip()
            expanded = '<mrk' + attrs + '>' + _esc_xml_content(target_map[mid]) + '</mrk>'
            abs_start = t_start + mrk_match.start()
            abs_end = abs_start + len(full)
            replacements.append((abs_start, abs_end, expanded))
            translated_mids.add(mid)

    # Apply replacements from end to start to preserve offsets
    replacements.sort(key=lambda r: r[0], reverse=True)
    result = raw_xml
    for start, end, content in replacements:
        result = result[:start] + content + result[end:]

    # Post-pass: set conf="Draft" on <sdl:seg> elements for every injected mid
    result = _update_sdl_seg_conf(result, translated_mids)

    return result


# ── MQXLIFF — raw string surgery (no XML parser) ───────────────────────────


def _update_mq_status(xml, ids):
    """Set mq:status="PreTranslated" on every <trans-unit id="N"> whose id is in *ids*."""
    if not ids:
        return xml
    tu_re = re.compile(r'<trans-unit\b([^>]*)>')

    def _replace(m):
        attrs = m.group(1)
        id_m = re.search(r'\bid\s*=\s*"([^"]*)"', attrs)
        if not id_m or id_m.group(1) not in ids:
            return m.group(0)
        if re.search(r'\bmq:status\s*=\s*"[^"]*"', attrs):
            attrs = re.sub(r'\bmq:status\s*=\s*"[^"]*"', 'mq:status="PreTranslated"', attrs)
        else:
            attrs += ' mq:status="PreTranslated"'
        return '<trans-unit' + attrs + '>'

    return tu_re.sub(_replace, xml)


def generate_mqxliff(raw_xml, segments):
    target_map = _build_target_map(segments)
    if not target_map:
        return raw_xml

    tu_open_re = re.compile(r'<trans-unit\b[^>]*\bid\s*=\s*"([^"]*)"[^>]*>')
    tgt_open_close_re = re.compile(
        r'<target\b([^>]*)(?<!/)>(.*?)</target>', re.DOTALL)
    tgt_self_close_re = re.compile(r'<target\b([^>]*)/>')

    replacements = []
    translated_ids = set()

    for tu_match in tu_open_re.finditer(raw_xml):
        tu_id = tu_match.group(1)
        if tu_id not in target_map:
            continue

        search_start = tu_match.end()
        new_content = _esc_xml_content(target_map[tu_id])

        # Try open-close <target ...>...</target> first
        tgt_match = tgt_open_close_re.search(raw_xml, search_start)
        if tgt_match:
            replacements.append((tgt_match.start(2), tgt_match.end(2), new_content))
            translated_ids.add(tu_id)
            continue

        # Fallback: self-closing <target ... />
        tgt_sc = tgt_self_close_re.search(raw_xml, search_start)
        if tgt_sc:
            attrs = tgt_sc.group(1)
            if not attrs or attrs.isspace():
                attrs = ' xml:space="preserve"'
            expanded = '<target' + attrs + '>' + new_content + '</target>'
            replacements.append((tgt_sc.start(), tgt_sc.end(), expanded))
            translated_ids.add(tu_id)

    # Apply replacements from end to start to preserve offsets
    replacements.sort(key=lambda r: r[0], reverse=True)
    result = raw_xml
    for start, end, content in replacements:
        result = result[:start] + content + result[end:]

    # Post-pass: set mq:status="PreTranslated" on injected trans-units
    result = _update_mq_status(result, translated_ids)

    return result


# ── TMX — string construction ────────────────────────────────────────────────

def generate_tmx(segments, source_lang, target_lang):
    now = datetime.now(timezone.utc).strftime('%Y%m%dT%H%M%S') + 'Z'
    tus = []
    for s in segments:
        if not s.get('target'):
            continue
        tus.append(
            f'  <tu tuid="{_esc_xml(s["mid"])}">\n'
            f'    <tuv xml:lang="{source_lang}"><seg>{_esc_xml(s["source"])}</seg></tuv>\n'
            f'    <tuv xml:lang="{target_lang}"><seg>{_esc_xml(s["target"])}</seg></tuv>\n'
            f'  </tu>'
        )
    tus_str = '\n'.join(tus)
    return (
        '<?xml version="1.0" encoding="UTF-8"?>\n'
        '<!DOCTYPE tmx SYSTEM "tmx14.dtd">\n'
        '<tmx version="1.4">\n'
        f'  <header creationtool="I\'M TranslAID" creationtoolversion="1.0"\n'
        f'    datatype="plaintext" segtype="sentence"\n'
        f'    adminlang="en-US" srclang="{source_lang}" creationdate="{now}"/>\n'
        f'  <body>\n{tus_str}\n  </body>\n'
        '</tmx>'
    )


# ── Standalone XLIFF 1.2 — string construction ───────────────────────────────

def _js_iso_date():
    dt = datetime.now(timezone.utc)
    return dt.strftime('%Y-%m-%dT%H:%M:%S.') + f'{dt.microsecond // 1000:03d}Z'


def generate_standalone_xliff12(segments, source_lang, target_lang):
    now = _js_iso_date()
    tus = []
    for s in segments:
        mid = _esc_xml(str(s['mid']))
        src = _esc_xml(s.get('source', ''))
        tgt = _esc_xml(s.get('target', ''))
        state_attr = ' state="translated"' if s.get('target') else ''
        tus.append(
            f'    <trans-unit id="{mid}">\n'
            f'      <source xml:lang="{source_lang}">{src}</source>\n'
            f'      <target xml:lang="{target_lang}"{state_attr}>{tgt}</target>\n'
            f'    </trans-unit>'
        )
    tus_str = '\n'.join(tus)
    return (
        '<?xml version="1.0" encoding="UTF-8"?>\n'
        '<xliff version="1.2" xmlns="urn:oasis:names:tc:xliff:document:1.2">\n'
        f'  <file source-language="{source_lang}" target-language="{target_lang}"\n'
        f'    datatype="plaintext" original="imlaid_export" date="{now}">\n'
        f'    <body>\n{tus_str}\n    </body>\n'
        '  </file>\n'
        '</xliff>'
    )


# ── Standalone XLIFF 2.0 — string construction ───────────────────────────────

def generate_standalone_xliff20(segments, source_lang, target_lang):
    units = []
    for s in segments:
        mid = _esc_xml(str(s['mid']))
        src = _esc_xml(s.get('source', ''))
        tgt = _esc_xml(s.get('target', ''))
        state = 'translated' if s.get('target') else 'initial'
        units.append(
            f'  <unit id="{mid}">\n'
            f'    <segment state="{state}">\n'
            f'      <source>{src}</source>\n'
            f'      <target>{tgt}</target>\n'
            f'    </segment>\n'
            f'  </unit>'
        )
    units_str = '\n'.join(units)
    return (
        '<?xml version="1.0" encoding="UTF-8"?>\n'
        '<xliff xmlns="urn:oasis:names:tc:xliff:document:2.0"\n'
        f'  version="2.0" srcLang="{source_lang}" trgLang="{target_lang}">\n'
        f'  <file id="f1" original="imlaid_export">\n{units_str}\n'
        '  </file>\n'
        '</xliff>'
    )
