from __future__ import annotations

import os
import json
from pathlib import Path
from textual.app import App, ComposeResult
from textual.containers import Horizontal, VerticalScroll
from textual.widgets import Tree, Static, Input, Button, RadioSet, RadioButton

class JSONEditorApp(App):
    """A Textual UI application for editing JSON data."""

    CSS_PATH = "app.tcss"

    def __init__(self):
        super().__init__()
        self.json_data = []
        self.selected_entry_index = None
        self.json_file_path = Path("output.json")

    async def on_mount(self):
        """Initialize the app, load JSON data, and populate the tree."""
        await self.load_json()
        self.populate_json_tree()

    async def load_json(self):
        """Load JSON data from a file."""
        if self.json_file_path.exists():
            with open(self.json_file_path, "r") as file:
                self.json_data = json.load(file)
        else:
            self.json_data = []

    async def save_json(self):
        """Save JSON data to a file."""
        with open(self.json_file_path, "w") as file:
            json.dump(self.json_data, file, indent=4)

    def populate_json_tree(self):
        """Populate the JSON tree widget with structured JSON data."""
        tree: Tree[dict] = self.query_one("#json-tree", Tree)  # Specify data type as dict
        tree.clear()  # Clear any existing nodes

        for idx, entry in enumerate(self.json_data):
            entry_label = f"Entry {idx + 1}"
            parent_node = tree.root.add(entry_label, entry)  # Add entry as a root-level node

            # Add key-value pairs as children of the parent node
            for key, value in entry.items():
                if isinstance(value, str) and len(value) > 100:  # Truncate long strings
                    value = value[:100] + "..."
                parent_node.add(f"{key}: {value}")

        tree.root.expand()  # Expand the root node to show all entries

    def compose(self) -> ComposeResult:
        """Compose the app layout."""
        yield Horizontal(
            VerticalScroll(
                Tree("JSON Data", id="json-tree"),
                id="tree-panel"
            ),
            VerticalScroll(
                Static("Editor", id="editor-title"),
                Static("Word Problem 1:", id="wp1-label"),
                Input(id="wp1-input"),
                Static("Word Problem 2:", id="wp2-label"),
                Input(id="wp2-input"),
                Static("Word Problem 3:", id="wp3-label"),
                Input(id="wp3-input"),
                Static("Python Script:", id="script-label"),
                Input(id="script-input"),
                Static("Grading:", id="grading-label"),
                RadioSet(
                    RadioButton("Non-Graded", id="non-graded-rb"),
                    RadioButton("Synthetic", id="synthetic-rb"),
                    RadioButton("Human-Derived", id="human-rb"),
                    id="grading-options"
                ),
                Button("Save Changes", id="save-btn"),
                id="editor-panel"
            ),
            id="main-layout"
        )

    def on_tree_node_selected(self, event: Tree.NodeSelected):
        """Update editor fields when a tree node is selected."""
        entry = event.node.data
        if entry is not None:
            self.selected_entry_index = self.json_data.index(entry)
            self.update_editor(entry)

    def update_editor(self, entry: dict):
        """Populate the editor fields with the selected entry's data."""
        self.query_one("#wp1-input", Input).value = entry.get("word_problem_01", "")
        self.query_one("#wp2-input", Input).value = entry.get("word_problem_02", "")
        self.query_one("#wp3-input", Input).value = entry.get("word_problem_03", "")
        self.query_one("#script-input", Input).value = entry.get("python_script", "")

        grading_options = self.query_one("#grading-options", RadioSet)
        grading_options.value = entry.get("graded", "non-graded")

    def on_button_pressed(self, event: Button.Pressed):
        """Handle Save Changes button press."""
        if event.button.id == "save-btn":
            self.save_changes()

    def save_changes(self):
        """Save changes from the editor to the selected JSON entry."""
        if self.selected_entry_index is None:
            return

        entry = self.json_data[self.selected_entry_index]

        entry["word_problem_01"] = self.query_one("#wp1-input", Input).value
        entry["word_problem_02"] = self.query_one("#wp2-input", Input).value
        entry["word_problem_03"] = self.query_one("#wp3-input", Input).value
        entry["python_script"] = self.query_one("#script-input", Input).value
        entry["graded"] = self.query_one("#grading-options", RadioSet).value

        self.save_json()
        self.populate_json_tree()

if __name__ == "__main__":
    JSONEditorApp().run()