from docx import Document
from src.detectors.hybrid_detector import HybridDetector
from src.redactors.mapper import StatefulMapper

class DocxHandler:
    """
    Handles reading, redacting, and saving DOCX files safely without destroying
    document structure, tables, headers, or footers.
    """
    def __init__(self, detector: HybridDetector, mapper: StatefulMapper):
        self.detector = detector
        self.mapper = mapper

    def process_document(self, input_path: str, output_path: str):
        """
        Loads a document, redacts PII across all sections, and saves it.
        """
        doc = Document(input_path)

        # 1. Process Core Paragraphs
        self._process_paragraphs(doc.paragraphs)

        # 2. Process Tables
        for table in doc.tables:
            for row_idx, row in enumerate(table.rows):
                for cell_idx, cell in enumerate(row.cells):
                    is_address_cell = False
                    if row_idx > 0:
                        try:
                            above_cell = table.cell(row_idx - 1, cell_idx)
                            if above_cell.text.strip().upper() in ["REGISTERED OFFICE", "CORPORATE OFFICE"]:
                                is_address_cell = True
                        except Exception:
                            pass
                            
                    self._process_paragraphs(cell.paragraphs, is_address_cell=is_address_cell)

        # 3. Process Headers and Footers
        for section in doc.sections:
            self._process_paragraphs(section.header.paragraphs)
            self._process_paragraphs(section.footer.paragraphs)

        doc.save(output_path)
        print(f"Redacted document successfully saved to {output_path}")

    def _process_paragraphs(self, paragraphs, is_address_cell=False):
        """
        Iterates over a list of paragraphs and redacts PII within them.
        """
        for paragraph in paragraphs:
            runs_text = "".join(r.text for r in paragraph.runs)
            if not runs_text.strip():
                continue

            if is_address_cell:
                entities = [{"entity_type": "ADDRESS", "start": 0, "end": len(runs_text), "text": runs_text}]
            else:
                # Detect entities in the run-reconstructed paragraph text
                entities = self.detector.analyze(runs_text)
            
            if not entities:
                continue

            # We process entities in reverse order (right to left) so that modifying
            # the text of a run does not invalidate the start/end character indices 
            # of entities that appear earlier in the paragraph.
            entities.sort(key=lambda x: x['start'], reverse=True)
            
            self._redact_runs(paragraph, entities)

    def _redact_runs(self, paragraph, entities):
        """
        Replaces detected PII across Word 'Runs' safely.
        A single paragraph is composed of multiple runs (sections of text with consistent formatting).
        An entity might span across multiple runs, so we must map character indices to runs.
        """
        for entity in entities:
            e_start = entity['start']
            e_end = entity['end']
            e_type = entity['entity_type']
            e_text = entity['text']

            fake_text = self.mapper.get_replacement(e_type, e_text)

            # Re-calculate run indices dynamically just to be safe, 
            # since we are modifying runs in-place.
            # Because we process right-to-left, the indices to the left remain intact.
            current_len = 0
            runs_info = []
            for idx, run in enumerate(paragraph.runs):
                runs_info.append({
                    "index": idx,
                    "run": run,
                    "start": current_len,
                    "end": current_len + len(run.text),
                    "text": run.text
                })
                current_len += len(run.text)

            # Find which runs the entity overlaps with
            overlapping_runs = []
            for r_info in runs_info:
                # Overlap condition:
                # Entity starts before run ends AND Entity ends after run starts
                if e_start < r_info['end'] and e_end > r_info['start']:
                    overlapping_runs.append(r_info)

            if not overlapping_runs:
                continue

            # If the entity is entirely contained in one run
            if len(overlapping_runs) == 1:
                r_info = overlapping_runs[0]
                run = r_info['run']
                local_start = e_start - r_info['start']
                local_end = e_end - r_info['start']
                
                # Replace the text inside the run
                run.text = run.text[:local_start] + fake_text + run.text[local_end:]
            
            # If the entity spans multiple runs
            else:
                first_r = overlapping_runs[0]
                last_r = overlapping_runs[-1]
                
                local_start = e_start - first_r['start']
                local_end = e_end - last_r['start']

                # Place the replacement text in the first run, keep the text before the entity
                first_r['run'].text = first_r['run'].text[:local_start] + fake_text
                
                # Clear all intermediate runs completely
                for r_info in overlapping_runs[1:-1]:
                    r_info['run'].text = ""
                
                # Keep the text after the entity in the last run
                last_r['run'].text = last_r['run'].text[local_end:]