# Copyright The Lightning AI team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import Any
from unittest import mock

import pytest
import torch

from lightning.pytorch import Trainer
from lightning.pytorch.callbacks import RichModelSummary, RichProgressBar
from lightning.pytorch.demos.boring_classes import BoringModel
from lightning.pytorch.utilities.model_summary import summarize
from tests_pytorch.helpers.runif import RunIf


@RunIf(rich=True)
def test_rich_model_summary_callback():
    trainer = Trainer(callbacks=RichProgressBar())

    assert any(isinstance(cb, RichModelSummary) for cb in trainer.callbacks)
    assert isinstance(trainer.progress_bar_callback, RichProgressBar)


def test_rich_progress_bar_import_error(monkeypatch):
    import lightning.pytorch.callbacks.rich_model_summary as imports

    monkeypatch.setattr(imports, "_RICH_AVAILABLE", False)
    with pytest.raises(ModuleNotFoundError, match="`RichModelSummary` requires `rich` to be installed."):
        RichModelSummary()


@RunIf(rich=True)
@mock.patch("rich.console.Console.print", autospec=True)
@mock.patch("rich.table.Table.add_row", autospec=True)
def test_rich_summary_tuples(mock_table_add_row, mock_console):
    """Ensure that tuples are converted into string, and print is called correctly."""
    model_summary = RichModelSummary()

    class TestModel(BoringModel):
        @property
        def example_input_array(self) -> Any:
            return torch.randn(4, 32)

    model = TestModel()
    summary = summarize(model)
    summary_data = summary._get_summary_data()

    model_summary.summarize(
        summary_data=summary_data,
        total_parameters=1,
        trainable_parameters=1,
        model_size=1,
        total_training_modes=summary.total_training_modes,
        total_flops=1,
    )

    # ensure that summary was logged + the breakdown of model parameters
    assert mock_console.call_count == 2
    # assert that the input summary data was converted correctly
    args, _ = mock_table_add_row.call_args_list[0]
    assert args[1:] == ("0", "layer", "Linear", "66  ", "train", "512  ", "[4, 32]", "[4, 2]")


@RunIf(rich=True)
def test_rich_summary_model_size_formatting():
    """Ensure model_size uses get_formatted_model_size, not get_human_readable_count."""
    from io import StringIO

    from rich.console import Console

    model_summary = RichModelSummary()
    model = BoringModel()
    summary = summarize(model)
    summary_data = summary._get_summary_data()

    output = StringIO()
    console = Console(file=output, force_terminal=True)

    with mock.patch("rich.get_console", return_value=console):
        model_summary.summarize(
            summary_data=summary_data,
            total_parameters=1,
            trainable_parameters=1,
            model_size=5500.0,
            total_training_modes=summary.total_training_modes,
            total_flops=1,
        )

    result = output.getvalue()
    # model_size=5500.0 should display as "5,500.000" (formatted), not "5.5 K" (human readable count)
    assert "5,500.000" in result
    assert "5.5 K" not in result
