#!/usr/bin/env python3
"""
Excel和CSV文件操作MCP服务器
支持Excel/CSV文件的读取、搜索、枚举值获取等操作
参考multi_keyword_search_server.py的实现方式
"""

import json
import os
import sys
import asyncio
import re
import chardet
from typing import Any, Dict, List, Optional, Union
import pandas as pd
from mcp_common import (
    get_allowed_directory, 
    load_tools_from_json, 
    resolve_file_path,
    find_file_in_project,
    is_regex_pattern,
    compile_pattern,
    create_error_response,
    create_success_response,
    create_initialize_response,
    create_ping_response,
    create_tools_list_response,
    handle_mcp_streaming
)


def detect_encoding(file_path: str) -> str:
    """检测文件编码"""
    try:
        with open(file_path, 'rb') as f:
            raw_data = f.read(10000)  # 读取前10KB来检测编码
            result = chardet.detect(raw_data)
            return result['encoding'] or 'utf-8'
    except:
        return 'utf-8'




class ExcelCSVOperator:
    """Excel和CSV文件操作核心类"""
    
    def __init__(self):
        self.supported_extensions = ['.xlsx', '.xls', '.csv']
        self.encoding_cache = {}
    
    def _validate_file(self, file_path: str) -> str:
        """验证并处理文件路径"""
        # 解析文件路径，支持 folder/document.txt 和 document.txt 格式
        resolved_path = resolve_file_path(file_path)
        
        # 验证文件扩展名
        file_ext = os.path.splitext(resolved_path)[1].lower()
        if file_ext not in self.supported_extensions:
            raise ValueError(f"Unsupported file format: {file_ext}, supported formats: {self.supported_extensions}")
        
        return resolved_path
    
    def load_data(self, file_path: str, sheet_name: str = None) -> pd.DataFrame:
        """加载Excel或CSV文件数据"""
        file_path = self._validate_file(file_path)
        file_ext = os.path.splitext(file_path)[1].lower()
        
        try:
            if file_ext == '.csv':
                encoding = detect_encoding(file_path)
                df = pd.read_csv(file_path, encoding=encoding)
            else:
                # Excel文件
                if sheet_name:
                    df = pd.read_excel(file_path, sheet_name=sheet_name)
                else:
                    # 读取第一个sheet
                    df = pd.read_excel(file_path)
            
            # 处理空值
            df = df.fillna('')
            
            return df
        except Exception as e:
            raise ValueError(f"File loading failed: {str(e)}")
    
    def get_sheets(self, file_path: str) -> List[str]:
        """获取Excel文件的所有sheet名称"""
        file_path = self._validate_file(file_path)
        file_ext = os.path.splitext(file_path)[1].lower()
        
        if file_ext == '.csv':
            return ['default']  # CSV文件只有一个默认sheet
        
        try:
            excel_file = pd.ExcelFile(file_path)
            return excel_file.sheet_names
        except Exception as e:
            raise ValueError(f"Failed to read Excel sheet list: {str(e)}")
    
    def get_schema(self, file_path: str, sheet_name: str = None) -> List[str]:
        """获取文件的schema字段列表"""
        try:
            df = self.load_data(file_path, sheet_name)
            return df.columns.tolist()
        except Exception as e:
            raise ValueError(f"Failed to get schema: {str(e)}")
    
    def full_text_search(self, file_path: str, keywords: List[str], 
                        top_k: int = 10, case_sensitive: bool = False) -> str:
        """全文搜索功能"""
        if not keywords:
            return "Error: Keyword list cannot be empty"
        
        # 预处理和验证关键词中的正则表达式
        valid_keywords = []
        regex_errors = []
        
        for keyword in keywords:
            compiled = compile_pattern(keyword)
            if compiled is None:
                regex_errors.append(keyword)
            else:
                valid_keywords.append(keyword)
        
        if regex_errors:
            error_msg = f"Warning: The following regular expressions failed to compile and will be ignored: {', '.join(regex_errors)}"
            print(error_msg)
        
        if not valid_keywords:
            return "Error: No valid search keywords"
        
        try:
            # 验证文件路径
            validated_path = self._validate_file(file_path)
            file_ext = os.path.splitext(validated_path)[1].lower()
            
            all_results = []
            
            if file_ext == '.csv':
                # CSV文件只有一个数据集
                results = self._search_in_file(validated_path, valid_keywords, case_sensitive, 'default')
                all_results.extend(results)
            else:
                # Excel文件，搜索所有sheet
                sheets = self.get_sheets(validated_path)
                for sheet in sheets:
                    results = self._search_in_file(validated_path, valid_keywords, case_sensitive, sheet)
                    all_results.extend(results)
            
            # 按匹配数量排序（降序）
            all_results.sort(key=lambda x: x['match_count'], reverse=True)
            
            # 限制结果数量
            limited_results = all_results[:top_k]
            
            # 格式化为CSV输出
            if not limited_results:
                return "No matching results found"
            
            # 构建CSV格式输出
            csv_lines = []
            headers = ["sheet", "row_index", "match_count", "matched_content", "match_details"]
            csv_lines.append(",".join(headers))
            
            for result in limited_results:
                # 转义CSV中的特殊字符
                sheet = str(result.get('sheet', '')).replace(',', '，')
                row_index = str(result.get('row_index', ''))
                match_count = str(result.get('match_count', 0))
                matched_content = str(result.get('matched_content', '')).replace(',', '，').replace('\n', ' ')
                match_details = str(result.get('match_details', '')).replace(',', '，')
                
                csv_lines.append(f"{sheet},{row_index},{match_count},{matched_content},{match_details}")
            
            return "\n".join(csv_lines)
            
        except Exception as e:
            return f"Search failed: {str(e)}"
    
    def _search_in_file(self, file_path: str, keywords: List[str], 
                      case_sensitive: bool, sheet_name: str = None) -> List[Dict[str, Any]]:
        """在文件中搜索关键词"""
        results = []
        
        try:
            df = self.load_data(file_path, sheet_name)
            
            # 预处理所有模式
            processed_patterns = []
            for keyword in keywords:
                compiled = compile_pattern(keyword)
                if compiled is not None:
                    processed_patterns.append({
                        'original': keyword,
                        'pattern': compiled,
                        'is_regex': isinstance(compiled, re.Pattern)
                    })
            
            # 逐行搜索
            for row_index, row in df.iterrows():
                # 将整行内容合并为字符串进行搜索
                row_content = " ".join([str(cell) for cell in row.values if str(cell).strip()])
                search_content = row_content if case_sensitive else row_content.lower()
                
                # 统计匹配的模式数量
                matched_patterns = []
                for pattern_info in processed_patterns:
                    pattern = pattern_info['pattern']
                    is_regex = pattern_info['is_regex']
                    
                    match_found = False
                    match_details = None
                    
                    if is_regex:
                        # 正则表达式匹配
                        if case_sensitive:
                            match = pattern.search(row_content)
                        else:
                            # 对于不区分大小写的正则，需要重新编译
                            if isinstance(pattern, re.Pattern):
                                flags = pattern.flags | re.IGNORECASE
                                case_insensitive_pattern = re.compile(pattern.pattern, flags)
                                match = case_insensitive_pattern.search(row_content)
                            else:
                                match = pattern.search(search_content)
                        
                        if match:
                            match_found = True
                            match_details = match.group(0)
                    else:
                        # 普通字符串匹配
                        search_keyword = pattern if case_sensitive else pattern.lower()
                        if search_keyword in search_content:
                            match_found = True
                            match_details = pattern
                    
                    if match_found:
                        matched_patterns.append({
                            'original': pattern_info['original'],
                            'type': 'regex' if is_regex else 'keyword',
                            'match': match_details
                        })
                
                match_count = len(matched_patterns)
                
                if match_count > 0:
                    # 构建匹配详情
                    match_details = []
                    for pattern in matched_patterns:
                        if pattern['type'] == 'regex':
                            match_details.append(f"[regex:{pattern['original']}={pattern['match']}]")
                        else:
                            match_details.append(f"[keyword:{pattern['match']}]")
                    
                    match_info = " ".join(match_details)
                    
                    results.append({
                        'sheet': sheet_name,
                        'row_index': row_index,
                        'match_count': match_count,
                        'matched_content': row_content,
                        'match_details': match_info
                    })
        
        except Exception as e:
            print(f"Error searching file {file_path} (sheet: {sheet_name}): {str(e)}")
        
        return results
    
    def filter_search(self, file_path: str, filters: Dict, 
                     sheet_name: str = None) -> str:
        """字段过滤搜索功能"""
        if not filters:
            return "Error: Filter conditions cannot be empty"
        
        try:
            df = self.load_data(file_path, sheet_name)
            
            # 应用过滤条件
            filtered_df = df.copy()
            
            for field_name, filter_condition in filters.items():
                if field_name not in df.columns:
                    return f"Error: Field '{field_name}' does not exist"
                
                operator = filter_condition.get('operator', 'eq')
                value = filter_condition.get('value')
                
                if operator == 'eq':
                    # 等于
                    filtered_df = filtered_df[filtered_df[field_name] == value]
                elif operator == 'gt':
                    # 大于
                    filtered_df = filtered_df[filtered_df[field_name] > value]
                elif operator == 'lt':
                    # 小于
                    filtered_df = filtered_df[filtered_df[field_name] < value]
                elif operator == 'gte':
                    # 大于等于
                    filtered_df = filtered_df[filtered_df[field_name] >= value]
                elif operator == 'lte':
                    # 小于等于
                    filtered_df = filtered_df[filtered_df[field_name] <= value]
                elif operator == 'contains':
                    # 包含
                    filtered_df = filtered_df[filtered_df[field_name].astype(str).str.contains(str(value), na=False)]
                elif operator == 'regex':
                    # 正则表达式
                    try:
                        pattern = re.compile(str(value))
                        filtered_df = filtered_df[filtered_df[field_name].astype(str).str.match(pattern, na=False)]
                    except re.error as e:
                        return f"Error: Regular expression '{value}' compilation failed: {str(e)}"
                else:
                    return f"Error: Unsupported operator '{operator}'"
            
            # 格式化为CSV输出
            if filtered_df.empty:
                return "No records matching conditions found"
            
            # 转换为CSV字符串
            csv_result = filtered_df.to_csv(index=False, encoding='utf-8')
            return csv_result
            
        except Exception as e:
            return f"Filter search failed: {str(e)}"
    
    def get_field_enums(self, file_path: str, field_names: List[str], 
                       sheet_name: str = None, max_enum_count: int = 100,
                       min_occurrence: int = 1) -> str:
        """获取指定字段的枚举值列表"""
        if not field_names:
            return "Error: Field name list cannot be empty"
        
        try:
            df = self.load_data(file_path, sheet_name)
            
            # 验证字段存在性
            missing_fields = [field for field in field_names if field not in df.columns]
            if missing_fields:
                return f"Error: Fields do not exist: {', '.join(missing_fields)}"
            
            # 计算每个字段的枚举值
            enum_results = {}
            
            for field in field_names:
                # 统计值出现次数
                value_counts = df[field].value_counts()
                
                # 过滤出现次数过少的值
                filtered_counts = value_counts[value_counts >= min_occurrence]
                
                # 限制返回数量
                top_values = filtered_counts.head(max_enum_count)
                
                # 格式化结果
                enum_values = []
                for value, count in top_values.items():
                    enum_values.append(f"{value}({count})")
                
                enum_results[field] = {
                    'enum_values': enum_values,
                    'total_unique': len(value_counts),
                    'total_filtered': len(filtered_counts),
                    'total_rows': len(df)
                }
            
            # 格式化输出
            output_lines = []
            for field, data in enum_results.items():
                enum_str = ", ".join(data['enum_values'])
                field_info = f"{field}: [{enum_str}] (总计: {data['total_unique']}个唯一值, 过滤后: {data['total_filtered']}个, 总行数: {data['total_rows']})"
                output_lines.append(field_info)
            
            return "\n".join(output_lines)
            
        except Exception as e:
            return f"Failed to get enum values: {str(e)}"


