# FILE: src/adeu/redline/mapper.py

import re
from copy import deepcopy
from dataclasses import dataclass
from typing import List, Optional, Tuple

import structlog
from docx.document import Document as DocumentObject
from docx.oxml import OxmlElement
from docx.oxml.ns import qn
from docx.table import Table
from docx.text.paragraph import Paragraph
from docx.text.run import Run

from adeu.redline.comments import CommentsManager
from adeu.utils.docx import (
    DocxEvent,
    get_paragraph_prefix,  # ENSURE THIS IMPORT IS PRESENT
    get_run_style_markers,
    get_run_text,
    iter_block_items,
    iter_document_parts,
    iter_paragraph_content,
)

logger = structlog.get_logger(__name__)


@dataclass
class TextSpan:
    start: int
    end: int
    text: str
    run: Optional[Run]
    paragraph: Optional[Paragraph]
    ins_id: Optional[str] = None
    del_id: Optional[str] = None


class DocumentMapper:
    def __init__(self, doc: DocumentObject, clean_view: bool = False):
        self.doc = doc
        self.clean_view = clean_view
        self.comments_mgr = CommentsManager(doc)
        self.comments_map = self.comments_mgr.extract_comments_data()
        self.full_text = ""
        self.spans: List[TextSpan] = []
        self._build_map()

    def _build_map(self):
        current_offset = 0
        self.spans = []
        self.full_text = ""

        for part in iter_document_parts(self.doc):
            current_offset = self._map_blocks(part, current_offset)

            # Add part separator if needed, or rely on block separators
            if self.spans and self.spans[-1].text != "\n\n":
                self._add_virtual_text("\n\n", current_offset, None)
                current_offset += 2

        # Cleanup trailing newlines
        while self.spans and self.spans[-1].text == "\n\n":
            self.spans.pop()
            self.full_text = self.full_text[:-2]

    def _map_blocks(self, container, offset: int) -> int:
        current = offset

        for item in iter_block_items(container):
            if isinstance(item, Paragraph):
                # FIX: Include paragraph prefix (e.g. "# ") in the mapper
                prefix = get_paragraph_prefix(item)
                if prefix:
                    self._add_virtual_text(prefix, current, item)
                    current += len(prefix)

                current = self._map_paragraph_content(item, current)

                # Separator between paragraphs
                self._add_virtual_text("\n\n", current, item)
                current += 2

            elif isinstance(item, Table):
                current = self._map_table(item, current)
                # Separator after table
                if self.spans and self.spans[-1].text != "\n\n":
                    self._add_virtual_text("\n\n", current, None)
                    current += 2

        return current

    def _map_table(self, table: Table, offset: int) -> int:
        current = offset
        rows_processed = 0

        for row in table.rows:
            if rows_processed > 0:
                # Newline separator BETWEEN rows (matches "\n".join in ingest)
                self._add_virtual_text("\n", current, None)
                current += 1

            seen_cells = set()
            cells_processed = 0

            for cell in row.cells:
                if cell in seen_cells:
                    continue
                seen_cells.add(cell)

                if cells_processed > 0:
                    self._add_virtual_text(" | ", current, None)
                    current += 3

                current = self._map_blocks(cell, current)
                cells_processed += 1

            rows_processed += 1

        return current

    def _strip_markdown_formatting(self, text: str) -> str:
        """
        Strips markdown formatting markers from text for matching purposes.
        Handles: **bold**, __bold__, _italic_, *italic*, # headers
        Only strips when content looks like actual formatted text (2+ word chars).
        """
        result = text

        # Strip header markers at start of lines
        result = re.sub(r"^#+\s*", "", result, flags=re.MULTILINE)

        # Strip bold markers - only when wrapping word content (not single chars)
        result = re.sub(r"\*\*(\w[\w\s]*\w|\w{2,})\*\*", r"\1", result)
        result = re.sub(r"__(\w[\w\s]*\w|\w{2,})__", r"\1", result)

        # Strip italic markers - only when wrapping word content
        result = re.sub(r"(?<!\w)_(\w[\w\s]*\w|\w{2,})_(?!\w)", r"\1", result)
        result = re.sub(r"(?<!\w)\*(\w[\w\s]*\w|\w{2,})\*(?!\w)", r"\1", result)

        return result

    def _map_paragraph_content(self, paragraph: Paragraph, start_offset: int) -> int:
        """
        Maps Runs to Spans, handling Flattened CriticMarkup generation.
        """
        current = start_offset

        active_ids: set[str] = set()
        active_ins: dict[str, DocxEvent] = {}
        active_del: dict[str, DocxEvent] = {}

        deferred_meta_states: List[Tuple] = []
        current_wrappers = ("", "")
        pending_runs: List[Tuple[str, str, Optional[Run], Optional[str], Optional[str]]] = []

        items = list(iter_paragraph_content(paragraph))

        for i, item in enumerate(items):
            if isinstance(item, Run):
                prefix, suffix = get_run_style_markers(item)
                run_parts: List[Tuple[str, str, Optional[Run]]] = []

                text = get_run_text(item)

                if "\n" in text and (prefix or suffix):
                    parts = text.split("\n")
                    for idx, part in enumerate(parts):
                        if idx > 0:
                            run_parts.append(("real", "\n", item))
                        if part:
                            if prefix:
                                run_parts.append(("virtual", prefix, None))
                            run_parts.append(("real", part, item))
                            if suffix:
                                run_parts.append(("virtual", suffix, None))
                else:
                    if prefix:
                        run_parts.append(("virtual", prefix, None))
                    if text:
                        run_parts.append(("real", text, item))
                    if suffix:
                        run_parts.append(("virtual", suffix, None))

                if self.clean_view and active_del:
                    pass

                full_seg_text = "".join(x[1] for x in run_parts)

                # Initialize IDs safely before use
                curr_ins_id = list(active_ins.keys())[-1] if active_ins else None
                curr_del_id = list(active_del.keys())[-1] if active_del else None

                if full_seg_text and not (self.clean_view and curr_del_id):
                    if self.clean_view:
                        new_wrappers = ("", "")
                    else:
                        start_token, end_token = self._get_wrappers(curr_ins_id, curr_del_id, active_ids)
                        new_wrappers = (start_token, end_token)

                    if pending_runs and new_wrappers == current_wrappers:
                        for kind, txt, r_obj in run_parts:
                            pending_runs.append((kind, txt, r_obj, curr_ins_id, curr_del_id))
                    else:
                        if pending_runs:
                            s_tok, e_tok = current_wrappers
                            if s_tok:
                                self._add_virtual_text(s_tok, current, paragraph)
                                current += len(s_tok)
                            for kind, txt, r_obj, i_id, d_id in pending_runs:
                                if kind == "virtual":
                                    self._add_virtual_text(txt, current, paragraph)
                                else:
                                    span = TextSpan(
                                        start=current,
                                        end=current + len(txt),
                                        text=txt,
                                        run=r_obj,
                                        paragraph=paragraph,
                                        ins_id=i_id,
                                        del_id=d_id,
                                    )
                                    self.spans.append(span)
                                    self.full_text += txt
                                current += len(txt)
                            if e_tok:
                                self._add_virtual_text(e_tok, current, paragraph)
                                current += len(e_tok)

                        current_wrappers = new_wrappers
                        pending_runs = []
                        for kind, txt, r_obj in run_parts:
                            pending_runs.append((kind, txt, r_obj, curr_ins_id, curr_del_id))

                # Metadata Handling
                if not self.clean_view:
                    state_snapshot = (
                        active_ins.copy(),
                        active_del.copy(),
                        active_ids.copy(),
                    )
                    deferred_meta_states.append(state_snapshot)

                    should_defer = False
                    is_redline = bool(curr_ins_id) or bool(curr_del_id)

                    if is_redline:
                        j = i + 1
                        next_is_redline = False
                        # Use counters to handle nested tag closures safely during lookahead
                        temp_ins_count = len(active_ins)
                        temp_del_count = len(active_del)

                        while j < len(items):
                            next_item = items[j]
                            if isinstance(next_item, Run):
                                if temp_ins_count > 0 or temp_del_count > 0:
                                    next_is_redline = True
                                break
                            elif isinstance(next_item, DocxEvent):
                                if next_item.type == "ins_start":
                                    temp_ins_count += 1
                                elif next_item.type == "ins_end":
                                    temp_ins_count = max(0, temp_ins_count - 1)
                                elif next_item.type == "del_start":
                                    temp_del_count += 1
                                elif next_item.type == "del_end":
                                    temp_del_count = max(0, temp_del_count - 1)
                            j += 1

                        if next_is_redline:
                            should_defer = True

                    if not should_defer:
                        if pending_runs:
                            s_tok, e_tok = current_wrappers
                            if s_tok:
                                self._add_virtual_text(s_tok, current, paragraph)
                                current += len(s_tok)
                            for kind, txt, r_obj, i_id, d_id in pending_runs:
                                if kind == "virtual":
                                    self._add_virtual_text(txt, current, paragraph)
                                else:
                                    span = TextSpan(
                                        start=current,
                                        end=current + len(txt),
                                        text=txt,
                                        run=r_obj,
                                        paragraph=paragraph,
                                        ins_id=i_id,
                                        del_id=d_id,
                                    )
                                    self.spans.append(span)
                                    self.full_text += txt
                                current += len(txt)
                            if e_tok:
                                self._add_virtual_text(e_tok, current, paragraph)
                                current += len(e_tok)
                            pending_runs = []
                            current_wrappers = ("", "")

                        meta_block = self._build_merged_meta_block(deferred_meta_states)
                        if meta_block:
                            full_meta = f"{{>>{meta_block}<<}}"
                            self._add_virtual_text(full_meta, current, paragraph)
                            current += len(full_meta)
                        deferred_meta_states = []

            elif isinstance(item, DocxEvent):
                if pending_runs:
                    s_tok, e_tok = current_wrappers
                    if s_tok:
                        self._add_virtual_text(s_tok, current, paragraph)
                        current += len(s_tok)
                    for kind, txt, r_obj, i_id, d_id in pending_runs:
                        if kind == "virtual":
                            self._add_virtual_text(txt, current, paragraph)
                        else:
                            span = TextSpan(
                                start=current,
                                end=current + len(txt),
                                text=txt,
                                run=r_obj,
                                paragraph=paragraph,
                                ins_id=i_id,
                                del_id=d_id,
                            )
                            self.spans.append(span)
                            self.full_text += txt
                        current += len(txt)
                    if e_tok:
                        self._add_virtual_text(e_tok, current, paragraph)
                        current += len(e_tok)
                    pending_runs = []
                    current_wrappers = ("", "")

                if item.type == "start":
                    active_ids.add(item.id)
                elif item.type == "end":
                    if item.id in active_ids:
                        active_ids.remove(item.id)
                elif item.type == "ins_start":
                    active_ins[item.id] = item
                elif item.type == "ins_end":
                    active_ins.pop(item.id, None)
                elif item.type == "del_start":
                    active_del[item.id] = item
                elif item.type == "del_end":
                    active_del.pop(item.id, None)

        # Final Flush
        if pending_runs:
            s_tok, e_tok = current_wrappers
            if s_tok:
                self._add_virtual_text(s_tok, current, paragraph)
                current += len(s_tok)
            for kind, txt, r_obj, i_id, d_id in pending_runs:
                if kind == "virtual":
                    self._add_virtual_text(txt, current, paragraph)
                else:
                    span = TextSpan(
                        start=current,
                        end=current + len(txt),
                        text=txt,
                        run=r_obj,
                        paragraph=paragraph,
                        ins_id=i_id,
                        del_id=d_id,
                    )
                    self.spans.append(span)
                    self.full_text += txt
                current += len(txt)
            if e_tok:
                self._add_virtual_text(e_tok, current, paragraph)
                current += len(e_tok)

        if deferred_meta_states:
            meta_block = self._build_merged_meta_block(deferred_meta_states)
            if meta_block:
                full_meta = f"{{>>{meta_block}<<}}"
                self._add_virtual_text(full_meta, current, paragraph)
                current += len(full_meta)

        return current

    def _get_wrappers(self, ins_id, del_id, active_ids):
        if del_id:
            return "{--", "--}"
        elif ins_id:
            return "{++", "++}"
        elif active_ids:
            return "{==", "==}"
        return "", ""

    def _build_merged_meta_block(self, states_list) -> str:
        change_lines = []
        comment_lines = []
        seen_sigs = set()

        for ins_map, del_map, comments_set in states_list:
            for map_obj in (ins_map, del_map):
                for uid, meta in map_obj.items():
                    sig = f"Chg:{uid}"
                    if sig not in seen_sigs:
                        auth = meta.author or "Unknown"
                        change_lines.append(f"[{sig}] {auth}")
                        seen_sigs.add(sig)

            sorted_ids = sorted(list(comments_set))
            for c_id in sorted_ids:
                if c_id not in self.comments_map:
                    continue
                sig = f"Com:{c_id}"
                if sig not in seen_sigs:
                    data = self.comments_map[c_id]
                    header = f"[{sig}] {data['author']}"
                    if data["date"]:
                        short_date = data["date"].split("T")[0]
                        header += f" @ {short_date}"
                    if data["resolved"]:
                        header += "(RESOLVED)"
                    comment_lines.append(f"{header}: {data['text']}")
                    seen_sigs.add(sig)

        return "\n".join(change_lines + comment_lines)

    def _add_virtual_text(self, text: str, offset: int, context_paragraph: Optional[Paragraph]):
        span = TextSpan(
            start=offset,
            end=offset + len(text),
            text=text,
            run=None,  # Virtual
            paragraph=context_paragraph,
        )
        self.spans.append(span)
        self.full_text += text

    def _replace_smart_quotes(self, text: str) -> str:
        return text.replace("“", '"').replace("”", '"').replace("‘", "'").replace("’", "'")

    def _make_fuzzy_regex(self, target_text: str) -> str:
        # First strip markdown from the target for cleaner matching
        target_text = self._strip_markdown_formatting(target_text)

        # Normalize quotes in target for consistency
        target_text = self._replace_smart_quotes(target_text)

        parts = []
        # Tokenize: Placeholder brackets with underscores, Whitespace, Quotes, Punctuation
        token_pattern = re.compile(r"(\[_+\])|(\s+)|(['\"])|([.,;:])")

        last_idx = 0
        for match in token_pattern.finditer(target_text):
            # Add literal text
            literal = target_text[last_idx : match.start()]
            if literal:
                escaped = re.escape(literal)
                parts.append(escaped)

            g_placeholder, g_space, g_quote, g_punct = match.groups()

            if g_placeholder:
                # [___] placeholder - allow variable underscore count
                parts.append(r"\[_+\]")
            elif g_space:
                # Allow optional markdown markers around whitespace (between words)
                # This handles cases like "Terms are **Net 90**" matching "Terms are Net 90"
                parts.append(r"(?:\*\*|__|\*|_)?")
                parts.append(r"\s+")
                parts.append(r"(?:\*\*|__|\*|_)?")
            elif g_quote:
                if g_quote == "'":
                    parts.append(r"[''']")
                else:
                    parts.append(r"[\"" "]")
            elif g_punct:
                # Allow optional markdown markers around punctuation
                parts.append(r"(?:\*\*|__|\*|_)?")
                parts.append(re.escape(g_punct))
                parts.append(r"(?:\*\*|__|\*|_)?")

            last_idx = match.end()

        remaining = target_text[last_idx:]
        if remaining:
            parts.append(re.escape(remaining))

        return "".join(parts)

    def find_match_index(self, target_text: str) -> Tuple[int, int]:
        """
        Returns (start_index, match_length).
        Returns (-1, 0) if not found.
        """
        # 1. Exact Match
        start_idx = self.full_text.find(target_text)
        if start_idx != -1:
            return start_idx, len(target_text)

        # 2. Smart Quote Normalization
        norm_full = self._replace_smart_quotes(self.full_text)
        norm_target = self._replace_smart_quotes(target_text)
        start_idx = norm_full.find(norm_target)
        if start_idx != -1:
            return start_idx, len(target_text)

        # 3. Strip markdown from target and try matching (ADDED)
        stripped_target = self._strip_markdown_formatting(target_text)

        # We can't use index from stripped_full directly on full_text,
        # but if it matches, it suggests we should try a fuzzy approach or fallback
        # This fallback is primarily for Header matching (#)
        if stripped_target in self.full_text:
            start_idx = self.full_text.find(stripped_target)
            return start_idx, len(stripped_target)

        # 4. Fuzzy Regex Match
        try:
            pattern = self._make_fuzzy_regex(target_text)
            match = re.search(pattern, self.full_text)
            if match:
                return match.start(), match.end() - match.start()
        except re.error:
            pass

        return -1, 0

    def find_all_match_indices(self, target_text: str) -> List[Tuple[int, int]]:
        """
        Returns a list of all non-overlapping matches as (start_index, match_length).
        Returns an empty list if not found.
        """
        if not target_text:
            return []

        # 1. Exact Match
        matches = [m.span() for m in re.finditer(re.escape(target_text), self.full_text)]
        if matches:
            return [(s, e - s) for s, e in matches]

        # 2. Smart Quote Normalization
        norm_full = self._replace_smart_quotes(self.full_text)
        norm_target = self._replace_smart_quotes(target_text)
        matches = [m.span() for m in re.finditer(re.escape(norm_target), norm_full)]
        if matches:
            return [(s, e - s) for s, e in matches]

        # 3. Strip markdown from target
        stripped_target = self._strip_markdown_formatting(target_text)
        matches = [m.span() for m in re.finditer(re.escape(stripped_target), self.full_text)]
        if matches:
            return [(s, e - s) for s, e in matches]

        # 4. Fuzzy Regex Match
        try:
            pattern = self._make_fuzzy_regex(target_text)
            matches = [m.span() for m in re.finditer(pattern, self.full_text)]
            if matches:
                return [(s, e - s) for s, e in matches]
        except re.error:
            pass

        return []

    def find_target_runs(self, target_text: str) -> List[Run]:
        start_idx, length = self.find_match_index(target_text)
        if start_idx == -1:
            return []
        return self._resolve_runs_at_range(start_idx, start_idx + length)

    def find_target_runs_by_index(self, start_index: int, length: int) -> List[Run]:
        end_index = start_index + length
        return self._resolve_runs_at_range(start_index, end_index)

    def _resolve_runs_at_range(self, start_idx: int, end_idx: int) -> List[Run]:
        affected_spans = [s for s in self.spans if s.end > start_idx and s.start < end_idx]
        if not affected_spans:
            return []

        working_runs = [s.run for s in affected_spans if s.run is not None]
        if not working_runs:
            return []

        dom_modified = False

        # 1. Start Split
        first_real_span = next((s for s in affected_spans if s.run is not None), None)
        start_split_adjustment = 0

        if first_real_span:
            local_start = start_idx - first_real_span.start
            if local_start > 0:
                idx_in_working = 0
                _, right_run = self._split_run_at_index(working_runs[idx_in_working], local_start)
                working_runs[idx_in_working] = right_run
                dom_modified = True
                start_split_adjustment = local_start

        # 2. End Split
        last_real_span = next((s for s in reversed(affected_spans) if s.run is not None), None)

        if last_real_span:
            is_same_run = first_real_span is last_real_span
            run_to_split = working_runs[-1]
            overlap_end = min(last_real_span.end, end_idx)
            local_end = overlap_end - last_real_span.start

            if is_same_run and start_split_adjustment > 0:
                local_end -= start_split_adjustment

            if 0 < local_end < len(run_to_split.text):
                left_run, _ = self._split_run_at_index(run_to_split, local_end)
                working_runs[-1] = left_run
                dom_modified = True

        if dom_modified:
            self._build_map()

        return working_runs

    def get_insertion_anchor(self, index: int) -> Optional[Run]:
        preceding = [s for s in self.spans if s.end == index]
        if preceding:
            if preceding[-1].run:
                return preceding[-1].run
        containing = [s for s in self.spans if s.start < index < s.end]
        if containing:
            span = containing[0]
            if span.run is None:
                pass
            else:
                offset = index - span.start
                left, _ = self._split_run_at_index(span.run, offset)
                return left

        if index == 0 and self.spans:
            for s in self.spans:
                if s.run:
                    return s.run
            return None

        preceding_gap = [s for s in self.spans if s.end < index]
        if preceding_gap:
            for s in reversed(preceding_gap):
                if s.run:
                    return s.run
        return None

    def _split_run_at_index(self, run: Run, split_index: int) -> Tuple[Run, Run]:
        text = run.text
        left_text = text[:split_index]
        right_text = text[split_index:]

        run.text = left_text
        new_r_element = deepcopy(run._element)
        t_list = new_r_element.findall(qn("w:t"))
        for t in t_list:
            new_r_element.remove(t)

        new_t = OxmlElement("w:t")
        new_t.text = right_text
        if right_text.strip() != right_text:
            new_t.set(qn("xml:space"), "preserve")
        new_r_element.append(new_t)
        run._element.addnext(new_r_element)
        new_run = Run(new_r_element, run._parent)
        return run, new_run

    def get_context_at_range(self, start_idx: int, end_idx: int) -> Optional[TextSpan]:
        real_spans = [s for s in self.spans if s.run and s.end > start_idx and s.start < end_idx]
        if real_spans:
            return real_spans[0]
        return None
