from __future__ import annotations

from pathlib import Path

from config import Settings
from prompt_enhancer import PromptEnhancer

BASE = Path(__file__).resolve().parents[1] / "workflows"


class FakeEnhanceComfy:
    def __init__(self, text: str = "A red apple on a wooden table, soft light") -> None:
        self.text = text
        self.workflow: dict | None = None
        self.submits = 0

    async def submit(self, workflow: dict, client_id: str) -> str:
        self.workflow = workflow
        self.submits += 1
        return "prompt-1"

    async def wait_for_history(self, prompt_id: str, *, timeout: float) -> dict:
        return {"outputs": {"3": {"text": [self.text]}}}

    def find_output_text(self, entry: dict) -> str | None:
        return entry["outputs"]["3"]["text"][0]


def make_settings(enabled: bool = True) -> Settings:
    return Settings(
        bot_token="",
        prompt_enhance_enabled=enabled,
        workflow_prompt=BASE / "prompt_enhance.json",
    )


async def test_enhance_t2i_prompt() -> None:
    comfy = FakeEnhanceComfy()
    enhancer = PromptEnhancer(comfy, make_settings())

    result = await enhancer.enhance("красное яблоко на столе, мягкий свет", "t2i")

    assert result == "A red apple on a wooden table, soft light"
    assert comfy.submits == 1
    assert comfy.workflow is not None
    prompt = comfy.workflow["2"]["inputs"]["prompt"]
    assert "<|im_start|>system" in prompt
    assert "image generation model" in prompt
    assert "красное яблоко" in prompt
    assert prompt.endswith("<|im_start|>assistant\n")


async def test_enhance_i2i_prompt() -> None:
    comfy = FakeEnhanceComfy("Make the apple blue, keep the composition")
    enhancer = PromptEnhancer(comfy, make_settings())

    result = await enhancer.enhance("Сделай яблоко синим", "i2i")

    assert result == "Make the apple blue, keep the composition"
    assert comfy.workflow is not None
    prompt = comfy.workflow["2"]["inputs"]["prompt"]
    assert "image-editing prompt translator" in prompt


async def test_enhance_with_context_uses_refine_prompt() -> None:
    comfy = FakeEnhanceComfy("A blue apple on a wooden table, soft light")
    enhancer = PromptEnhancer(comfy, make_settings())

    result = await enhancer.enhance(
        "Сделай яблоко синим",
        "t2i",
        previous_prompt="A red apple on a wooden table, soft light",
        base_prompt="A red apple on a wooden table",
    )

    assert result == "A blue apple on a wooden table, soft light"
    assert comfy.workflow is not None
    prompt = comfy.workflow["2"]["inputs"]["prompt"]
    assert "editing an existing image prompt" in prompt
    assert "A red apple on a wooden table, soft light" in prompt
    assert "Requested change: Сделай яблоко синим" in prompt


async def test_disabled_returns_original_without_comfy_call() -> None:
    comfy = FakeEnhanceComfy()
    enhancer = PromptEnhancer(comfy, make_settings(enabled=False))

    result = await enhancer.enhance("красное яблоко", "t2i")

    assert result == "красное яблоко"
    assert comfy.submits == 0


async def test_clean_strips_prefix_and_empty_output() -> None:
    assert PromptEnhancer._clean('Prompt: "A red apple"') == "A red apple"
    assert PromptEnhancer._clean("user") == ""