import re
import logging
from typing import Optional, List, Dict, Any

logger = logging.getLogger(__name__)

class TextProcessor:
    """大模型文本处理工具类"""
    
    @staticmethod
    def clean_model_output(output: str) -> str:
        """清理模型输出文本
        
        Args:
            output: 模型输出的原始文本
            
        Returns:
            str: 清理后的文本
        """
        try:
            # 移除引用标记
            output = re.sub(r'\[[^\]]*\]', '', output)
            
            # 移除多余的空白字符
            output = re.sub(r'\s+', ' ', output)
            
            # 移除首尾的空白字符
            output = output.strip()
            
            # 移除可能存在的HTML标签
            output = re.sub(r'<[^>]+>', '', output)
            
            # 移除可能存在的Markdown格式
            output = re.sub(r'[*_`~]', '', output)
            
            return output
            
        except Exception as e:
            logger.error(f"清理模型输出失败: {str(e)}")
            return output

    @staticmethod
    def clean_thought_process(output: str) -> str:
        """清理模型输出，移除推理过程，只保留最终结果

        Args:
            output: 模型原始输出文本

        Returns:
            str: 清理后的输出文本
        """
        try:
            # 找到</think>标签的位置
            think_end = output.find('</think>')
            if think_end != -1:
                # 移除</think>标签及其之前的所有内容
                output = output[think_end + len('</think>'):]

            # 处理可能存在的空行
            lines = output.split('\n')
            cleaned_lines = []
            for line in lines:
                line = line.strip()
                if line:  # 只保留非空行
                    cleaned_lines.append(line)

            # 重新组合文本
            output = '\n'.join(cleaned_lines)

            return output.strip()

        except Exception as e:
            logger.error(f"清理模型输出失败: {str(e)}")
            return output.strip()

    @staticmethod
    def extract_numeric_value_from_response(response: str) -> str:
        """从模型响应中提取数值
        
        Args:
            response: 模型响应的文本
            
        Returns:
            str: 提取出的数值
        """
        try:
            # 清理响应文本
            cleaned_response = TextProcessor.clean_model_output(response)
            
            # 尝试提取数值
            numeric_patterns = [
                r'[-+]?\d*\.\d+',  # 浮点数
                r'[-+]?\d+',       # 整数
                r'[-+]?\d+\.\d*'   # 可能带小数点的数
            ]
            
            for pattern in numeric_patterns:
                matches = re.findall(pattern, cleaned_response)
                if matches:
                    # 返回第一个匹配的数值
                    return matches[0]
            
            # 如果没有找到数值，返回原始响应
            return cleaned_response
            
        except Exception as e:
            logger.error(f"提取数值失败: {str(e)}")
            return response
    
    @staticmethod
    def extract_list_from_response(response: str) -> List[str]:
        """从模型响应中提取列表项
        
        Args:
            response: 模型响应的文本
            
        Returns:
            List[str]: 提取出的列表项
        """
        try:
            # 清理响应文本
            cleaned_response = TextProcessor.clean_model_output(response)
            
            # 尝试匹配列表项
            # 匹配数字编号的列表项
            numbered_items = re.findall(r'\d+\.\s*([^\n]+)', cleaned_response)
            if numbered_items:
                return [item.strip() for item in numbered_items]
            
            # 匹配带符号的列表项
            bullet_items = re.findall(r'[-*•]\s*([^\n]+)', cleaned_response)
            if bullet_items:
                return [item.strip() for item in bullet_items]
            
            # 如果没有找到列表项，返回空列表
            return []
            
        except Exception as e:
            logger.error(f"提取列表项失败: {str(e)}")
            return []
    
    @staticmethod
    def extract_key_value_pairs(response: str) -> Dict[str, str]:
        """从模型响应中提取键值对
        
        Args:
            response: 模型响应的文本
            
        Returns:
            Dict[str, str]: 提取出的键值对
        """
        try:
            # 清理响应文本
            cleaned_response = TextProcessor.clean_model_output(response)
            
            # 尝试匹配键值对
            # 匹配冒号分隔的键值对
            pairs = re.findall(r'([^:]+):\s*([^\n]+)', cleaned_response)
            if pairs:
                return {key.strip(): value.strip() for key, value in pairs}
            
            # 匹配等号分隔的键值对
            pairs = re.findall(r'([^=]+)=\s*([^\n]+)', cleaned_response)
            if pairs:
                return {key.strip(): value.strip() for key, value in pairs}
            
            # 如果没有找到键值对，返回空字典
            return {}
            
        except Exception as e:
            logger.error(f"提取键值对失败: {str(e)}")
            return {} 