copykar/services/ai_processor.py

import os
import logging
import json
from dataclasses import dataclass
from typing import List, Optional, Dict, Any
from db.models import Post, TargetChannel
from db.repository import Repository
from core.llm import LLMClient
from core.metrics import DUPLICATES_DETECTED_TOTAL
from core.error_logger import log_exception

logger = logging.getLogger(__name__)

@dataclass
class RewriteResult:
    decision: str = "accept"  # "accept" or "reject"
    rejection_reason: str = ""
    rewritten_text: str = ""

    @property
    def is_rejected(self) -> bool:
        return self.decision.lower() == "reject"

    def __str__(self) -> str:
        return self.rewritten_text


SUPPORTED_LANGUAGES = {
    "fa": {
        "name": "Persian / Farsi",
        "instruction": "Translate or rewrite into natural, highly engaging, and fluent Persian (فارسی روان، جذاب و حرفه‌ای).",
        "default_personality": "لحن رسمی، جذاب و روان به همراه ایموجی‌های مرتبط و پاراگراف‌بندی مرتب.",
        "label": "فارسی 🇮🇷",
    },
    "en": {
        "name": "English",
        "instruction": "Translate or rewrite into natural, highly engaging, and fluent English.",
        "default_personality": "Professional, engaging, and clear tone with appropriate emojis and clean paragraph formatting.",
        "label": "English 🇬🇧",
    },
    "es": {
        "name": "Spanish / Español",
        "instruction": "Translate or rewrite into natural, highly engaging, and fluent Spanish (español natural, fluido y profesional).",
        "default_personality": "Tono profesional, atractivo y fluido con emojis apropiados y párrafos claros.",
        "label": "Español 🇪🇸",
    },
    "ar": {
        "name": "Arabic / العربية",
        "instruction": "Translate or rewrite into natural, highly engaging, and fluent Arabic (عربي فصيح وسلس وجذاب).",
        "default_personality": "أسلوب مهني وجذاب وسلس مع إيموجي مناسبة وفقرات واضحة.",
        "label": "العربية 🇸🇦",
    },
    "tr": {
        "name": "Turkish / Türkçe",
        "instruction": "Translate or rewrite into natural, highly engaging, and fluent Turkish (akıcı, doğal ve ilgi çekici Türkçe).",
        "default_personality": "İlgili emojiler ve düzenli paragraflarla profesyonel, akıcı ve ilgi çekici bir ton.",
        "label": "Türkçe 🇹🇷",
    },
}

# The prompt embeds a literal JSON example, so the placeholders are substituted by name
# rather than through str.format(), which would try to read the braces as fields.
CHANNEL_REWRITE_SYSTEM_PROMPT = """
You are a professional __LANGUAGE_NAME__ Telegram copywriter and content adapter.
Your job is to transform and rewrite raw posts specifically to match the target channel's personality and tone: "__CHANNEL_TITLE__".

TARGET CHANNEL DESIRED OUTPUT PERSONALITY & TONE:
__PERSONALITY__
__CUSTOM_INSTRUCTIONS_BLOCK__

IMPORTANT RULES FOR PERSONALITY & TONE:
- You MUST actively, thoroughly rewrite and restyle the post into the requested personality and tone.
- If the personality requires satire, humor, intense sarcasm/slang, formal journalism, or energetic marketing, YOU MUST WRITE IN THAT EXACT VOICE.
- Do NOT output bland, neutral, or unstyled text.
- Do NOT reject incoming posts because their original style, length, or tone is different from this personality. Your task is to transform ANY valid content into this target voice.

CRITICAL RULES:
1. Completely REMOVE all original channel usernames (e.g. @source_channel), sponsor tags, author watermarks, and source links.
2. __LANGUAGE_INSTRUCTION__
3. Rewrite and format the post to fully embody the target personality with appropriate emojis and clear paragraph spacing.
4. If a custom footer/tag is provided below, append it cleanly at the very end of the post:
__CUSTOM_FOOTER__
5. TELEGRAM CHARACTER LIMIT CONSTRAINT:
__LENGTH_LIMIT_RULE__
6. POST EVALUATION & REJECTION CRITERIA:
- Default decision is "accept". You should adapt and rewrite normal posts even if their original tone or subject is diverse.
- You should ONLY reject a post if it is:
  * Pure spam, unrelated scam/gambling/phishing ads
  * Completely empty, corrupt, or meaningless text
  * Explicitly forbidden by specific negative constraints in the custom commands above
- If REJECTING: set "decision": "reject" and specify the exact reason in "rejection_reason" (in Persian, e.g. "تبلیغات نامرتبط و اسپم").
- If ACCEPTING (Default): set "decision": "accept", set "rejection_reason": "", and provide the rewritten output in "rewritten_text".

Respond ONLY in valid JSON format:
{
  "decision": "accept",
  "rejection_reason": "",
  "rewritten_text": "..."
}
"""



