capabilities/templates.py

"""
Jinja Templates and MDX Components for MiniMind Max2
Enhanced code generation with rich formatting support.
"""

from dataclasses import dataclass, field
from typing import List, Optional, Dict, Any, Union
import re
import json
from enum import Enum


class MDXComponentType(Enum):
    """Supported MDX component types."""
    LINEAR_PROCESS_FLOW = "LinearProcessFlow"
    QUIZ = "Quiz"
    MATH = "math"
    THINKING = "Thinking"
    CODE_BLOCK = "CodeBlock"
    MERMAID = "Mermaid"
    CALLOUT = "Callout"
    TABS = "Tabs"
    ACCORDION = "Accordion"


@dataclass
class CodeBlockMeta:
    """Metadata for extended code blocks."""
    project: str = ""
    file: str = ""
    block_type: str = "code"  # react, nodejs, html, markdown, diagram, code
    language: str = "python"
    title: str = ""
    show_line_numbers: bool = True
    highlight_lines: List[int] = field(default_factory=list)
    executable: bool = False


@dataclass
class ThinkingBlock:
    """Represents a thinking/reasoning block."""
    content: str
    step: int = 0
    is_final: bool = False
    confidence: float = 1.0


class JinjaTemplateEngine:
    """
    Jinja-style template engine for code generation.
    Supports variables, loops, conditionals, and filters.
    """

    def __init__(self):
        self.filters = {
            "upper": str.upper,
            "lower": str.lower,
            "capitalize": str.capitalize,
            "title": str.title,
            "strip": str.strip,
            "length": len,
            "default": lambda x, d: x if x else d,
            "join": lambda x, sep=", ": sep.join(str(i) for i in x),
            "first": lambda x: x[0] if x else None,
            "last": lambda x: x[-1] if x else None,
            "json": json.dumps,
        }
        self.globals = {}

    def register_filter(self, name: str, func):
        """Register a custom filter."""
        self.filters[name] = func

    def register_global(self, name: str, value):
        """Register a global variable."""
        self.globals[name] = value

    def render(self, template: str, context: Dict[str, Any]) -> str:
        """
        Render a Jinja-style template.

        Args:
            template: Template string with {{ }} and {% %} blocks
            context: Variables to substitute

        Returns:
            Rendered string
        """
        # Merge context with globals
        full_context = {**self.globals, **context}

        # Process control structures first
        result = self._process_control_structures(template, full_context)

        # Then substitute variables
        result = self._substitute_variables(result, full_context)

        return result

    def _substitute_variables(self, template: str, context: Dict[str, Any]) -> str:
        """Substitute {{ variable }} expressions."""
        pattern = r'\{\{\s*(.+?)\s*\}\}'

        def replace(match):
            expr = match.group(1)
            return str(self._evaluate_expression(expr, context))

        return re.sub(pattern, replace, template)

    def _evaluate_expression(self, expr: str, context: Dict[str, Any]) -> Any:
        """Evaluate a Jinja expression."""
        # Handle filters (e.g., name|upper)
        if '|' in expr:
            parts = expr.split('|')
            value = self._evaluate_expression(parts[0].strip(), context)
            for filter_expr in parts[1:]:
                filter_expr = filter_expr.strip()
                # Handle filter with arguments
                if '(' in filter_expr:
                    filter_name = filter_expr[:filter_expr.index('(')]
                    args_str = filter_expr[filter_expr.index('(')+1:-1]
                    if filter_name in self.filters:
                        value = self.filters[filter_name](value, args_str.strip('"\''))
                else:
                    if filter_expr in self.filters:
                        value = self.filters[filter_expr](value)
            return value

        # Handle attribute access (e.g., user.name)
        if '.' in expr:
            parts = expr.split('.')
            value = context.get(parts[0])
            for part in parts[1:]:
                if isinstance(value, dict):
                    value = value.get(part)
                elif hasattr(value, part):
                    value = getattr(value, part)
                else:
                    return None
            return value

        # Handle index access (e.g., items[0])
        if '[' in expr:
            match = re.match(r'(\w+)\[(.+)\]', expr)
            if match:
                var_name = match.group(1)
                index = match.group(2)
                value = context.get(var_name)
                if value:
                    try:
                        idx = int(index)
                        return value[idx]
                    except (ValueError, IndexError):
                        return None
            return None

        # Simple variable lookup
        return context.get(expr, '')

    def _process_control_structures(self, template: str, context: Dict[str, Any]) -> str:
        """Process {% for %}, {% if %}, etc."""
        result = template

        # Process for loops
        for_pattern = r'\{%\s*for\s+(\w+)\s+in\s+(\w+)\s*%\}(.*?)\{%\s*endfor\s*%\}'
        while True:
            match = re.search(for_pattern, result, re.DOTALL)
            if not match:
                break

            var_name = match.group(1)
            iterable_name = match.group(2)
            body = match.group(3)

            iterable = context.get(iterable_name, [])
            rendered_parts = []

            for item in iterable:
                item_context = {**context, var_name: item}
                rendered_parts.append(self._substitute_variables(body, item_context))

            result = result[:match.start()] + ''.join(rendered_parts) + result[match.end():]

        # Process if statements
        if_pattern = r'\{%\s*if\s+(.+?)\s*%\}(.*?)(?:\{%\s*else\s*%\}(.*?))?\{%\s*endif\s*%\}'
        while True:
            match = re.search(if_pattern, result, re.DOTALL)
            if not match:
                break

            condition = match.group(1)
            true_block = match.group(2)
            false_block = match.group(3) or ''

            if self._evaluate_condition(condition, context):
                replacement = true_block
            else:
                replacement = false_block

            result = result[:match.start()] + replacement + result[match.end():]

        return result

    def _evaluate_condition(self, condition: str, context: Dict[str, Any]) -> bool:
        """Evaluate a condition expression."""
        # Simple conditions
        if ' == ' in condition:
            left, right = condition.split(' == ', 1)
            left_val = self._evaluate_expression(left.strip(), context)
            right_val = right.strip().strip('"\'')
            return str(left_val) == right_val

        if ' != ' in condition:
            left, right = condition.split(' != ', 1)
            left_val = self._evaluate_expression(left.strip(), context)
            right_val = right.strip().strip('"\'')
            return str(left_val) != right_val

        if ' > ' in condition:
            left, right = condition.split(' > ', 1)
            left_val = self._evaluate_expression(left.strip(), context)
            right_val = float(right.strip())
            return float(left_val) > right_val

        if ' < ' in condition:
            left, right = condition.split(' < ', 1)
            left_val = self._evaluate_expression(left.strip(), context)
            right_val = float(right.strip())
            return float(left_val) < right_val

        # Truthiness check
        value = self._evaluate_expression(condition.strip(), context)
        return bool(value)


