#!/usr/bin/env python3
import sys
import json
import argparse
import jinja2.ext as jinja2_ext
from PySide6.QtWidgets import (
    QApplication,
    QMainWindow,
    QWidget,
    QVBoxLayout,
    QHBoxLayout,
    QLabel,
    QPlainTextEdit,
    QTextEdit,
    QPushButton,
    QFileDialog,
)
from PySide6.QtGui import QColor, QColorConstants, QTextCursor, QTextFormat
from PySide6.QtCore import Qt, QRect, QSize
from jinja2 import TemplateSyntaxError
from jinja2.sandbox import ImmutableSandboxedEnvironment
from datetime import datetime
from typing import Callable


def format_template_content(template_content):
    """Format the Jinja template content using Jinja2's lexer."""
    if not template_content.strip():
        return template_content

    env = ImmutableSandboxedEnvironment()
    tc_rstrip = template_content.rstrip()
    tokens = list(env.lex(tc_rstrip))
    result = ""
    indent_level = 0
    i = 0

    while i < len(tokens):
        token = tokens[i]
        _, token_type, token_value = token

        if token_type == "block_begin":
            block_start = i
            # Collect all tokens for this block construct
            construct_content = token_value
            end_token_type = token_type.replace("_begin", "_end")
            j = i + 1
            while j < len(tokens) and tokens[j][1] != end_token_type:
                construct_content += tokens[j][2]
                j += 1

            if j < len(tokens):  # Found the end token
                construct_content += tokens[j][2]
                i = j  # Skip to the end token

                # Check for control structure keywords for indentation
                stripped_content = construct_content.strip()
                instr = block_start + 1
                while tokens[instr][1] == "whitespace":
                    instr = instr + 1

                instruction_token = tokens[instr][2]
                start_control_tokens = ["if", "for", "macro", "call", "block"]
                end_control_tokens = ["end" + t for t in start_control_tokens]
                is_control_start = any(
                    instruction_token.startswith(kw) for kw in start_control_tokens
                )
                is_control_end = any(
                    instruction_token.startswith(kw) for kw in end_control_tokens
                )

                # Adjust indentation for control structures
                # For control end blocks, decrease indent BEFORE adding the content
                if is_control_end:
                    indent_level = max(0, indent_level - 1)

                # Remove all previous whitespace before this block
                result = result.rstrip()

                # Add proper indent, but only if this is not the first token
                added_newline = False
                if result:  # Only add newline and indent if there's already content
                    result += (
                        "\n" + "  " * indent_level
                    )  # Use 2 spaces per indent level
                    added_newline = True
                else:  # For the first token, don't add any indent
                    result += ""

                # Add the block content
                result += stripped_content

                # Add '-' after '%' if it wasn't there and we added a newline or indent
                if (
                    added_newline
                    and stripped_content.startswith("{%")
                    and not stripped_content.startswith("{%-")
                ):
                    # Add '-' at the beginning
                    result = (
                        result[: result.rfind("{%")]
                        + "{%-"
                        + result[result.rfind("{%") + 2 :]
                    )
                if stripped_content.endswith("%}") and not stripped_content.endswith(
                    "-%}"
                ):
                    # Only add '-' if this is not the last token or if there's content after
                    if i + 1 < len(tokens) and tokens[i + 1][1] != "eof":
                        result = result[:-2] + "-%}"

                # For control start blocks, increase indent AFTER adding the content
                if is_control_start:
                    indent_level += 1
            else:
                # Malformed template, just add the token
                result += token_value
        elif token_type == "variable_begin":
            # Collect all tokens for this variable construct
            construct_content = token_value
            end_token_type = token_type.replace("_begin", "_end")
            j = i + 1
            while j < len(tokens) and tokens[j][1] != end_token_type:
                construct_content += tokens[j][2]
                j += 1

            if j < len(tokens):  # Found the end token
                construct_content += tokens[j][2]
                i = j  # Skip to the end token

                # For variable constructs, leave them alone
                # Do not add indent or whitespace before or after them
                result += construct_content
            else:
                # Malformed template, just add the token
                result += token_value
        elif token_type == "data":
            # Handle data (text between Jinja constructs)
            # For data content, preserve it as is
            result += token_value
        else:
            # Handle any other tokens
            result += token_value

        i += 1

    # Clean up trailing newlines and spaces
    result = result.rstrip()

    # Copy the newline / space count from the original
    if (trailing_length := len(template_content) - len(tc_rstrip)):
        result += template_content[-trailing_length:]

    return result


