from textual.app import App, ComposeResult
from textual.widgets import TextArea, Footer, Static, Input
from textual.containers import Horizontal, Vertical, VerticalScroll
from textual.screen import Screen
from textual.events import Key
import json


class WelcomeScreen(Screen):

    CSS = """
    WelcomeScreen {
        align: center middle;
    }
    #welcome {
        width: auto;
        height: auto;
        border: tall cyan;
        padding: 2 4;
        content-align: center middle;
    }
    """

    def compose(self) -> ComposeResult:
        yield Static("""
        [bold cyan] ██████╗██╗   ██╗██████╗ ██╗  ██╗   ██╗[/bold cyan]
        [bold cyan]██╔════╝██║   ██║██╔══██╗██║  ╚██╗ ██╔╝[/bold cyan]
        [bold cyan]██║     ██║   ██║██████╔╝██║   ╚████╔╝ [/bold cyan]
        [bold cyan]██║     ██║   ██║██╔══██╗██║    ╚██╔╝  [/bold cyan]
        [bold cyan]╚██████╗╚██████╔╝██║  ██║███████╗██║   [/bold cyan]
        [bold cyan] ╚═════╝ ╚═════╝ ╚═╝  ╚═╝╚══════╝╚═╝  [/bold cyan]

        [dim]your local JSON workbench[/dim]

        [bold]format[/bold]    pretty prints pasted JSON         [dim]F2[/dim]
        [bold]preview[/bold]   browse the structure as a tree    [dim]F3[/dim]
        [bold]search[/bold]    find any key or value             [dim]F4[/dim]
        [bold]diff[/bold]      compare two JSON documents        [dim]F5[/dim]
        [bold]stringify[/bold] compact single-line JSON          [dim]F6[/dim]

        [dim]press any key to start · [/dim][bold]ctrl+q[/bold][dim] to quit[/dim]
        """, id="welcome")

    def on_key(self, event: Key) -> None:
        if event.key == "ctrl+q":
            self.app.exit()
            return
        event.stop()
        self.app.pop_screen()

def format_json(raw: str) -> tuple[str, str | None]:
    if not raw:
        return "", None
    try:
        json_data = json.loads(raw)
        return json.dumps(json_data, indent=2), None
    except json.JSONDecodeError as e:
        return raw, str(e)

def search_tree(node, query: str, path: str = "") -> set[str]:
    matches = set()
    query = query.lower()

    if isinstance(node, dict):
        for key, value in node.items():
            current_path = f"{path}.{key}" if path else key

            if query in str(key).lower():
                matches.add(current_path)
            matches |= search_tree(value, query, current_path)
    elif isinstance(node, list):
        for i, value in enumerate(node):
            current_path = f"{path}[{i}]"
            matches |= search_tree(value, query, current_path)

    else:
        if query in str(node).lower():
            matches.add(path)

    return matches

def render_tree(node, indent=0, path="", matches=None) -> str:
    if matches is None:
        matches = set()

    pad = " " * indent

    if isinstance(node, dict):
        lines = []
        for key, value in node.items():
            current_path = f"{path}.{key}" if path else key
            highlight = "[yellow]" if current_path in matches else ""
            end = "[/yellow]" if current_path in matches else ""

            if isinstance(value, (dict, list)):
                lines.append(f"{pad}{highlight}[b]{key}[/b]{end}:")
                lines.append(render_tree(value, indent + 1, current_path, matches))
            else:
                lines.append(f"{pad}{highlight}[b]{key}[/b]: {value}{end}")
        return "\n".join(lines)
    elif isinstance(node, list):
        lines = []
        for i, value in enumerate(node):
            current_path = f"{path}[{i}]"
            highlight = "[yellow]" if current_path in matches else ""
            end = "[/yellow]" if current_path in matches else ""
            if isinstance(value, (dict, list)):
                lines.append(f"{pad}{highlight}[dim]{i}[/dim]{end}:")
                lines.append(render_tree(value, indent + 1, current_path, matches))
            else:
                lines.append(f"{pad}{highlight}[dim]{i}[/dim]: {value}{end}")
        return "\n".join(lines)
    else:
        return f"{pad}{node}"