def build_rewrite_system_prompt(
    channel_title: str,
    personality: str,
    custom_footer: str,
    custom_prompt: str = "",
    language: str = "fa",
    has_media: bool = False,
) -> str:
    lang_info = SUPPORTED_LANGUAGES.get(language, SUPPORTED_LANGUAGES["fa"])
    if has_media:
        length_rule = "This message contains media (photo/video/document). The total output (including footer) MUST NOT exceed 1000 characters to strictly fit Telegram's 1024-character caption limit."
    else:
        length_rule = "This is a text-only message. The total output (including footer) MUST NOT exceed 4000 characters to strictly fit Telegram's 4096-character message limit."

    custom_block = ""
    if custom_prompt and custom_prompt.strip():
        custom_block = f"\nTARGET CHANNEL SPECIFIC COMMANDS & FORMATTING RULES:\n{custom_prompt.strip()}\n"

    return (
        CHANNEL_REWRITE_SYSTEM_PROMPT
        .replace("__CHANNEL_TITLE__", channel_title)
        .replace("__LANGUAGE_NAME__", lang_info["name"])
        .replace("__LANGUAGE_INSTRUCTION__", lang_info["instruction"])
        .replace("__PERSONALITY__", personality)
        .replace("__CUSTOM_INSTRUCTIONS_BLOCK__", custom_block)
        .replace("__CUSTOM_FOOTER__", custom_footer)
        .replace("__LENGTH_LIMIT_RULE__", length_rule)
    )


VERIFY_REWRITE_SYSTEM_PROMPT = """
You are a senior editor and quality-assurance reviewer for Telegram channels.
Your task is to double-check and clean up an AI-rewritten post for the target channel: "__CHANNEL_TITLE__".

TARGET CHANNEL PERSONALITY & TONE (MUST BE PRESERVED):
__PERSONALITY__
__CUSTOM_INSTRUCTIONS_BLOCK__

TARGET LANGUAGE REQUIREMENT:
The ENTIRE post MUST be written 100% in __LANGUAGE_NAME__ (__LANGUAGE_INSTRUCTION__).
NO mixed languages, NO foreign sentences/phrases accidentally left over from other languages, and NO unwanted translations.

QUALITY & SAFETY AUDIT CHECKLIST:
1. PRESERVE THE PERSONALITY & HUMOR: Do NOT sanitize, formalize, or flatten the draft! Keep the exact tone, humor, satire, slang, and stylistic persona of the draft intact.
2. Fix Language Leaks: If any foreign words or mixed language leaked into the draft, translate them into __LANGUAGE_NAME__ while strictly preserving the persona and tone.
3. Remove any surviving source channel tags, usernames (@...), sponsor links, or URLs from the original post.
4. Keep the custom footer intact if present:
__CUSTOM_FOOTER__
5. TELEGRAM CHARACTER LIMIT CONSTRAINT:
__LENGTH_LIMIT_RULE__

Respond ONLY in valid JSON format:
{
  "final_text": "..."
}
"""


def build_verify_system_prompt(
    channel_title: str,
    personality: str,
    custom_footer: str,
    custom_prompt: str = "",
    language: str = "fa",
    has_media: bool = False,
) -> str:
    lang_info = SUPPORTED_LANGUAGES.get(language, SUPPORTED_LANGUAGES["fa"])
    if has_media:
        length_rule = "This message contains media (photo/video/document). The total output (including footer) MUST NOT exceed 1000 characters to strictly fit Telegram's 1024-character caption limit."
    else:
        length_rule = "This is a text-only message. The total output (including footer) MUST NOT exceed 4000 characters to strictly fit Telegram's 4096-character message limit."

    custom_block = ""
    if custom_prompt and custom_prompt.strip():
        custom_block = f"\nTARGET CHANNEL SPECIFIC COMMANDS & FORMATTING RULES:\n{custom_prompt.strip()}\n"

    return (
        VERIFY_REWRITE_SYSTEM_PROMPT
        .replace("__CHANNEL_TITLE__", channel_title)
        .replace("__LANGUAGE_NAME__", lang_info["name"])
        .replace("__LANGUAGE_INSTRUCTION__", lang_info["instruction"])
        .replace("__PERSONALITY__", personality)
        .replace("__CUSTOM_INSTRUCTIONS_BLOCK__", custom_block)
        .replace("__CUSTOM_FOOTER__", custom_footer)
        .replace("__LENGTH_LIMIT_RULE__", length_rule)
    )


