import textwrap
import unittest
from unittest import mock

from rich.console import Console
from rich.table import Table

from atrace.interpreter import (
    UNASSIGN,
    Call,
    History,
    Line,
    LineEffects,
    Raise,
    Return,
    Var,
)
from atrace.reporter import LeftAligned, history_to_table, history_to_table_data


def capture_report(table: Table) -> str:
    console = Console(color_system=None, width=88)
    with console.capture() as capture:
        console.print(table)
    return capture.get()


class TestReporter(unittest.TestCase):
    def test_assign_none_unassign(self):
        history: History = [
            (1, Line()),
            (1, LineEffects({Var("<module>", "x"): 1}, None)),
            (2, Line()),
            (2, LineEffects({Var("<module>", "x"): None}, None)),
            (3, Line()),
            (3, LineEffects({Var("<module>", "x"): UNASSIGN}, None)),
            (4, Line()),
            (4, LineEffects({Var("<module>", "x"): "bob"}, None)),
        ]
        expected_table_data = (
            ["line", "x"],
            [
                ["1", "1"],
                ["2", "None"],
                ["3", "✖"],
                ["4", '"bob"'],
            ],
        )
        self.assertEqual(expected_table_data, history_to_table_data(history))

    def test_assign_then_print(self):
        history: History = [
            (1, Line()),
            (1, LineEffects({Var("<module>", "x"): 1}, None)),
            (2, Line()),
            (2, LineEffects({}, "1\n")),
        ]
        expected_table_data = (
            ["line", "x", "output"],
            [
                ["1", "1", ""],
                ["2", "", "1"],
            ],
        )
        self.assertEqual(expected_table_data, history_to_table_data(history))

    def test_function(self):
        history: History = [
            (1, Line()),
            # We need to pass a callable, otherwise it gets displayed
            (1, LineEffects({Var("<module>", "double"): lambda: None}, None)),
            (6, Line()),
            (1, Call("double", {Var("double", "a"): 3})),
            (2, Line()),
            (2, LineEffects({Var("double", "result"): 6}, None)),
            (3, Line()),
            (3, Return(6)),
            (6, LineEffects({Var("<module>", "x"): 6}, None)),
        ]
        expected_table_data = (
            ["line", LeftAligned("double"), "(double) a", "(double) result", "x"],
            [
                ["1", "double(3)", "3", "", ""],
                ["2", "│  ", "", "6", ""],
                ["3", "└─ 6", "", "", ""],
                ["6", "", "", "", "6"],
            ],
        )
        self.assertEqual(expected_table_data, history_to_table_data(history))

    def test_function_print_before_call(self):
        """This verifies that events happening after we initiated a call get
        assigned to the line we get to after the return"""
        history: History = [
            (1, Line()),
            # We need to pass a callable, otherwise it gets displayed
            (1, LineEffects({Var("<module>", "double"): lambda: None}, None)),
            (5, Line()),
            (5, LineEffects({}, "hahaha\n")),
            (6, Line()),
            (1, Call("double", {Var("double", "a"): 5})),
            (2, Line()),
            (2, Return(10)),
            (6, LineEffects({Var("<module>", "x"): 10}, None)),
        ]
        expected_table_data = (
            ["line", LeftAligned("double"), "(double) a", "x", "output"],
            [
                ["5", "", "", "", "hahaha"],
                ["1", "double(5)", "5", "", ""],
                ["2", "└─ 10", "", "", ""],
                ["6", "", "", "10", ""],
            ],
        )
        self.assertEqual(expected_table_data, history_to_table_data(history))

    def test_procedure(self):
        """A function that takes and returns nothing."""
        history: History = [
            (1, Line()),
            # We need to pass a callable, otherwise it gets displayed
            (1, LineEffects({Var("<module>", "p"): lambda: None}, None)),
            (5, Line()),
            (1, Call("p", {})),
            (2, Line()),
            (2, LineEffects({}, "hello\n")),
            (2, Return(None)),
        ]
        expected_table_data = (
            ["line", LeftAligned("p"), "output"],
            [
                ["1", "p()", ""],
                ["2", "│  ", "hello"],
                ["2", "└─ ", ""],
            ],
        )
        self.assertEqual(expected_table_data, history_to_table_data(history))

    def test_uncaught_exception(self):
        the_exception = Exception("an exception")
        history: History = [
            (1, Line()),
            (1, LineEffects({}, "hai\n")),
            (2, Line()),
            (2, LineEffects({Var("<module>", "x"): 1}, None)),
            (3, Line()),
            (3, Raise(type(the_exception), the_exception, mock.ANY)),
        ]
        expected_table_data = (
            ["line", "x", "output", "exception"],
            [
                ["1", "", "hai", ""],
                ["2", "1", "", ""],
                ["3", "", "", "Exception('an exception')"],
            ],
        )
        self.assertEqual(expected_table_data, history_to_table_data(history))

    def test_display(self):
        history: History = [
            (1, Line()),
            (3, Line()),
            (3, LineEffects({Var("<module>", "x"): 1, Var("<module>", "y"): 3}, None)),
            (5, Line()),
            (6, Line()),
            (6, LineEffects({Var("<module>", "x"): 2}, None)),
            (5, Line()),
            (6, Line()),
            (6, LineEffects({Var("<module>", "x"): 3}, None)),
            (5, Line()),
            (8, Line()),
            (8, LineEffects({}, "x: 3\n")),
            (11, Line()),
            # We need to pass a callable, otherwise it gets displayed
            (11, LineEffects({Var("<module>", "greet"): lambda: None}, None)),
            (16, Line()),
            (11, Call("greet", {Var("greet", "name"): "Bob"})),
            (12, Line()),
            (12, LineEffects({Var("greet", "message"): "Hi Bob!"}, None)),
            (13, Line()),
            (13, Return("Hi Bob!")),
            (16, LineEffects({}, "Hi Bob!\n")),
        ]
        expected_result = """\
        ╭──────┬───┬───┬──────────────┬──────────────┬─────────────────┬─────────╮
        │ line │ x │ y │ greet        │ (greet) name │ (greet) message │  output │
        ├──────┼───┼───┼──────────────┼──────────────┼─────────────────┼─────────┤
        │    3 │ 1 │ 3 │              │              │                 │         │
        │    6 │ 2 │   │              │              │                 │         │
        │    6 │ 3 │   │              │              │                 │         │
        │    8 │   │   │              │              │                 │    x: 3 │
        │   11 │   │   │ greet("Bob") │        "Bob" │                 │         │
        │   12 │   │   │ │            │              │       "Hi Bob!" │         │
        │   13 │   │   │ └─ "Hi Bob!" │              │                 │         │
        │   16 │   │   │              │              │                 │ Hi Bob! │
        ╰──────┴───┴───┴──────────────┴──────────────┴─────────────────┴─────────╯
        """
        self.assertEqual(
            textwrap.dedent(expected_result), capture_report(history_to_table(history))
        )