class MDXRenderer:
    """
    Render MDX components for rich documentation and UI.
    """

    def __init__(self):
        self.components = {}
        self._register_default_components()

    def _register_default_components(self):
        """Register default MDX components."""
        self.components = {
            MDXComponentType.LINEAR_PROCESS_FLOW: self._render_linear_process_flow,
            MDXComponentType.QUIZ: self._render_quiz,
            MDXComponentType.MATH: self._render_math,
            MDXComponentType.THINKING: self._render_thinking,
            MDXComponentType.MERMAID: self._render_mermaid,
            MDXComponentType.CALLOUT: self._render_callout,
            MDXComponentType.TABS: self._render_tabs,
        }

    def render(self, component_type: MDXComponentType, props: Dict[str, Any]) -> str:
        """Render an MDX component."""
        if component_type in self.components:
            return self.components[component_type](props)
        return f"<!-- Unknown component: {component_type} -->"

    def _render_linear_process_flow(self, props: Dict[str, Any]) -> str:
        """Render a linear process flow diagram."""
        steps = props.get("steps", [])
        title = props.get("title", "Process Flow")

        html = f'<div class="linear-process-flow">\n'
        html += f'  <h3>{title}</h3>\n'
        html += '  <div class="steps">\n'

        for i, step in enumerate(steps):
            html += f'    <div class="step" data-step="{i+1}">\n'
            html += f'      <div class="step-number">{i+1}</div>\n'
            html += f'      <div class="step-content">\n'
            html += f'        <h4>{step.get("title", f"Step {i+1}")}</h4>\n'
            html += f'        <p>{step.get("description", "")}</p>\n'
            html += '      </div>\n'
            if i < len(steps) - 1:
                html += '      <div class="step-arrow">→</div>\n'
            html += '    </div>\n'

        html += '  </div>\n'
        html += '</div>'

        return html

    def _render_quiz(self, props: Dict[str, Any]) -> str:
        """Render an interactive quiz."""
        questions = props.get("questions", [])
        title = props.get("title", "Quiz")

        html = f'<div class="quiz" data-quiz-id="{props.get("id", "quiz")}">\n'
        html += f'  <h3>{title}</h3>\n'

        for i, q in enumerate(questions):
            html += f'  <div class="question" data-question="{i}">\n'
            html += f'    <p class="question-text">{q.get("question", "")}</p>\n'
            html += '    <div class="options">\n'

            for j, opt in enumerate(q.get("options", [])):
                correct = j == q.get("correct", 0)
                html += f'      <label class="option" data-correct="{str(correct).lower()}">\n'
                html += f'        <input type="radio" name="q{i}" value="{j}">\n'
                html += f'        <span>{opt}</span>\n'
                html += '      </label>\n'

            html += '    </div>\n'
            if q.get("explanation"):
                html += f'    <div class="explanation" hidden>{q["explanation"]}</div>\n'
            html += '  </div>\n'

        html += '</div>'
        return html

    def _render_math(self, props: Dict[str, Any]) -> str:
        """Render LaTeX math expressions."""
        expression = props.get("expression", "")
        display = props.get("display", False)

        if display:
            return f'$${expression}$$'
        return f'${expression}$'

    def _render_thinking(self, props: Dict[str, Any]) -> str:
        """Render a thinking/reasoning block."""
        content = props.get("content", "")
        step = props.get("step", 0)
        is_visible = props.get("visible", False)

        css_class = "thinking-block" + (" visible" if is_visible else " hidden")

        html = f'<div class="{css_class}" data-step="{step}">\n'
        html += '  <div class="thinking-header">\n'
        html += f'    <span class="thinking-icon">🤔</span>\n'
        html += f'    <span class="thinking-label">Thinking (Step {step})</span>\n'
        html += '  </div>\n'
        html += f'  <div class="thinking-content">{content}</div>\n'
        html += '</div>'

        return html

    def _render_mermaid(self, props: Dict[str, Any]) -> str:
        """Render a Mermaid diagram."""
        code = props.get("code", "")
        title = props.get("title", "")

        html = '<div class="mermaid-diagram">\n'
        if title:
            html += f'  <div class="diagram-title">{title}</div>\n'
        html += f'  <pre class="mermaid">\n{code}\n  </pre>\n'
        html += '</div>'

        return html

    def _render_callout(self, props: Dict[str, Any]) -> str:
        """Render a callout/alert box."""
        callout_type = props.get("type", "info")  # info, warning, error, success
        title = props.get("title", "")
        content = props.get("content", "")

        icons = {
            "info": "ℹ️",
            "warning": "⚠️",
            "error": "❌",
            "success": "✅",
        }

        html = f'<div class="callout callout-{callout_type}">\n'
        html += f'  <div class="callout-icon">{icons.get(callout_type, "ℹ️")}</div>\n'
        html += '  <div class="callout-content">\n'
        if title:
            html += f'    <div class="callout-title">{title}</div>\n'
        html += f'    <div class="callout-body">{content}</div>\n'
        html += '  </div>\n'
        html += '</div>'

        return html

    def _render_tabs(self, props: Dict[str, Any]) -> str:
        """Render a tabbed interface."""
        tabs = props.get("tabs", [])

        html = '<div class="tabs-container">\n'
        html += '  <div class="tabs-header">\n'

        for i, tab in enumerate(tabs):
            active = "active" if i == 0 else ""
            html += f'    <button class="tab-button {active}" data-tab="{i}">{tab.get("label", f"Tab {i+1}")}</button>\n'

        html += '  </div>\n'
        html += '  <div class="tabs-content">\n'

        for i, tab in enumerate(tabs):
            active = "active" if i == 0 else ""
            html += f'    <div class="tab-panel {active}" data-tab="{i}">{tab.get("content", "")}</div>\n'

        html += '  </div>\n'
        html += '</div>'

        return html


