from typing import Optional, List, Set,Iterable,Any
try:
    import ahocorasick
except ImportError:
    ahocorasick = None
import re
import json
import logging

try:
    from langchain_core.prompts import FewShotPromptTemplate, PromptTemplate
except ImportError:
    try:
        from langchain.prompts import FewShotPromptTemplate, PromptTemplate
    except ImportError:
        class PromptTemplate:
            def __init__(self, input_variables=None, template: str = ""):
                self.input_variables = input_variables or []
                self.template = template

            def format(self, **kwargs):
                return self.template.format(**kwargs)

        class FewShotPromptTemplate:
            def __init__(
                self,
                examples=None,
                example_prompt=None,
                prefix: str = "",
                suffix: str = "",
                input_variables=None,
                example_separator: str = "\n\n",
            ):
                self.examples = examples or []
                self.example_prompt = example_prompt
                self.prefix = prefix
                self.suffix = suffix
                self.input_variables = input_variables or []
                self.example_separator = example_separator

            def format(self, **kwargs):
                rendered_examples = []
                for example in self.examples:
                    rendered_examples.append(self.example_prompt.format(**example))
                parts = [self.prefix]
                if rendered_examples:
                    parts.append(self.example_separator.join(rendered_examples))
                parts.append(self.suffix.format(**kwargs))
                return self.example_separator.join(part for part in parts if part)

from openai.types.chat import ChatCompletionSystemMessageParam, ChatCompletionUserMessageParam, ChatCompletionMessageParam
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


def cotprompt(query):
    # 精简 Few-Shot：只示范「输入文本 -> 结构化修改结果」，不再包含推理过程/思考链，
    # 从源头杜绝模型输出解释、分析、括号备注等多余内容，同时大幅缩短每次请求的前缀长度。
    examples = [
        {
            # 综合示例：覆盖错别字/错别词语/标点不规范/重复内容，
            # 重点示范：多余标点(。。。)、全半角括号 -> 标点不规范（不是逻辑不通）
            "question": "现,代教育中，智能教学系统【】让每个雪生受益，自动文达系统很方便。现就安全生产大检查工作提出如下意见。。。(本报记者 小陈)",
            "answer": (
                "###\n"
                "错误：标点不规范\n原文：现,代教育\n建议：现代教育\n\n"
                "错误：标点不规范\n原文：智能教学系统【】\n建议：智能教学系统\n\n"
                "错误：错别字\n原文：每个雪生\n建议：每个学生\n\n"
                "错误：错别词语\n原文：自动文达系统\n建议：自动问答系统\n\n"
                "错误：标点不规范\n原文：如下意见。。。\n建议：如下意见。\n\n"
                "错误：标点不规范\n原文：(本报记者 小陈)\n建议：（本报记者 小陈）\n"
                "###"
            ),
        },
        {
            # 重复内容的原文必须带唯一上下文，建议返回删除后的完整片段。
            "question": "2026年通知：2025年通知内容从下述角度开战：一、项目推进方面。各部门需加快项目进度，确保按时交付。\n\n二、培训学习方面：公司将组织系列培训课程，提升员工专业技能。具体时间、地点另行通知。员工需按时参加，不得无辜缺席。自即日起，严禁在工作时间从事与工作无关的活动。\n\n二、纪律要求。近期发现部分员工上班时间做与工作无关的事情，如炒股、聊天等，严重影响工作效率和公司形象。自即日起，严禁在工作时间从事与工作无关的活动。违反规定者，将按照公司规章制度严肃处理。\n\n特此通知：",
            "answer": (
                "###\n"
                "错误：逻辑错误\n"
                "原文：2026年通知：2025年通知内容从下述角度开战：\n"
                "建议：2026年通知：2026年通知内容从下述角度开展：\n\n"
                "错误：标点不规范\n"
                "原文：一、项目推进方面。\n"
                "建议：一、项目推进方面：\n\n"
                "错误：错别词语、重复内容\n"
                "原文：培训学习方面：公司将组织系列培训课程，提升员工专业技能。具体时间、地点另行通知。员工需按时参加，不得无辜缺席。自即日起，严禁在工作时间从事与工作无关的活动。\n"
                "建议：培训学习方面：公司将组织系列培训课程，提升员工专业技能。具体时间、地点另行通知。员工需按时参加，不得无故缺席。\n\n"
                "错误：格式不规范\n"
                "原文：二、纪律要求。\n"
                "建议：三、纪律要求：\n\n"
                "错误：标点不规范\n"
                "原文：特此通知：\n"
                "建议：特此通知。\n"
                "###"
            ),
        },
        {
            # 无明显错误示例：只返回空标记，禁止输出任何说明、分析或括号备注
            "question": "全年审核采购合同、对账单据数百份，均按规定完成登记归档。",
            "answer": "###\n###",
        },
    ]

    # 单个示例模板：只保留 问题/答案，删除“推理过程”
    example_prompt = PromptTemplate(
        input_variables=["question", "answer"],
        template="问题：{question}\n答案：{answer}",
    )

    few_shot_prompt = FewShotPromptTemplate(
        examples=examples,
        example_prompt=example_prompt,
        prefix=(
            "你是专业文本校对专家，负责检查错别字、错别词语、错误标点、标题格式与序号、重复内容、逻辑不通及合规问题。\n"
            "输出规则（务必严格遵守）：\n"
            "1. 每个错误必须连续占三行：第一行“错误：<错误类型>”，第二行“原文：<有错的原文片段>”，第三行“建议：<修改后的正确文本>”。三行之间不得插入空行；不同错误之间空一行。整体用 ### 包裹。\n"
            "2. 错误类型只能从以下选择：错别字、错别词语、标点不规范、格式不规范、重复内容、逻辑不通、合规问题。同一连续原文片段存在修改范围重叠的多类错误时，必须合并为一条，错误类型用中文顿号连接，并在一条建议中一次改正。标题序号重复、跳号、顺序错误或层级不一致归为“格式不规范”；标题末尾的纯标点问题归为“标点不规范”。\n"
            "3. 凡是标点问题（多余或重复标点如“。。。”、全角/半角标点混用如英文括号()应为中文（）、标点缺失或误用），一律归为“标点不规范”，禁止归为“逻辑不通”；“逻辑不通”只用于前后文语义矛盾、指代不清等真正的逻辑问题。\n"
            "4. 建议字段只写修改后的正确内容，禁止输出任何解释、分析、推理、评论或括号备注（例如禁止出现“（删除该句……）”“（此处无明显错误……）”“（通常……可接受）”这类内容）。\n"
            "5. 判定重复内容时，必须结合全文语义判断是否造成冗余，不能仅凭文字相同判错。若两处分别承担必要的说明、强调、引用或衔接作用，则不得输出校对项。\n"
            "6. 确认属于重复内容后，必须比较每次出现位置与标题主题、前后文和表达作用的匹配程度，保留主题最匹配、逻辑最完整的一处，禁止机械删除第一次或最后一次。\n"
            "7. 删除局部重复内容时，原文必须包含足以唯一定位的连续上下文；建议必须返回删除重复内容后的同一完整片段，不得留空。只有整个原文片段均为应删除的纯冗余内容时，建议才可留空。\n"
            "8. 必须按错别字和错别词语、标点、标题格式与序号、重复内容、逻辑、合规的顺序逐项检查全文。发现某类错误后仍须继续检查其他类别。\n"
            "9. 若整段文本没有任何错误，只返回：###\n###，不要输出其它任何字符。\n"
            "请参照以下示例完成校对："
        ),
        suffix=(
            "问题：{question}\n"
            "输出答案前必须在内部完成以下扫描，但不要输出扫描过程：\n"
            "A. 逐字检查错别字和错别词语；\n"
            "B. 逐行检查标题及结尾标点，同级标题格式必须一致；标题后直接接正文时检查是否应使用冒号；\n"
            "C. 按正文顺序检查各级标题序号是否重复、跳号、倒序或层级错误；\n"
            "D. 全文检索重复句子和连续文本，结合标题主题判断是否冗余；\n"
            "E. 继续检查逻辑与合规问题。\n"
            "必须完成 A 至 E 后再输出全部问题，不得提前结束或限制问题数量。\n"
            "每条原文必须从本次输入正文中逐字、连续、原样复制，严禁提前采用修改结果。建议只能写可直接替换的最终正文。\n"
            "答案："
        ),
        input_variables=["question"],
        example_separator="\n\n",
    )
    return few_shot_prompt.format(question=query)


