from io import StringIO
from pathlib import Path
from unittest.mock import patch

import pytest

from mflux.cli.parser.parsers import CommandLineParser
from mflux.utils.prompt_util import PromptUtil


@pytest.fixture
def mflux_generate_parser() -> CommandLineParser:
    parser = CommandLineParser(description="Generate an image based on a prompt.")
    parser.add_general_arguments()
    parser.add_model_arguments(require_model_arg=False)
    parser.add_image_generator_arguments(supports_metadata_config=True)
    parser.add_lora_arguments()
    parser.add_image_to_image_arguments(required=False)
    parser.add_image_outpaint_arguments()
    parser.add_output_arguments()
    return parser


@pytest.fixture
def temp_output_dir(tmp_path_factory) -> Path:
    return tmp_path_factory.mktemp("mflux_stdin_test")


@pytest.mark.fast
def test_prompt_from_stdin(mflux_generate_parser):
    stdin_content = "A beautiful sunset over the ocean"

    # Simulate stdin input
    with patch("sys.stdin", StringIO(stdin_content)):
        with patch("sys.argv", ["mflux-generate", "--prompt", "-", "--model", "dev"]):
            args = mflux_generate_parser.parse_args()

            # The parser returns the raw args, read_prompt handles stdin
            assert args.prompt == "-"


@pytest.mark.fast
def test_prompt_stdin_vs_regular(mflux_generate_parser):
    regular_prompt = "A regular prompt not from stdin"

    with patch("sys.argv", ["mflux-generate", "--prompt", regular_prompt, "--model", "dev"]):
        args = mflux_generate_parser.parse_args()
        assert args.prompt == regular_prompt
        assert PromptUtil.read_prompt(args) == regular_prompt


@pytest.mark.fast
def test_prompt_stdin_with_whitespace(mflux_generate_parser):
    stdin_content = "\n\n   A prompt with whitespace   \n\n"
    expected_prompt = "A prompt with whitespace"

    with patch("sys.stdin", StringIO(stdin_content)):
        with patch("sys.argv", ["mflux-generate", "--prompt", "-", "--model", "dev"]):
            args = mflux_generate_parser.parse_args()
            effective_prompt = PromptUtil.read_prompt(args)
            assert effective_prompt == expected_prompt


@pytest.mark.fast
def test_prompt_file_takes_precedence_over_stdin(mflux_generate_parser, temp_output_dir):
    # Create a prompt file
    prompt_file = temp_output_dir / "prompt.txt"
    file_prompt = "Prompt from file"
    prompt_file.write_text(file_prompt)
    stdin_content = "This should not be used because --prompt is not provided in the command."

    with patch("sys.stdin", StringIO(stdin_content)):
        with patch("sys.argv", ["mflux-generate", "--prompt-file", str(prompt_file), "--model", "dev"]):
            args = mflux_generate_parser.parse_args()
            effective_prompt = PromptUtil.read_prompt(args)
            assert effective_prompt == file_prompt