# 全局操作器实例
operator = ExcelCSVOperator()


async def handle_request(request: Dict[str, Any]) -> Dict[str, Any]:
    """Handle MCP request"""
    try:
        method = request.get("method")
        params = request.get("params", {})
        request_id = request.get("id")
        
        if method == "initialize":
            return create_initialize_response(request_id, "excel-csv-operator")
        
        elif method == "ping":
            return create_ping_response(request_id)
        
        elif method == "tools/list":
            # 从 JSON 文件加载工具定义
            tools = load_tools_from_json("excel_csv_operator_tools.json")
            return create_tools_list_response(request_id, tools)
        
        elif method == "tools/call":
            tool_name = params.get("name")
            arguments = params.get("arguments", {})
            
            if tool_name == "get_excel_sheets":
                file_path = arguments.get("file_path")
                result = operator.get_sheets(file_path)
                
                return create_success_response(request_id, {
                    "content": [
                        {
                            "type": "text",
                            "text": json.dumps(result, ensure_ascii=False, indent=2)
                        }
                    ]
                })
            
            elif tool_name == "get_table_schema":
                file_path = arguments.get("file_path")
                sheet_name = arguments.get("sheet_name")
                result = operator.get_schema(file_path, sheet_name)
                
                return create_success_response(request_id, {
                    "content": [
                        {
                            "type": "text",
                            "text": json.dumps(result, ensure_ascii=False, indent=2)
                        }
                    ]
                })
            
            elif tool_name == "full_text_search":
                file_path = arguments.get("file_path")
                keywords = arguments.get("keywords", [])
                top_k = arguments.get("top_k", 10)
                case_sensitive = arguments.get("case_sensitive", False)
                
                result = operator.full_text_search(file_path, keywords, top_k, case_sensitive)
                
                return create_success_response(request_id, {
                    "content": [
                        {
                            "type": "text",
                            "text": result
                        }
                    ]
                })
            
            elif tool_name == "filter_search":
                file_path = arguments.get("file_path")
                sheet_name = arguments.get("sheet_name")
                filters = arguments.get("filters")
                
                result = operator.filter_search(file_path, filters, sheet_name)
                
                return create_success_response(request_id, {
                    "content": [
                        {
                            "type": "text",
                            "text": result
                        }
                    ]
                })
            
            elif tool_name == "get_field_enums":
                file_path = arguments.get("file_path")
                sheet_name = arguments.get("sheet_name")
                field_names = arguments.get("field_names", [])
                max_enum_count = arguments.get("max_enum_count", 100)
                min_occurrence = arguments.get("min_occurrence", 1)
                
                result = operator.get_field_enums(file_path, field_names, sheet_name, max_enum_count, min_occurrence)
                
                return create_success_response(request_id, {
                    "content": [
                        {
                            "type": "text",
                            "text": result
                        }
                    ]
                })
            
            else:
                return create_error_response(request_id, -32601, f"Unknown tool: {tool_name}")
        
        else:
            return create_error_response(request_id, -32601, f"Unknown method: {method}")
    
    except Exception as e:
        return create_error_response(request.get("id"), -32603, f"Internal error: {str(e)}")


async def main():
    """Main entry point."""
    await handle_mcp_streaming(handle_request)


if __name__ == "__main__":
    asyncio.run(main())