def build_review_prompt(types: List[str], content: str, require: Optional[str]) -> Iterable[ChatCompletionMessageParam]:
    if require is None:
        require = ""
    require = require.strip()
    review_t = "现在需要你帮我完成以下文本校对任务："
    for i , t in enumerate(types):
        if t =="逻辑校对":
            review_t += f"{i} 对这段文本进行逻辑校对，主要是内容中词语或者其他内容本身是个合规词语，但在文本用词不当，不符合当前内容中逻辑，例如：#2026年终总结：2025年对本公司的业务开战总结，正确的内容应该为：#2026年终总结：2026年对本公司的业务开展总结\n"
        if t =="基础校对":
            review_t += f"{i}.1 进行错别字校对，错别字错误主要以文本用字不当为主，例如：星光店电，正确的应该为：星光点点。\n"
            review_t += f"{i}.2 进行标点校对，主要以标点符号使用不当为主，例如：星光点点；万里无云，正确的应该为：星光点点，万里无云\n"
            review_t += f"{i}.3 进行格式规范校对\n"
            review_t += f"{i}.4 进行重复内容校对，主要以原文中出现重复词语、句子或文本为主，例如：“今天今天心情相当不错，我很开心。我很开心。“，正确的内容应该为：今天心情相当不错，我很开心。\n"
        if t=="合规性检查":
            review_t += f"{i} 对这段文本进行合规性检查\n"
    require_text = f"\n额外要求：{require}\n" if require else ""
    query = f"""
任务要求：
{review_t}
{require_text}
注意：只输出结构化校对结果（错误/原文/建议），每条结果的错误、原文、建议必须连续三行，三行之间不得有空行；建议只写修改后的正确文本；没有错误时只返回 ###\n###。
必须逐项检查错别字和错别词语、标点、标题格式与序号、重复内容、逻辑及合规。重复内容必须结合标题和上下文判断；删除局部重复内容时，原文必须带唯一上下文，建议必须返回删除后的完整上下文。同一片段中修改范围重叠的多类错误必须合并为一条。

待校对正文开始：
{content}
待校对正文结束。
    """.strip()
    combined_prompt = cotprompt(query)
    return [
        ChatCompletionUserMessageParam(role="user", content=combined_prompt),
    ]