def build_preview(raw: str) -> str:
    try:
        data = json.loads(raw)
        return render_tree(data)
    except json.JSONDecodeError as e:
        return f"[red]Invalid JSON:[/red]\n{e}"

def diff(a, b, path="") -> list[tuple]:
    results = []

    if isinstance(a, dict) and isinstance(b, dict):
        all_keys = set(a.keys()) | set(b.keys())
        for key in all_keys:
            current_path = f"{path}.{key}" if path else key
            if key not in b:
                results.append(("removed", current_path, a[key], None))
            elif key not in a:
                results.append(("added", current_path, None, b[key]))
            else:
                results.extend(diff(a[key], b[key], current_path))

    elif isinstance(a, list) and isinstance(b, list):
        for i in range(max(len(a), len(b))):
            current_path = f"{path}[{i}]"
            if i >= len(a):
                results.append(("added", current_path, None, b[i]))
            elif i >= len(b):
                results.append(("removed", current_path, a[i], None))
            else:
                results.extend(diff(a[i], b[i], current_path))
            
    elif a == b:
        results.append(("same", path, a, b))

    else:
        results.append(("changed", path, a, b))

    return results

def render_diff(diff_list: list[tuple]) -> str:
    lines = []
    for kind, path, a, b in diff_list:
        if kind == "same":
            lines.append(f"[dim]  {path}: {a}[/dim]")
        elif kind == "changed":
            lines.append(f"[yellow]~ {path}: {a} -> {b}[/yellow]")
        elif kind == "added":
            lines.append(f"[green]+ {path}: {b}[/green]")
        elif kind == "removed":
            lines.append(f"[red]- {path}: {a}[/red]")
    return "\n".join(lines)


