import re
import json
import logging
from typing import List, Dict
from enum import Enum

logger = logging.getLogger("MathParser")

class MathMode(Enum):
    INLINE = "inline"
    DISPLAY = "display"

class MathSegmenter:
    @staticmethod
    def segment_text_with_math(text: str) -> List[Dict]:
        """
        Convert natural text with $math$ to structured segments.
        Combines V68 protections with V70 logic.
        """
        # 1. Memory Protection (V68 Feature)
        if not text:
            return [{"type": "text", "value": "", "direction": "rtl"}]
        
        if len(text) > 50000:
             logger.warning(f"⚠️ Input text too long ({len(text)}), truncating.")
             text = text[:50000] + "... (truncated)"

        if '$' not in text:
            return [{"type": "text", "value": text.strip(), "direction": "rtl"}]
        
        segments = []
        pos = 0
        text_len = len(text)
        max_iterations = 5000 # Safety limit
        
        iteration = 0
        while pos < text_len and iteration < max_iterations:
            iteration += 1
            dollar_idx = text.find('$', pos)
            
            if dollar_idx == -1:
                remaining = text[pos:]
                if remaining.strip():
                    segments.append(MathSegmenter._create_text_segment(remaining))
                break
            
            is_double = (dollar_idx + 1 < text_len and text[dollar_idx + 1] == '$')
            
            if dollar_idx > pos:
                text_before = text[pos:dollar_idx]
                if text_before.strip():
                    segments.append(MathSegmenter._create_text_segment(text_before))
            
            if is_double:
                close_idx = text.find('$$', dollar_idx + 2)
                if close_idx == -1:
                    segments.append(MathSegmenter._create_text_segment('$$'))
                    pos = dollar_idx + 2
                    continue
                math_content = text[dollar_idx + 2:close_idx].strip()
                if math_content:
                    segments.append({"type": "math", "mode": "display", "value": math_content, "direction": "ltr"})
                pos = close_idx + 2
            else:
                close_idx = text.find('$', dollar_idx + 1)
                if close_idx == -1:
                    segments.append(MathSegmenter._create_text_segment('$'))
                    pos = dollar_idx + 1
                    continue
                math_content = text[dollar_idx + 1:close_idx].strip()
                if math_content:
                    segments.append({"type": "math", "mode": "inline", "value": math_content, "direction": "ltr"})
                pos = close_idx + 1
        
        if not segments:
            return [{"type": "text", "value": text, "direction": "rtl"}]
        
        return MathSegmenter._merge_text_segments(segments)
    
    @staticmethod
    def _create_text_segment(text: str) -> Dict:
        cleaned = text
        # Safe RLM injection (V70 Logic)
        RLM = chr(0x200f)
        if cleaned.strip() and MathSegmenter._ends_with_hebrew(cleaned.strip()):
            cleaned = cleaned.rstrip() + RLM
        return {"type": "text", "value": cleaned, "direction": "rtl"}
    
    @staticmethod
    def _ends_with_hebrew(text: str) -> bool:
        if not text: return False
        last_char = text[-1]
        return '\u0590' <= last_char <= '\u05FF'
    
    @staticmethod
    def _merge_text_segments(segments: List[Dict]) -> List[Dict]:
        if not segments: return []
        merged = []
        curr_text = ""
        for seg in segments:
            if seg["type"] == "text":
                curr_text += seg["value"]
            else:
                if curr_text:
                    merged.append(MathSegmenter._create_text_segment(curr_text))
                    curr_text = ""
                merged.append(seg)
        if curr_text:
            merged.append(MathSegmenter._create_text_segment(curr_text))
        return merged

    @staticmethod
    def segments_to_safe_string(segments: List[Dict]) -> str:
        """
        Reconstructs string safely.
        Uses ONLY RLM (V70 Safe Logic) - NO LRE/PDF to avoid stream crashes.
        """
        result_parts = []
        RLM = chr(0x200f)
        
        for i, segment in enumerate(segments):
            if segment["type"] == "text":
                text = segment["value"]
                # Ensure Hebrew ends with RLM
                if text.strip() and MathSegmenter._ends_with_hebrew(text.strip()):
                     if not text.endswith(RLM):
                        text += RLM
                result_parts.append(text)
            
            elif segment["type"] == "math":
                math_val = segment["value"]
                
                # Spacing logic
                prefix = ""
                suffix = ""
                if i > 0 and segments[i-1]["type"] == "text":
                    if not segments[i-1]["value"].endswith(" ") and not segments[i-1]["value"].endswith(RLM):
                         prefix = " "
                if i < len(segments)-1 and segments[i+1]["type"] == "text":
                    if not segments[i+1]["value"].startswith(" "):
                        suffix = " "

                # Simple wrap: Just Math + RLM.
                if segment.get("mode") == "display":
                    formatted_math = f"$${math_val}$${RLM}"
                else:
                    formatted_math = f"${math_val}${RLM}"
                
                result_parts.append(prefix + formatted_math + suffix)
                
        return "".join(result_parts)

    @staticmethod
    def aggregate_math(segments: List[Dict]) -> str:
        """
        Collects all math for the gray box.
        Uses simple keywords (V70 Safe Logic) to avoid Regex OOM.
        """
        math_blocks = []
        keywords = ['=', '\\frac', '\\sqrt', '\\cdot', 'A', 'B', 'C', 'D', 'M', 'x', 'y']
        
        for seg in segments:
            if seg["type"] == "math":
                val = seg["value"]
                if seg.get("mode") == "display":
                    clean_val = val.replace('$$', '')
                    math_blocks.append(clean_val)
                else:
                    if len(val) > 1 and any(k in val for k in keywords):
                         math_blocks.append(val)
        
        seen = set()
        unique_blocks = []
        for m in math_blocks:
            if m not in seen:
                unique_blocks.append(m)
                seen.add(m)
        
        # Guard against huge aggregation (V68 Protection)
        result = " \\\\ ".join(unique_blocks)
        if len(result) > 2000:
            return result[:2000] + "..."
        return result