from __future__ import annotations

import copy
import json
import logging
import re
import uuid
from pathlib import Path
from typing import Any

from comfy_client import ComfyClient
from config import Settings

logger = logging.getLogger(__name__)

SYSTEM_PROMPTS = {
    "t2i": (
        "You are a prompt engineer for an image generation model. "
        "Translate the user's request into English and rewrite it as a concise, "
        "descriptive image prompt. Preserve every requested subject, action, style, "
        "color, composition and constraint. Do not invent new objects, camera angles "
        "or details. Return only the final English prompt, without quotes or "
        "explanations."
    ),
    "i2i": (
        "You are an image-editing prompt translator. Translate the user's editing "
        "instruction into English and rewrite it clearly. Preserve the requested "
        "changes, composition, pose, lighting and all constraints. Do not add new "
        "objects, camera angles or details. Return only the final English "
        "instruction, without quotes or explanations."
    ),
}


class PromptEnhanceError(RuntimeError):
    pass


class PromptEnhancer:
    def __init__(self, comfy: ComfyClient, settings: Settings) -> None:
        self._comfy = comfy
        self._settings = settings
        self._template = self._load_template(settings.workflow_prompt)
        self._text_node_id = self._find_text_node(self._template)
        if not any(
            node.get("class_type") == "PreviewAny"
            for node in self._template.values()
            if isinstance(node, dict)
        ):
            raise PromptEnhanceError(
                "В workflow улучшения промпта нет узла PreviewAny"
            )

    @property
    def enabled(self) -> bool:
        return self._settings.prompt_enhance_enabled

    @staticmethod
    def _load_template(path: Path) -> dict[str, Any]:
        try:
            data = json.loads(path.read_text(encoding="utf-8"))
        except FileNotFoundError as exc:
            raise PromptEnhanceError(f"Не найден workflow: {path}") from exc
        except json.JSONDecodeError as exc:
            raise PromptEnhanceError(f"Некорректный JSON в {path}: {exc}") from exc
        if not isinstance(data, dict) or not data:
            raise PromptEnhanceError("Workflow улучшения промпта пустой")
        return data

    @staticmethod
    def _find_text_node(workflow: dict[str, Any]) -> str:
        node_ids = [
            node_id
            for node_id, node in workflow.items()
            if isinstance(node, dict) and node.get("class_type") == "TextGenerate"
        ]
        if len(node_ids) != 1:
            raise PromptEnhanceError(
                "В workflow улучшения промпта должен быть один TextGenerate"
            )
        return str(node_ids[0])

    async def enhance(self, prompt: str, mode: str) -> str:
        if not self.enabled:
            return prompt

        user_prompt = prompt.strip()
        for token in ("<|im_start|>", "<|im_end|>"):
            user_prompt = user_prompt.replace(token, "")
        if not user_prompt:
            return prompt

        workflow = copy.deepcopy(self._template)
        workflow[self._text_node_id]["inputs"]["prompt"] = self._chat_prompt(
            user_prompt, mode
        )
        prompt_id = await self._comfy.submit(workflow, uuid.uuid4().hex)
        entry = await self._comfy.wait_for_history(
            prompt_id,
            timeout=self._settings.prompt_enhance_timeout,
        )
        text = self._comfy.find_output_text(entry)
        if not text:
            raise PromptEnhanceError("ComfyUI вернул пустой промпт")

        cleaned = self._clean(text)
        if not cleaned:
            raise PromptEnhanceError("Не удалось очистить сгенерированный промпт")
        return cleaned[: self._settings.max_prompt_length]

    @staticmethod
    def _chat_prompt(prompt: str, mode: str) -> str:
        system = SYSTEM_PROMPTS.get(mode, SYSTEM_PROMPTS["t2i"])
        return (
            f"<|im_start|>system\n{system}<|im_end|>\n"
            f"<|im_start|>user\n{prompt}<|im_end|>\n"
            "<|im_start|>assistant\n"
        )

    @staticmethod
    def _clean(text: str) -> str:
        text = " ".join(text.split()).strip()
        for token in ("<|im_end|>", "<|im_start|>"):
            text = text.split(token)[0].strip()
        text = text.strip("`").strip()
        for prefix in ("English prompt:", "Prompt:", "Assistant:"):
            if text.lower().startswith(prefix.lower()):
                text = text[len(prefix):].strip()
        text = text.strip(" \"'")
        text = text.lstrip(" ,.;:-")
        if text.lower() in {"user", "assistant", "system"}:
            return ""
        if not re.search(r"[A-Za-z]", text):
            return ""
        return text