class JsonTool(App):
    mode = "output"

    CSS = """
    Screen {
        layout: vertical;
    }
    Horizontal {
        height: 1fr;
    }
    TextArea {
        width: 1fr;
        height: 1fr;
    }
    #output-scroll {
        display: none;
        width: 1fr;
        height: 1fr;
        border: tall green;
        padding: 1;
    }
    #output-scroll.visible {
        display: block;
    }
    #editor-b {
        display: none;
    }
    #editor-b.visible {
        display: block;
    }
    #diff-scroll {
        display: none;
        height: 12;
        border: tall yellow;
        padding: 1;
    }
    #diff-scroll.visible {
        display: block;
    }
    Input {
        display: none;
    }
    Input.visible {
        display: block;
    }
    """

    BINDINGS = [
        ("f2", "format", "Format"),
        ("f3", "preview", "Preview"),
        ("f4", "search", "Search"),
        ("f5", "diff",      "Diff"),
        ("f6", "stringify", "Stringify"),
        ("ctrl+q", "quit", "Quit")
    ]

    def compose(self) -> ComposeResult:
        with Horizontal():
            yield TextArea(language="json", id="editor-a")
            with VerticalScroll(id="output-scroll"):
                yield Static("", id="output")
            yield TextArea(language="json", id="editor-b")
        with VerticalScroll(id="diff-scroll"):
            yield Static("", id="diff-panel")
        yield Input(placeholder="Search keys and values...", id="search-bar")
        yield Footer()


    def _editor_text(self, editor_id: str) -> str:
        return self.query_one(f"#{editor_id}", TextArea).text.strip()

    def switch_mode(self, mode: str):
        self.mode = mode
        output_scroll = self.query_one("#output-scroll")
        editor_b = self.query_one("#editor-b", TextArea)
        diff_scroll = self.query_one("#diff-scroll")

        output_scroll.set_class(mode != "diff", "visible")
        editor_b.set_class(mode == "diff", "visible")
        diff_scroll.set_class(mode == "diff", "visible")

        self.query_one("#editor-b").border_title = "JSON B"
        output_scroll.border_title = mode.capitalize()


    def action_format(self):
        try:
            self.exit_diff_mode()
            raw = self._editor_text("editor-a")
            result, error = format_json(raw)
            output = self.query_one("#output", Static)
            if error:
                output.update(f"[red]Format error: {error}[/red]")
            else:
                output.update(result)
        except Exception as e:
            self.query_one("#output", Static).update(f"[red]Error: {e}[/red]")

    def action_stringify(self):
        try:
            self.exit_diff_mode()
            raw = self._editor_text("editor-a")
            output = self.query_one("#output", Static)
            if not raw:
                output.update("")
                return
            data = json.loads(raw)
            output.update(json.dumps(data, separators=(",", ":")))
        except json.JSONDecodeError as e:
            self.query_one("#output", Static).update(f"[red]Format error: {e}[/red]")
        except Exception as e:
            self.query_one("#output", Static).update(f"[red]Error: {e}[/red]")

    def action_preview(self):
        try:
            self.exit_diff_mode()
            raw = self._editor_text("editor-a")
            self.query_one("#output", Static).update(build_preview(raw))
        except Exception as e:
            self.query_one("#output", Static).update(f"[red]Error: {e}[/red]")

    def action_search(self):
        self.exit_diff_mode()
        search_bar = self.query_one("#search-bar", Input)
        search_bar.toggle_class("visible")
        if search_bar.has_class("visible"):
            search_bar.focus()
        else:
            search_bar.blur()

    def run_diff(self):
        diff_panel = self.query_one("#diff-panel", Static)
        try:
            a_text = self._editor_text("editor-a")
            b_text = self._editor_text("editor-b")

            if not a_text or not b_text:
                diff_panel.update("[dim]Paste JSON into both panes then press F5[/dim]")
                return

            a_data = json.loads(a_text)
            b_data = json.loads(b_text)
            diff_list = diff(a_data, b_data)
            diff_panel.update(render_diff(diff_list))
        except json.JSONDecodeError as e:
            diff_panel.update(f"[red]Invalid JSON:[/red]\n{e}")
        except Exception as e:
            diff_panel.update(f"[red]Error: {e}[/red]")

    def action_diff(self):
        self.switch_mode("diff")
        self.run_diff()

    def exit_diff_mode(self):
        self.query_one("#editor-b", TextArea).remove_class("visible")
        self.query_one("#diff-scroll").remove_class("visible")
        self.query_one("#output-scroll").add_class("visible")

    def on_input_changed(self, event: Input.Changed) -> None:
        try:
            raw = self._editor_text("editor-a")
            output = self.query_one("#output", Static)
            data = json.loads(raw)
            matches = search_tree(data, event.value) if event.value else set()
            output.update(render_tree(data, matches=matches))
        except Exception:
            pass

    def on_mount(self) -> None:
        if not getattr(self, "file_a", "") and not getattr(self, "file_b", ""):
            self.push_screen(WelcomeScreen())

        self.query_one("#editor-a").border_title = "JSON A"
        self.query_one("#editor-b").border_title = "JSON B"
        self.query_one("#diff-scroll").border_title = "Diff"
        self.query_one("#output-scroll").border_title = "Output"
        self.query_one("#output-scroll").add_class("visible")
        self.mode = "output"

        if getattr(self, "file_a", ""):
            self.query_one("#editor-a", TextArea).load_text(self.file_a)
        if getattr(self, "file_b", ""):
            self.query_one("#editor-b", TextArea).load_text(self.file_b)
            self.switch_mode("diff")
            self.run_diff()

    def on_text_area_changed(self, event: TextArea.Changed) -> None:
        try:
            if event.text_area.id == "editor-a":
                if self.mode != "diff":
                    raw = event.text_area.text.strip()
                    result, error = format_json(raw)
                    output = self.query_one("#output", Static)
                    if error:
                        output.update(f"[red]Format error: {error}[/red]")
                    else:
                        self.switch_mode("output")
                        output.update(result)

            elif event.text_area.id == "editor-b":
                self.run_diff()
        except Exception as e:
            self.query_one("#output", Static).update(f"[red]Error: {e}[/red]")


def main():
    import sys

    app = JsonTool()

    if len(sys.argv) == 3:
        try:
            app.file_a = open(sys.argv[1], "r").read()
            app.file_b = open(sys.argv[2], "r").read()
        except FileNotFoundError as e:
            print(f"Error: {e}")
            sys.exit(1)
    else:
        app.file_a = ""
        app.file_b = ""

    app.run()

if __name__ == "__main__":
    main()
