"""
PromptEngineer — enriches user prompts with heritage style keywords.

SDXL-aware: builds long prompts (~75-150 tokens) that exploit SDXL's
two-encoder architecture (G+XL text encoders). Negative prompts are
also style-aware.
"""
from __future__ import annotations

import logging
from typing import List

from .base import AgentClient, AgentResponse
from config.styles import StyleSpec

log = logging.getLogger(__name__)

SYSTEM_PROMPT = """You are an expert prompt engineer for Stable Diffusion XL, specialized in Indian heritage art.

Your task: take a user's plain description and enrich it with concrete visual keywords
for the chosen heritage style. Aim for 60-120 words. Be specific about:
- Composition (perspective, framing)
- Color palette (named pigments, not hex codes)
- Linework and brushwork quality
- Iconographic motifs authentic to the style
- Border / frame treatment
- Negative prompt should list things to avoid (modern, photographic, 3D, etc.)

Output STRICT JSON:
{"prompt": "<enriched prompt 60-120 words>",
 "negative_prompt": "<negative prompt 20-40 words>",
 "style_summary": "<3-word summary>"}"""


class PromptEngineer:
    def __init__(self, client: AgentClient | None = None) -> None:
        self.client = client or AgentClient(temperature=0.6, max_tokens=900)

    def enrich(self, user_prompt: str, style: StyleSpec) -> AgentResponse:
        """Return enriched prompt + negative + summary."""
        if self.client.enabled:
            user_msg = (
                f"User prompt: {user_prompt!r}\n"
                f"Style: {style.id} ({style.display_name})\n"
                f"Style region: {style.region}\n"
                f"Authentic motifs: {', '.join(style.cultural_keywords)}\n"
                f"Reference palette: {', '.join(style.palette)}"
            )
            resp = self.client.chat(system_prompt=SYSTEM_PROMPT, user_prompt=user_msg)
            if resp.ok:
                resp.content = self._extract_prompt_text(resp.content, user_prompt, style)
                return resp

        # Heuristic fallback
        return AgentResponse(
            content=self._heuristic_enrich(user_prompt, style),
            model=self.client.model,
            ok=True,
            error="heuristic_fallback",
        )

    @staticmethod
    def _extract_prompt_text(raw: str, user_prompt: str, style: StyleSpec) -> str:
        """Parse agent JSON; if parsing fails, return a heuristic fallback."""
        import json, re
        try:
            data = json.loads(raw)
            return data.get("prompt", raw)
        except Exception:
            m = re.search(r'"prompt"\s*:\s*"([^"]+)"', raw)
            if m:
                return m.group(1).replace("\\n", " ").replace('\\"', '"')
            return PromptEngineer._heuristic_enrich(user_prompt, style)

    @staticmethod
    def build_negative(style: StyleSpec, extra: List[str] | None = None) -> str:
        """Compose the negative prompt for a style."""
        base = list(style.negative_tags)
        # SDXL-universal negatives
        universal = [
            "low quality", "jpeg artifacts", "watermark", "signature",
            "text", "logo", "cropped", "out of frame", "duplicate",
            "extra limbs", "deformed", "blurry",
        ]
        parts = base + universal
        if extra:
            parts.extend(extra)
        return ", ".join(parts)

    @staticmethod
    def _heuristic_enrich(user_prompt: str, style: StyleSpec) -> str:
        tags = ", ".join(style.prompt_tags)
        return (
            f"{user_prompt.strip().rstrip('.')}, {tags}, "
            "intricate detail, traditional composition, "
            "museum-quality heritage artwork, high resolution"
        )