# ------------------------
# Line Number Widget
# ------------------------
class LineNumberArea(QWidget):
    def __init__(self, editor):
        super().__init__(editor)
        self.code_editor = editor

    def sizeHint(self):
        return QSize(self.code_editor.line_number_area_width(), 0)

    def paintEvent(self, event):
        self.code_editor.line_number_area_paint_event(event)


class CodeEditor(QPlainTextEdit):
    def __init__(self):
        super().__init__()
        self.line_number_area = LineNumberArea(self)

        self.blockCountChanged.connect(self.update_line_number_area_width)
        self.updateRequest.connect(self.update_line_number_area)
        self.cursorPositionChanged.connect(self.highlight_current_line)

        self.update_line_number_area_width(0)
        self.highlight_current_line()

    def line_number_area_width(self):
        digits = len(str(self.blockCount()))
        space = 3 + self.fontMetrics().horizontalAdvance("9") * digits
        return space

    def update_line_number_area_width(self, _):
        self.setViewportMargins(self.line_number_area_width(), 0, 0, 0)

    def update_line_number_area(self, rect, dy):
        if dy:
            self.line_number_area.scroll(0, dy)
        else:
            self.line_number_area.update(
                0, rect.y(), self.line_number_area.width(), rect.height()
            )

        if rect.contains(self.viewport().rect()):
            self.update_line_number_area_width(0)

    def resizeEvent(self, event):
        super().resizeEvent(event)
        cr = self.contentsRect()
        self.line_number_area.setGeometry(
            QRect(cr.left(), cr.top(), self.line_number_area_width(), cr.height())
        )

    def line_number_area_paint_event(self, event):
        from PySide6.QtGui import QPainter

        painter = QPainter(self.line_number_area)
        painter.fillRect(event.rect(), QColorConstants.LightGray)

        block = self.firstVisibleBlock()
        block_number = block.blockNumber()
        top = int(
            self.blockBoundingGeometry(block).translated(self.contentOffset()).top()
        )
        bottom = top + int(self.blockBoundingRect(block).height())

        while block.isValid() and top <= event.rect().bottom():
            if block.isVisible() and bottom >= event.rect().top():
                number = str(block_number + 1)
                painter.setPen(QColorConstants.Black)
                painter.drawText(
                    0,
                    top,
                    self.line_number_area.width() - 2,
                    self.fontMetrics().height(),
                    Qt.AlignmentFlag.AlignRight,
                    number,
                )
            block = block.next()
            top = bottom
            bottom = top + int(self.blockBoundingRect(block).height())
            block_number += 1

    def highlight_current_line(self):
        extra_selections = []
        if not self.isReadOnly():
            selection = QTextEdit.ExtraSelection()
            line_color = QColorConstants.Yellow.lighter(160)
            selection.format.setBackground(line_color)  # pyright: ignore[reportAttributeAccessIssue] # ty: ignore[unresolved-attribute]
            selection.format.setProperty(QTextFormat.Property.FullWidthSelection, True)  # pyright: ignore[reportAttributeAccessIssue] # ty: ignore[unresolved-attribute]
            selection.cursor = self.textCursor()  # pyright: ignore[reportAttributeAccessIssue] # ty: ignore[unresolved-attribute]
            selection.cursor.clearSelection()  # pyright: ignore[reportAttributeAccessIssue] # ty: ignore[unresolved-attribute]
            extra_selections.append(selection)
        self.setExtraSelections(extra_selections)

    def highlight_position(self, lineno: int, col: int, color: QColor):
        block = self.document().findBlockByLineNumber(lineno - 1)
        if block.isValid():
            cursor = QTextCursor(block)
            text = block.text()
            start = block.position() + max(0, col - 1)
            cursor.setPosition(start)
            if col <= len(text):
                cursor.movePosition(
                    QTextCursor.MoveOperation.NextCharacter,
                    QTextCursor.MoveMode.KeepAnchor,
                )

            extra = QTextEdit.ExtraSelection()
            extra.format.setBackground(color.lighter(160))  # pyright: ignore[reportAttributeAccessIssue] # ty: ignore[unresolved-attribute]
            extra.cursor = cursor  # pyright: ignore[reportAttributeAccessIssue] # ty: ignore[unresolved-attribute]

            self.setExtraSelections(self.extraSelections() + [extra])

    def highlight_line(self, lineno: int, color: QColor):
        block = self.document().findBlockByLineNumber(lineno - 1)
        if block.isValid():
            cursor = QTextCursor(block)
            cursor.select(QTextCursor.SelectionType.LineUnderCursor)

            extra = QTextEdit.ExtraSelection()
            extra.format.setBackground(color.lighter(160))  # pyright: ignore[reportAttributeAccessIssue] # ty: ignore[unresolved-attribute]
            extra.cursor = cursor  # pyright: ignore[reportAttributeAccessIssue] # ty: ignore[unresolved-attribute]

            self.setExtraSelections(self.extraSelections() + [extra])

    def clear_highlighting(self):
        self.highlight_current_line()


