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."
    ),
}

REFINE_SYSTEM_PROMPTS = {
    "t2i": (
        "You are editing an existing image prompt. The previous prompt describes "
        "the current image. Apply only the user's requested changes. Keep all "
        "subjects, style, colors, composition and constraints from the previous "
        "prompt unless the user explicitly asks to change them. Return only the "
        "final English prompt, without quotes or explanations."
    ),
    "i2i": (
        "You are updating an image-editing instruction. The previous prompt "
        "describes the current image. Apply only the user's requested changes. "
        "Keep the existing composition, subjects, style, lighting and constraints "
        "unless the user explicitly asks to change them. 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,
        *,
        previous_prompt: str | None = None,
        base_prompt: str | None = None,
    ) -> str:
        if not self.enabled:
            return prompt

        user_prompt = self._sanitize(prompt)
        if not user_prompt:
            return prompt

        previous = self._sanitize(previous_prompt or "")
        base = self._sanitize(base_prompt or "")
        if previous:
            chat_prompt = self._refine_chat_prompt(
                user_prompt,
                mode,
                previous_prompt=previous,
                base_prompt=base or previous,
            )
        else:
            chat_prompt = self._chat_prompt(user_prompt, mode)

        workflow = copy.deepcopy(self._template)
        workflow[self._text_node_id]["inputs"]["prompt"] = chat_prompt
        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 _sanitize(text: str) -> str:
        value = text.strip()
        for token in ("<|im_start|>", "<|im_end|>"):
            value = value.replace(token, "")
        return value

    @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 _refine_chat_prompt(
        prompt: str,
        mode: str,
        *,
        previous_prompt: str,
        base_prompt: str,
    ) -> str:
        system = REFINE_SYSTEM_PROMPTS.get(
            mode, REFINE_SYSTEM_PROMPTS["t2i"]
        )
        return (
            f"<|im_start|>system\n{system}<|im_end|>\n"
            "<|im_start|>user\n"
            f"Base scene: {base_prompt}\n"
            f"Current prompt: {previous_prompt}\n"
            f"Requested change: {prompt}\n"
            "<|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