class AIProcessor:
    def __init__(self, repo: Repository, llm: Optional[LLMClient] = None, double_check: Optional[bool] = None):
        self.repo = repo
        self.llm = llm or LLMClient(repo=self.repo)
        self.explicit_double_check = double_check
        if double_check is not None:
            self.double_check = double_check
        else:
            self.double_check = os.getenv("AI_DOUBLE_CHECK", "false").lower() in ("true", "1", "yes")

    async def _is_double_check_enabled(self) -> bool:
        if self.explicit_double_check is not None:
            return self.explicit_double_check
        if self.repo:
            db_dc = await self.repo.get_setting("ai_double_check")
            if db_dc is not None:
                return db_dc.strip().lower() in ("true", "1", "yes")
        return self.double_check

    async def _verify_rewrite(self, raw_text: str, draft_text: str, target: TargetChannel, has_media: bool = False) -> str:
        """Second-pass quality check to fix language mixing, tags, and length constraints without losing personality."""
        lang = getattr(target, "language", "fa") or "fa"
        lang_info = SUPPORTED_LANGUAGES.get(lang, SUPPORTED_LANGUAGES["fa"])
        personality_text = target.personality.strip() if target.personality else lang_info["default_personality"]
        footer_text = target.custom_footer.strip() if target.custom_footer else (f"@{target.username}" if target.username else "")
        custom_prompt_text = getattr(target, "custom_prompt", "") or ""

        sys_prompt = build_verify_system_prompt(
            channel_title=target.title or "کانال تلگرام",
            personality=personality_text,
            custom_footer=footer_text,
            custom_prompt=custom_prompt_text,
            language=lang,
            has_media=has_media,
        )


        user_prompt = (
            f"متن خام اولیه:\n{raw_text}\n\n"
            f"پیش‌نویس اولیه بازنویسی‌شده برای بررسی و اصلاح زبان/فرمت:\n{draft_text}"
        )
        try:
            res = await self.llm.generate_json(
                prompt=user_prompt,
                system_prompt=sys_prompt,
                action_name="verify_target_post"
            )
            final_text = res.get("final_text")
            if final_text and final_text.strip():
                return final_text.strip()
        except Exception as e:
            logger.warning(f"Double-check verification failed, keeping initial draft: {e}")
        return draft_text

    async def rewrite_for_target(
        self,
        raw_text: str,
        target: TargetChannel,
        has_media: bool = False,
        image_path: Optional[str] = None,
    ) -> RewriteResult:
        """Rewrite raw text according to target channel's language, personality, custom prompt commands, length limits, and optional image."""
        if not raw_text and not image_path:
            return RewriteResult(decision="reject", rejection_reason="متن پیام و تصویر هر دو خالی هستند", rewritten_text="")

        lang = getattr(target, "language", "fa") or "fa"
        lang_info = SUPPORTED_LANGUAGES.get(lang, SUPPORTED_LANGUAGES["fa"])
        personality_text = target.personality.strip() if target.personality else lang_info["default_personality"]
        footer_text = target.custom_footer.strip() if target.custom_footer else (f"@{target.username}" if target.username else "")
        custom_prompt_text = getattr(target, "custom_prompt", "") or ""

        sys_prompt = build_rewrite_system_prompt(
            channel_title=target.title or "کانال تلگرام",
            personality=personality_text,
            custom_footer=footer_text,
            custom_prompt=custom_prompt_text,
            language=lang,
            has_media=has_media,
        )

        user_prompt_text = f"متن اصلی پست برای بازنویسی و تبدیل به لحن و استایل کانال مقصد:\n\n{raw_text}" if raw_text else "لطفاً با توجه به تصویر پیوست، یک متن جذاب و مناسب برای کانال بنویسید."

        try:
            res = await self.llm.generate_json(
                prompt=user_prompt_text,
                system_prompt=sys_prompt,
                action_name="rewrite_target_post",
                image_path=image_path
            )
            decision = str(res.get("decision", "accept")).strip().lower()
            rejection_reason = str(res.get("rejection_reason", "")).strip()
            rewritten = res.get("rewritten_text") or res.get("text") or ""

            if decision == "reject":
                return RewriteResult(
                    decision="reject",
                    rejection_reason=rejection_reason or "رد شده طبق ارزیابی هوش مصنوعی",
                    rewritten_text=rewritten.strip() if rewritten else ""
                )

            if rewritten:
                rewritten = rewritten.strip()
                if await self._is_double_check_enabled():
                    rewritten = await self._verify_rewrite(raw_text, rewritten, target, has_media=has_media)
                return RewriteResult(
                    decision="accept",
                    rejection_reason="",
                    rewritten_text=rewritten
                )
        except Exception as e:
            await log_exception("ai_processor.rewrite", e, {"target_id": target.id, "target_title": target.title})

        # Fallback if AI fails: clean basic @mentions and append footer
        fallback = raw_text
        if footer_text:
            fallback = f"{fallback}\n\n{footer_text}"
        return RewriteResult(decision="accept", rejection_reason="", rewritten_text=fallback)