def length_convert(length: Optional[str]):
    """将篇幅解析为 {target, max}：target 为建议目标字数，max 为硬上限。"""
    if length == "短":
        return {"target": 700, "max": 1000}
    if length == "中":
        return {"target": 2500, "max": 3000}
    if length == "长":
        return {"target": 4500, "max": 5000}
    # 兜底：从自定义篇幅字符串里解析出最大字数
    if length:
        nums = re.findall(r"\d+", length)
        if nums:
            max_chars = int(nums[-1])
            return {"target": int(max_chars * 0.8), "max": max_chars}
    return None


def length_display(length: Optional[str]) -> str:
    """模板 {length} 占位符用的字符串。"""
    info = length_convert(length)
    if info:
        return f"{info['target']}字左右（不超过{info['max']}字）"
    return length or ""


def build_length_instruction(length: Optional[str]) -> str:
    """把 target/max 写进提示词：建议目标 + 硬上限 + 优先级。"""
    info = length_convert(length)
    if not info:
        return ""
    return (
        f"全文建议写到{info['target']}字左右，绝对不得超过{info['max']}字；"
        "字数为硬性要求，不得注水或重复凑字数。"
    )




def normalize_words(words: List[str]) -> Set[str]:
    return {word.strip() for word in words if word and word.strip()}


def find_keywords(text: str, keyword_set: Set[str]) -> List[str]:
    if not text or not keyword_set:
        return []

    if ahocorasick is None:
        return [keyword for keyword in keyword_set if keyword in text]

    automaton = ahocorasick.Automaton()

    for idx, keyword in enumerate(keyword_set):
        automaton.add_word(keyword, (idx, keyword))

    automaton.make_automaton()

    found = set()
    for _, (_, keyword) in automaton.iter(text):
        found.add(keyword)

    return list(found)


def append_dictionary_results(
    items: List[dict],
    content: str,
    sensitive_words: Set[str],
    negative_words: Set[str],
) -> List[dict]:
    for word in find_keywords(content, sensitive_words):
        items.append({
            "original": word,
            "error": "敏感词汇",
            "suggestion": "",
        })

    for word in find_keywords(content, negative_words):
        items.append({
            "original": word,
            "error": "错误词汇",
            "suggestion": "",
        })

    return items


def filter_positive_word_false_positives(
    items: List[dict],
    positive_words: Set[str],
) -> List[dict]:
    if not positive_words:
        return items

    result = []

    for item in items:
        original = str(item.get("original", ""))
        error = str(item.get("error", ""))

        has_positive_word = any(word in original for word in positive_words)
        is_typo_error = (
            "错别字" in error
            or "错别词" in error
            or "错别词语" in error
        )

        if has_positive_word and is_typo_error:
            continue

        result.append(item)

    return result


def dedupe_and_filter(items: List[dict], source_text: str) -> List[dict]:
    seen = set()
    result = []

    for item in items:
        original = str(item.get("original", "")).strip()
        error = str(item.get("error", "")).strip()
        suggestion = str(item.get("suggestion", "")).strip()

        if not original:
            continue

        if original == suggestion:
            continue

        if original not in source_text:
            continue

        key = (original, error, suggestion)
        if key in seen:
            continue

        seen.add(key)
        result.append({
            "original": original,
            "error": error,
            "suggestion": suggestion,
        })

    return result



# =============== 辅助函数 ===============
def process_rag_prompt(rag_prompt: Any) -> str:
    """
    处理 rag_prompt 参数，将各种类型转换为字符串

    Args:
        rag_prompt: 可以是字符串、列表、字典等任意类型

    Returns:
        处理后的字符串
    """
    if rag_prompt is None:
        return ""

    # 如果是字符串，直接返回
    if isinstance(rag_prompt, str):
        return rag_prompt.strip()

    # 如果是列表，转换为多行字符串
    if isinstance(rag_prompt, list):
        # 将列表中的每个元素转换为字符串，并用换行符连接
        return "\n".join(str(item) for item in rag_prompt if item)

    # 如果是字典，转换为 JSON 字符串
    if isinstance(rag_prompt, dict):
        try:
            return json.dumps(rag_prompt, ensure_ascii=False, indent=2)
        except Exception as e:
            logger.warning(f"rag_prompt 字典转换失败: {e}")
            return str(rag_prompt)

    # 其他类型，直接转换为字符串
    return str(rag_prompt).strip()