class ExtendedCodeBlockParser:
    """
    Parse and generate extended code blocks with metadata.
    """

    # Language to file extension mapping
    LANG_EXTENSIONS = {
        "python": ".py",
        "javascript": ".js",
        "typescript": ".ts",
        "tsx": ".tsx",
        "jsx": ".jsx",
        "rust": ".rs",
        "go": ".go",
        "java": ".java",
        "cpp": ".cpp",
        "c": ".c",
        "html": ".html",
        "css": ".css",
        "json": ".json",
        "yaml": ".yaml",
        "markdown": ".md",
        "sql": ".sql",
        "bash": ".sh",
        "shell": ".sh",
    }

    @classmethod
    def parse(cls, code_block: str) -> tuple[CodeBlockMeta, str]:
        """
        Parse an extended code block.

        Args:
            code_block: Full code block including ``` markers

        Returns:
            Tuple of (metadata, code_content)
        """
        lines = code_block.strip().split('\n')
        first_line = lines[0]

        # Parse first line for metadata
        meta = CodeBlockMeta()

        # Extract language
        lang_match = re.match(r'```(\w+)', first_line)
        if lang_match:
            meta.language = lang_match.group(1)

        # Extract attributes
        attrs = {}
        attr_pattern = r'(\w+)="([^"]*)"'
        for match in re.finditer(attr_pattern, first_line):
            attrs[match.group(1)] = match.group(2)

        meta.project = attrs.get("project", "")
        meta.file = attrs.get("file", "")
        meta.block_type = attrs.get("type", "code")
        meta.title = attrs.get("title", "")
        meta.executable = meta.block_type in ["react", "nodejs", "html"]

        # Handle highlight lines
        if "highlight" in attrs:
            try:
                meta.highlight_lines = [int(x) for x in attrs["highlight"].split(",")]
            except ValueError:
                pass

        # Extract code content
        code_lines = lines[1:-1] if lines[-1].strip() == '```' else lines[1:]
        code = '\n'.join(code_lines)

        return meta, code

    @classmethod
    def format(cls, code: str, meta: CodeBlockMeta) -> str:
        """
        Format code with extended code block syntax.

        Args:
            code: Code content
            meta: Metadata for the code block

        Returns:
            Formatted code block string
        """
        attrs = []

        if meta.project:
            attrs.append(f'project="{meta.project}"')
        if meta.file:
            attrs.append(f'file="{meta.file}"')
        if meta.block_type != "code":
            attrs.append(f'type="{meta.block_type}"')
        if meta.title:
            attrs.append(f'title="{meta.title}"')
        if meta.highlight_lines:
            attrs.append(f'highlight="{",".join(str(x) for x in meta.highlight_lines)}"')

        attr_str = " " + " ".join(attrs) if attrs else ""

        return f"```{meta.language}{attr_str}\n{code}\n```"

    @classmethod
    def create_react_component(
        cls,
        code: str,
        component_name: str,
        project: str = "MiniMind",
    ) -> str:
        """Create a React component code block."""
        meta = CodeBlockMeta(
            project=project,
            file=f"components/{component_name}.tsx",
            block_type="react",
            language="tsx",
        )
        return cls.format(code, meta)

    @classmethod
    def create_mermaid_diagram(
        cls,
        code: str,
        title: str = "",
    ) -> str:
        """Create a Mermaid diagram code block."""
        meta = CodeBlockMeta(
            block_type="diagram",
            language="mermaid",
            title=title,
        )
        return cls.format(code, meta)