# ------------------------
# Main App
# ------------------------
class JinjaTester(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("Jinja Template Tester")
        self.resize(1200, 800)

        central = QWidget()
        main_layout = QVBoxLayout(central)

        # -------- Top input area --------
        input_layout = QHBoxLayout()

        # Template editor with label
        template_layout = QVBoxLayout()
        template_label = QLabel("Jinja2 Template")
        template_layout.addWidget(template_label)
        self.template_edit = CodeEditor()
        template_layout.addWidget(self.template_edit)
        input_layout.addLayout(template_layout)

        # JSON editor with label
        json_layout = QVBoxLayout()
        json_label = QLabel("Context (JSON)")
        json_layout.addWidget(json_label)
        self.json_edit = CodeEditor()
        self.json_edit.setPlainText("""
{
    "add_generation_prompt": true,
    "bos_token": "",
    "eos_token": "",
    "messages": [
        {
            "role": "user",
            "content": "What is the capital of Poland?"
        }
    ]
}
        """.strip())
        json_layout.addWidget(self.json_edit)
        input_layout.addLayout(json_layout)

        main_layout.addLayout(input_layout)

        # -------- Rendered output area --------
        output_label = QLabel("Rendered Output")
        main_layout.addWidget(output_label)
        self.output_edit = QPlainTextEdit()
        self.output_edit.setReadOnly(True)
        main_layout.addWidget(self.output_edit)

        # -------- Render button and status --------
        btn_layout = QHBoxLayout()

        # Load template button
        self.load_btn = QPushButton("Load Template")
        self.load_btn.clicked.connect(self.load_template)
        btn_layout.addWidget(self.load_btn)

        # Format template button
        self.format_btn = QPushButton("Format")
        self.format_btn.clicked.connect(self.format_template)
        btn_layout.addWidget(self.format_btn)

        self.render_btn = QPushButton("Render")
        self.render_btn.clicked.connect(self.render_template)
        btn_layout.addWidget(self.render_btn)
        main_layout.addLayout(btn_layout)

        # Status label below buttons
        self.status_label = QLabel("Ready")
        main_layout.addWidget(self.status_label)

        self.setCentralWidget(central)

    def render_template(self):
        self.template_edit.clear_highlighting()
        self.output_edit.clear()

        template_str = self.template_edit.toPlainText()
        json_str = self.json_edit.toPlainText()

        # Parse JSON context
        try:
            context = json.loads(json_str) if json_str.strip() else {}
        except Exception as e:
            self.status_label.setText(f"❌ JSON Error: {e}")
            return

        def raise_exception(text: str) -> str:
            raise RuntimeError(text)

        env = ImmutableSandboxedEnvironment(
            trim_blocks=True,
            lstrip_blocks=True,
            extensions=[jinja2_ext.loopcontrols],
        )
        env.filters["tojson"] = (
            lambda x,
            indent=None,
            separators=None,
            sort_keys=False,
            ensure_ascii=False: json.dumps(
                x,
                indent=indent,
                separators=separators,
                sort_keys=sort_keys,
                ensure_ascii=ensure_ascii,
            )
        )
        env.globals["strftime_now"]: Callable[[str], str] = lambda format: datetime.now().strftime(format)
        env.globals["raise_exception"] = raise_exception  # ty: ignore[invalid-assignment]
        try:
            template = env.from_string(template_str)
            output = template.render(context)
            self.output_edit.setPlainText(output)
            self.status_label.setText("✅ Render successful")
        except TemplateSyntaxError as e:
            self.status_label.setText(f"❌ Syntax Error (line {e.lineno}): {e.message}")
            if e.lineno:
                self.template_edit.highlight_line(e.lineno, QColor("red"))
        except Exception as e:
            # Catch all runtime errors
            # Try to extract template line number
            lineno = None
            tb = e.__traceback__
            while tb:
                frame = tb.tb_frame
                if frame.f_code.co_filename == "<template>":
                    lineno = tb.tb_lineno
                    break
                tb = tb.tb_next

            error_msg = f"Runtime Error: {type(e).__name__}: {e}"
            if lineno:
                error_msg = f"Runtime Error at line {lineno} in template: {type(e).__name__}: {e}"
                self.template_edit.highlight_line(lineno, QColor("orange"))

            self.output_edit.setPlainText(error_msg)
            self.status_label.setText(f"❌ {error_msg}")

    def load_template(self):
        """Load a Jinja template from a file using a file dialog."""
        file_path, _ = QFileDialog.getOpenFileName(
            self,
            "Load Jinja Template",
            "",
            "Template Files (*.jinja *.j2 *.html *.txt);;All Files (*)",
        )

        if file_path:
            try:
                with open(file_path, "r", encoding="utf-8") as file:
                    content = file.read()
                    self.template_edit.setPlainText(content)
                    self.status_label.setText(f"✅ Loaded template from {file_path}")
            except Exception as e:
                self.status_label.setText(f"❌ Error loading file: {str(e)}")

    def format_template(self):
        """Format the Jinja template using Jinja2's lexer for proper parsing."""
        try:
            template_content = self.template_edit.toPlainText()
            if not template_content.strip():
                self.status_label.setText("⚠️ Template is empty")
                return

            formatted_content = format_template_content(template_content)
            self.template_edit.setPlainText(formatted_content)
            self.status_label.setText("✅ Template formatted")
        except Exception as e:
            self.status_label.setText(f"❌ Error formatting template: {str(e)}")


if __name__ == "__main__":
    if len(sys.argv) > 1:
        # CLI mode
        parser = argparse.ArgumentParser(description="Jinja Template Tester")
        parser.add_argument(
            "--template", required=True, help="Path to Jinja template file"
        )
        parser.add_argument("--context", required=True, help="JSON string for context")
        parser.add_argument(
            "--action",
            choices=["format", "render"],
            default="render",
            help="Action to perform",
        )
        args = parser.parse_args()

        # Load template
        with open(args.template, "r", encoding="utf-8") as f:
            template_content = f.read()

        # Load JSON
        context = json.loads(args.context)
        # Add missing variables
        context.setdefault("bos_token", "")
        context.setdefault("eos_token", "")
        context.setdefault("add_generation_prompt", False)

        env = ImmutableSandboxedEnvironment()

        if args.action == "format":
            formatted = format_template_content(template_content)
            print(formatted) # noqa: NP100
        elif args.action == "render":
            template = env.from_string(template_content)
            output = template.render(context)
            print(output) # noqa: NP100

    else:
        # GUI mode
        app = QApplication(sys.argv)
        window = JinjaTester()
        window.show()
        sys.exit(app.exec())