import json
import subprocess
import re
import logging
from typing import List, Dict, Tuple, Literal
from enum import Enum

logger = logging.getLogger(__name__)


class OutputFormat(str, Enum):
    """Supported output formats for LaTeX conversion."""
    SVG = "svg"
    HTML = "html"


class LatexConverter:
    """Efficient LaTeX converter using batched MathJax processing.
    
    Supports both SVG and HTML (CHTML) output formats.
    """

    node_path = "node"

    def __init__(
        self,
        script_path: str = "node_scripts/batch_latex_converter.mjs",
        output_format: OutputFormat = OutputFormat.SVG,
        cache_size: int = 1000,
    ):
        """
        Initialize the converter.

        Args:
            script_path: Path to the Node.js script
            output_format: Output format - 'svg' or 'html'
            cache_size: Unused in this implementation (kept for compatibility)
        """
        self.script_path = script_path
        self.output_format = output_format
        self._chtml_styles = None

    def with_node_path(self, node_path: str):
        """Use the node path to execute the script"""
        self.node_path = node_path
        return self

    def with_output_format(self, output_format: OutputFormat):
        """Set the output format (svg or html)"""
        self.output_format = output_format
        return self

    def find_latex_in_html(self, html_content: str) -> List[Tuple[str, str, bool]]:
        """
        Find all LaTeX formulas in HTML content.

        Returns:
            List of tuples (placeholder, latex, is_display)
        """
        formulas = []

        # Find display math ($$...$$ or \[...\])
        display_patterns = [
            (r"\$\$(.*?)\$\$", "$$"),
            (r"\\\[(.*?)\\\]", "\\[\\]"),
        ]

        for pattern, delim in display_patterns:
            for match in re.finditer(pattern, html_content, re.DOTALL):
                placeholder = match.group(0)
                latex = match.group(1).strip()
                formulas.append((placeholder, latex, True))

        # Find inline math ($...$ or \(...\))
        inline_patterns = [
            (r"(?<!\$)\$(?!\$)(.*?)(?<!\$)\$(?!\$)", "$"),
            (r"\\\((.*?)\\\)", "\\(\\)"),
        ]

        for pattern, delim in inline_patterns:
            for match in re.finditer(pattern, html_content, re.DOTALL):
                placeholder = match.group(0)
                latex = match.group(1).strip()
                formulas.append((placeholder, latex, False))

        logger.debug(f"formulas: {formulas}")

        return formulas

    def convert_batch(self, formulas: List[Tuple[str, bool]]) -> List[Dict]:
        """
        Convert multiple LaTeX formulas to SVG or HTML in a single call.

        Args:
            formulas: List of tuples (latex_string, is_display)

        Returns:
            List of dictionaries with 'svg'/'html' and 'output' or 'error' keys
        """
        if not formulas:
            return []

        # Prepare input data with format specification
        input_data = {
            "formulas": [],
            "format": self.output_format.value,
            "includeStyles": self.output_format == OutputFormat.HTML,
        }
        
        for i, (latex, display) in enumerate(formulas):
            input_data["formulas"].append({
                "id": i, 
                "latex": latex, 
                "display": display
            })

        # Call the Node.js script
        try:
            result = subprocess.run(
                [self.node_path, self.script_path],
                input=json.dumps(input_data),
                capture_output=True,
                text=True,
                check=True,
            )

            # Parse the results
            response = json.loads(result.stdout)
            logger.debug(f"Converted formulas response: {response}")

            # Handle both old array format and new object format
            if isinstance(response, list):
                outputs = response
            else:
                outputs = response.get("results", [])
                # Store CHTML styles if available
                if "styles" in response:
                    self._chtml_styles = response["styles"]

            # Sort by ID to maintain order
            outputs.sort(key=lambda x: x["id"])

            return outputs

        except subprocess.CalledProcessError as e:
            print(f"Error running Node.js script: {e.stderr}")
            logger.error(f"Error running Node.js script: {e.stderr}")
            # Return error for all formulas
            return [{"success": False, "error": str(e)} for _ in formulas]
        except json.JSONDecodeError as e:
            print(f"Error parsing output: {e}")
            logger.error(f"Error parsing output: {e}")
            return [{"success": False, "error": str(e)} for _ in formulas]

    def get_chtml_styles(self) -> str:
        """
        Get the CHTML styles required for HTML output.
        
        Returns:
            CSS styles as a string, or empty string if not available
        """
        return self._chtml_styles or ""

    def process_html(self, html_content: str, use_cache: bool = True) -> str:
        """
        Process HTML content and replace all LaTeX formulas with SVGs or HTML.

        Args:
            html_content: HTML content with LaTeX formulas
            use_cache: Unused in this implementation (kept for compatibility)

        Returns:
            HTML content with LaTeX replaced by SVG or HTML
        """
        # Find all formulas
        formula_data = self.find_latex_in_html(html_content)

        if not formula_data:
            return html_content

        # Deduplicate formulas to avoid redundant processing
        # Key: (latex, display), Value: output
        unique_formulas_map = {}
        unique_formulas_list = []

        for _, latex, display in formula_data:
            key = (latex, display)
            if key not in unique_formulas_map:
                unique_formulas_map[key] = None
                unique_formulas_list.append(key)

        # Batch convert all unique formulas
        if unique_formulas_list:
            results = self.convert_batch(unique_formulas_list)
            
            for (latex, display), result in zip(unique_formulas_list, results):
                if result["success"]:
                    # Use 'output' key which works for both formats, fallback to 'svg' or 'html'
                    output = result.get("output") or result.get("svg") or result.get("html")
                    unique_formulas_map[(latex, display)] = output
                else:
                    print(f"Failed to convert: {latex}")
                    # Leave as None to indicate failure

        # Replace all placeholders with converted output
        processed_html = html_content
        
        # Using a map of replacements to do it cleanly
        replacements = {}
        for placeholder, latex, display in formula_data:
            output = unique_formulas_map.get((latex, display))
            if output:
                is_display = display
                wrapper_class = "math-display" if is_display else "math-inline"
                wrapped_output = f'<span class="{wrapper_class}">{output}</span>'
                replacements[placeholder] = wrapped_output
        
        # Apply replacements
        # Sort keys by length descending to avoid substring replacement issues
        sorted_replacements = sorted(replacements.keys(), key=len, reverse=True)
        
        for placeholder in sorted_replacements:
            processed_html = processed_html.replace(placeholder, replacements[placeholder])

        # For HTML format, inject CHTML styles into the head if available
        if self.output_format == OutputFormat.HTML and self._chtml_styles:
            styles_tag = f'<style id="mjx-styles">{self._chtml_styles}</style>'
            # Try to inject before </head>, otherwise prepend to content
            if '</head>' in processed_html:
                processed_html = processed_html.replace('</head>', f'{styles_tag}</head>')
            elif '<head>' in processed_html:
                processed_html = processed_html.replace('<head>', f'<head>{styles_tag}')

        return processed_html


# Backward compatibility alias
LatexToSVGConverter = LatexConverter