class TemplateLibrary:
    """
    Pre-built templates for common code patterns.
    """

    TEMPLATES = {
        "react_component": '''
import React from 'react';

interface {{ name }}Props {
{% for prop in props %}
  {{ prop.name }}: {{ prop.type }};
{% endfor %}
}

export const {{ name }}: React.FC<{{ name }}Props> = ({ {% for prop in props %}{{ prop.name }}{% if not loop.last %}, {% endif %}{% endfor %} }) => {
  return (
    <div className="{{ name | lower }}">
      {{ content }}
    </div>
  );
};
''',
        "python_class": '''
class {{ name }}:
    """{{ description }}"""

    def __init__(self{% for param in params %}, {{ param.name }}: {{ param.type }}{% if param.default %} = {{ param.default }}{% endif %}{% endfor %}):
{% for param in params %}
        self.{{ param.name }} = {{ param.name }}
{% endfor %}

{% for method in methods %}
    def {{ method.name }}(self{% for arg in method.args %}, {{ arg.name }}: {{ arg.type }}{% endfor %}) -> {{ method.return_type }}:
        """{{ method.description }}"""
        {{ method.body }}

{% endfor %}
''',
        "api_endpoint": '''
from fastapi import APIRouter, HTTPException
from pydantic import BaseModel

router = APIRouter(prefix="/{{ prefix }}", tags=["{{ tag }}"])

{% for model in models %}
class {{ model.name }}(BaseModel):
{% for field in model.fields %}
    {{ field.name }}: {{ field.type }}
{% endfor %}

{% endfor %}

{% for endpoint in endpoints %}
@router.{{ endpoint.method }}("{{ endpoint.path }}")
async def {{ endpoint.name }}({% for param in endpoint.params %}{{ param.name }}: {{ param.type }}{% if not loop.last %}, {% endif %}{% endfor %}):
    """{{ endpoint.description }}"""
    {{ endpoint.body }}

{% endfor %}
''',
    }

    def __init__(self):
        self.engine = JinjaTemplateEngine()

    def render(self, template_name: str, context: Dict[str, Any]) -> str:
        """Render a pre-built template."""
        if template_name not in self.TEMPLATES:
            raise ValueError(f"Unknown template: {template_name}")

        template = self.TEMPLATES[template_name]
        return self.engine.render(template, context)

    def add_template(self, name: str, template: str):
        """Add a custom template."""
        self.TEMPLATES[name] = template