# llm_enhancer.py
from transformers import pipeline
import torch
from typing import List, Dict
import re


class LLMEnhancer:
    """LLM增强功能类"""

    def __init__(self, model_name="facebook/bart-large-cnn", device=None):
        if device is None:
            device = "cuda" if torch.cuda.is_available() else "cpu"

        print(f"加载LLM模型到 {device}...")
        self.device = device

        # 摘要模型
        self.summarizer = pipeline(
            "summarization",
            model=model_name,
            device=0 if device == "cuda" else -1
        )

        # QA模型（可选）
        try:
            self.qa_model = pipeline(
                "question-answering",
                model="distilbert-base-cased-distilled-squad",
                device=0 if device == "cuda" else -1
            )
        except:
            print("未加载QA模型")
            self.qa_model = None

    def generate_summary(self, text: str, max_length: int = 150, min_length: int = 50):
        """生成文本摘要"""
        if len(text) < 100:
            return text

        # 截断文本以适应模型限制
        if len(text) > 1024:
            text = text[:1024]

        try:
            summary = self.summarizer(
                text,
                max_length=max_length,
                min_length=min_length,
                do_sample=False
            )
            return summary[0]['summary_text']
        except Exception as e:
            print(f"摘要生成失败: {e}")
            return text[:max_length] + "..."

    def answer_question(self, context: str, question: str):
        """基于上下文回答问题"""
        if not self.qa_model:
            return "QA功能未启用"

        try:
            result = self.qa_model(question=question, context=context)
            return result['answer']
        except Exception as e:
            print(f"问答失败: {e}")
            return "无法回答问题"

    def extract_key_points(self, text: str, num_points: int = 5):
        """提取关键点"""
        prompt = f"""
        请从以下文本中提取{num_points}个关键点:

        文本: {text[:2000]}

        关键点:
        1. """

        # 这里可以调用本地LLM或API
        # 简化版：使用规则提取
        sentences = re.split(r'[.!?]+', text)
        key_sentences = []

        for sentence in sentences:
            sentence = sentence.strip()
            if len(sentence) > 50 and any(keyword in sentence.lower() for keyword in
                                          ['important', 'significant', 'found', 'showed', 'demonstrated']):
                key_sentences.append(sentence)
                if len(key_sentences) >= num_points:
                    break

        return key_sentences