#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Capricode Vision Pro - 修复版代码视觉定位系统
"""
import gradio as gr
import re
from typing import Dict, List, Tuple, Any, Optional
from dataclasses import dataclass
import time

@dataclass
class QuantumPosition:
    """量子级精准位置定位"""
    absolute_line: int
    relative_line: int
    column: int
    scope_depth: int
    context_hash: str
    semantic_position: str

class PrecisionCodeVision:
    def __init__(self):
        self.position_cache = {}
        
    def quantum_analyze(self, code: str, cursor_line: int, cursor_column: int) -> Dict[str, Any]:
        """
        量子级精准代码分析 - 修复测试发现的问题
        """
        lines = code.split('\n')
        
        # 边界检查增强
        if cursor_line < 0 or cursor_line >= len(lines) or cursor_column < 0:
            return self._handle_invalid_position(cursor_line, cursor_column)
        
        # 1. 绝对位置校准
        absolute_position = self._calibrate_absolute_position(lines, cursor_line, cursor_column)
        
        # 2. 语义位置分析
        semantic_position = self._analyze_semantic_position(code, cursor_line, cursor_column)
        
        # 3. 作用域精准定位 - 修复作用域边界问题
        scope_context = self._precise_scope_analysis(lines, cursor_line, cursor_column)
        
        # 4. 代码结构映射
        structure_map = self._create_structure_mapping(code, lines)
        
        # 5. 修复建议定位
        repair_locations = self._calculate_repair_locations(structure_map, absolute_position, lines, cursor_line)
        
        return {
            "quantum_position": absolute_position,
            "semantic_context": semantic_position,
            "scope_precision": scope_context,
            "structure_mapping": structure_map,
            "repair_targets": repair_locations,
            "confidence_score": self._calculate_confidence(lines, cursor_line, cursor_column)
        }
    
    def _handle_invalid_position(self, line: int, column: int) -> Dict[str, Any]:
        """处理无效位置"""
        return {
            "quantum_position": {
                "absolute_line": max(0, line),
                "relative_line": 0,
                "column": max(0, column),
                "scope_depth": 0,
                "context_hash": "invalid_position",
                "semantic_position": "无效位置"
            },
            "semantic_context": {
                "current_token": "",
                "statement_type": "invalid",
                "is_inside_string": False,
                "is_inside_comment": False,
                "surrounding_tokens": []
            },
            "scope_precision": {
                "scope_type": "invalid_scope",
                "scope_boundary": {"start": 0, "end": 0},
                "accessible_variables": [],
                "scope_hierarchy": []
            },
            "structure_mapping": {"syntactic_elements": {}},
            "repair_targets": [],
            "confidence_score": 0.1
        }
    
    def _calibrate_absolute_position(self, lines: List[str], line: int, column: int) -> QuantumPosition:
        """绝对位置校准 - 修复边界处理"""
        if line >= len(lines):
            line = len(lines) - 1
        if column >= len(lines[line]):
            column = len(lines[line]) - 1
            
        current_line = lines[line]
        context_hash = self._generate_position_hash(lines, line, column)
        
        return QuantumPosition(
            absolute_line=line + 1,
            relative_line=self._calculate_relative_line(lines, line),
            column=column + 1,
            scope_depth=self._calculate_scope_depth(lines, line),
            context_hash=context_hash,
            semantic_position=self._get_semantic_label(lines, line, column)
        )
    
    def _calculate_relative_line(self, lines: List[str], line: int) -> int:
        """计算相对行号 - 修复算法"""
        scope_start = self._find_scope_start(lines, line)
        return line - scope_start + 1
    
    def _calculate_scope_depth(self, lines: List[str], line: int) -> int:
        """计算作用域深度 - 修复深度计算"""
        depth = 0
        brace_count = 0
        for i in range(line + 1):
            current_line = lines[i]
            # 修复：只计算实际的作用域嵌套，不考虑控制流
            if re.search(r'\b(function|class)\b', current_line):
                brace_count += current_line.count('{')
                brace_count -= current_line.count('}')
                depth = brace_count
            elif current_line.strip().startswith(('if ', 'for ', 'while ')):
                # 控制流不增加深度，只在内部增加
                pass
        return max(0, depth)
    
    def _generate_position_hash(self, lines: List[str], line: int, column: int) -> str:
        """生成位置哈希"""
        context = ""
        start = max(0, line - 2)
        end = min(len(lines), line + 3)
        
        for i in range(start, end):
            context += lines[i].strip() + "|"
        
        context += f"L{line}C{column}"
        return str(hash(context))
    
    def _get_semantic_label(self, lines: List[str], line: int, column: int) -> str:
        """获取语义标签 - 增强识别精度"""
        if line >= len(lines):
            return "无效位置"
            
        current_line = lines[line]
        
        # 修复：更精确的语义识别
        if self._is_inside_comment(current_line, column):
            return "注释内"
        elif self._is_inside_string(current_line, column):
            return "字符串内"
        elif re.search(r'\bfunction\b', current_line):
            return "函数定义内"
        elif re.search(r'\bclass\b', current_line):
            return "类定义内"
        elif any(re.search(rf'\b{keyword}\b', current_line) for keyword in ['if', 'for', 'while']):
            return "控制流语句内"
        elif any(keyword in current_line for keyword in ['import', 'from']):
            return "导入语句"
        elif any(keyword in current_line for keyword in ['const', 'let', 'var']):
            return "变量声明"
        elif '=' in current_line and '==' not in current_line:
            return "赋值语句"
        else:
            return "代码语句内"
    
    def _analyze_semantic_position(self, code: str, line: int, column: int) -> Dict[str, Any]:
        """语义位置分析"""
        lines = code.split('\n')
        if line >= len(lines):
            return {
                "current_token": "",
                "statement_type": "invalid",
                "is_inside_string": False,
                "is_inside_comment": False,
                "surrounding_tokens": []
            }
            
        current_line = lines[line]
        
        return {
            "current_token": self._extract_current_token(current_line, column),
            "statement_type": self._classify_statement_type(current_line),
            "is_inside_string": self._is_inside_string(current_line, column),
            "is_inside_comment": self._is_inside_comment(current_line, column),
            "surrounding_tokens": self._get_surrounding_tokens(lines, line, column),
        }
    
    def _precise_scope_analysis(self, lines: List[str], line: int, column: int) -> Dict[str, Any]:
        """精准作用域分析 - 修复作用域边界问题"""
        if line >= len(lines):
            return {
                "scope_type": "invalid_scope",
                "scope_boundary": {"start": 0, "end": 0},
                "accessible_variables": [],
                "scope_hierarchy": []
            }
            
        # 修复：使用新的作用域查找算法
        scope_start, scope_end, scope_type = self._find_precise_scope(lines, line)
        
        return {
            "scope_type": scope_type,
            "scope_boundary": {
                "start": scope_start + 1,
                "end": scope_end + 1,
            },
            "accessible_variables": self._find_accessible_variables(lines, scope_start, scope_end, line),
            "scope_hierarchy": self._get_scope_hierarchy(lines, line),
        }
    
    def _find_precise_scope(self, lines: List[str], line: int) -> Tuple[int, int, str]:
        """精确查找作用域 - 修复边界计算"""
        # 首先查找当前行的直接作用域
        current_scope_start = self._find_scope_start(lines, line)
        current_scope_type = self._detect_scope_type(lines, current_scope_start)
        
        # 修复：根据作用域类型确定边界
        if current_scope_type == "class_scope":
            # 类作用域应该只包含类定义内的内容
            scope_end = self._find_class_end(lines, current_scope_start)
        elif current_scope_type == "function_scope":
            # 函数作用域
            scope_end = self._find_function_end(lines, current_scope_start)
        elif current_scope_type == "control_flow_scope":
            # 控制流作用域
            scope_end = self._find_control_flow_end(lines, current_scope_start)
        else:
            # 全局作用域
            scope_end = len(lines) - 1
            
        return current_scope_start, scope_end, current_scope_type
    
    def _find_class_end(self, lines: List[str], start_line: int) -> int:
        """查找类定义结束 - 修复类边界"""
        indent_level = len(lines[start_line]) - len(lines[start_line].lstrip())
        
        for i in range(start_line + 1, len(lines)):
            current_indent = len(lines[i]) - len(lines[i].lstrip())
            # 如果缩进小于类定义的缩进，说明类结束
            if current_indent <= indent_level and lines[i].strip():
                return i - 1
                
        return len(lines) - 1
    
    def _find_function_end(self, lines: List[str], start_line: int) -> int:
        """查找函数定义结束"""
        brace_count = 0
        for i in range(start_line, len(lines)):
            brace_count += lines[i].count('{')
            brace_count -= lines[i].count('}')
            if brace_count < 0:
                return i
        return len(lines) - 1
    
    def _find_control_flow_end(self, lines: List[str], start_line: int) -> int:
        """查找控制流结束"""
        indent_level = len(lines[start_line]) - len(lines[start_line].lstrip())
        
        for i in range(start_line + 1, len(lines)):
            current_indent = len(lines[i]) - len(lines[i].lstrip())
            if current_indent <= indent_level and lines[i].strip():
                return i - 1
                
        return len(lines) - 1
    
    def _extract_current_token(self, line: str, column: int) -> str:
        """提取当前token - 修复算法"""
        if column >= len(line) or column < 0:
            return ""
        
        # 修复：使用正则表达式精确提取token
        start = column
        while start > 0 and (line[start-1].isalnum() or line[start-1] in '_.'):
            start -= 1
            
        end = column
        while end < len(line) and (line[end].isalnum() or line[end] in '_.'):
            end += 1
            
        return line[start:end] if start < end else ""
    
    def _classify_statement_type(self, line: str) -> str:
        """分类语句类型"""
        line = line.strip()
        if re.search(r'^def\s+', line):
            return "function_definition"
        elif re.search(r'^class\s+', line):
            return "class_definition"
        elif any(re.search(rf'^{keyword}\s+', line) for keyword in ['if', 'for', 'while']):
            return "control_flow"
        elif any(keyword in line for keyword in ['import', 'from']):
            return "import_statement"
        elif any(keyword in line for keyword in ['const', 'let', 'var']):
            return "variable_declaration"
        elif '=' in line and '==' not in line:
            return "assignment"
        else:
            return "expression"
    
    def _is_inside_string(self, line: str, column: int) -> bool:
        """检查是否在字符串内"""
        if column >= len(line):
            return False
            
        in_single_quote = False
        in_double_quote = False
        escape_next = False
        
        for i, char in enumerate(line):
            if i > column:
                break
                
            if escape_next:
                escape_next = False
                continue
                
            if char == '\\':
                escape_next = True
                continue
                
            if char == "'" and not in_double_quote:
                in_single_quote = not in_single_quote
            elif char == '"' and not in_single_quote:
                in_double_quote = not in_double_quote
                
        return in_single_quote or in_double_quote
    
    def _is_inside_comment(self, line: str, column: int) -> bool:
        """检查是否在注释内"""
        if column >= len(line):
            return False
            
        line_before_cursor = line[:column]
        return '//' in line_before_cursor or '/*' in line_before_cursor or line.strip().startswith('#')
    
    def _get_surrounding_tokens(self, lines: List[str], line: int, column: int) -> List[str]:
        """获取周围token"""
        tokens = []
        start = max(0, line - 1)
        end = min(len(lines), line + 2)
        
        for i in range(start, end):
            line_tokens = re.findall(r'[a-zA-Z_]\w*|[0-9]+|\S', lines[i])
            tokens.extend(line_tokens)
        
        return tokens[:15]
    
    def _find_scope_start(self, lines: List[str], line: int) -> int:
        """查找作用域开始 - 修复算法"""
        # 修复：避免控制流错误包含类作用域
        for i in range(line, -1, -1):
            current_line = lines[i]
            if re.search(r'\b(class|function)\b', current_line):
                return i
            elif any(re.search(rf'\b{keyword}\b', current_line) for keyword in ['if', 'for', 'while']):
                # 只在没有找到类/函数时才返回控制流
                found_higher = False
                for j in range(i-1, -1, -1):
                    if re.search(r'\b(class|function)\b', lines[j]):
                        found_higher = True
                        break
                if not found_higher:
                    return i
        return 0
    
    def _find_scope_end(self, lines: List[str], line: int) -> int:
        """查找作用域结束"""
        brace_count = 0
        for i in range(line, len(lines)):
            brace_count += lines[i].count('{')
            brace_count -= lines[i].count('}')
            if brace_count < 0:
                return i
        return len(lines) - 1
    
    def _detect_scope_type(self, lines: List[str], scope_start: int) -> str:
        """检测作用域类型"""
        if scope_start >= len(lines):
            return "global_scope"
            
        line = lines[scope_start]
        if re.search(r'\bfunction\b', line):
            return "function_scope"
        elif re.search(r'\bclass\b', line):
            return "class_scope"
        elif any(re.search(rf'\b{keyword}\b', line) for keyword in ['if', 'for', 'while']):
            return "control_flow_scope"
        else:
            return "global_scope"
    
    def _find_accessible_variables(self, lines: List[str], scope_start: int, scope_end: int, current_line: int) -> List[str]:
        """查找可访问变量 - 修复变量识别"""
        variables = []
        for i in range(scope_start, min(current_line + 1, scope_end + 1)):
            line = lines[i]
            # 修复：增强变量识别模式
            patterns = [
                r'(?:const|let|var)\s+([a-zA-Z_]\w*)',
                r'def\s+([a-zA-Z_]\w*)',
                r'class\s+([a-zA-Z_]\w*)',
                r'([a-zA-Z_]\w*)\s*=',
                r'self\.([a-zA-Z_]\w*)\s*='
            ]
            
            for pattern in patterns:
                matches = re.findall(pattern, line)
                variables.extend(matches)
                
        return list(set(variables))
    
    def _get_scope_hierarchy(self, lines: List[str], line: int) -> List[Dict[str, Any]]:
        """获取作用域层次结构 - 修复层级关系"""
        hierarchy = []
        current_line = line
        
        while current_line >= 0:
            scope_start = self._find_scope_start(lines, current_line)
            scope_type = self._detect_scope_type(lines, scope_start)
            
            # 修复：使用精确的作用域边界
            if scope_type == "class_scope":
                scope_end = self._find_class_end(lines, scope_start)
            elif scope_type == "function_scope":
                scope_end = self._find_function_end(lines, scope_start)
            elif scope_type == "control_flow_scope":
                scope_end = self._find_control_flow_end(lines, scope_start)
            else:
                scope_end = len(lines) - 1
            
            hierarchy.append({
                "type": scope_type,
                "start_line": scope_start + 1,
                "end_line": scope_end + 1
            })
            
            # 如果已经是全局作用域，停止回溯
            if scope_start == 0 or scope_type == "global_scope":
                break
                
            current_line = scope_start - 1
        
        return hierarchy[::-1]
    
    def _create_structure_mapping(self, code: str, lines: List[str]) -> Dict[str, Any]:
        """创建代码结构映射"""
        return {
            "syntactic_elements": {
                "functions": self._map_functions(lines),
                "classes": self._map_classes(lines),
                "control_flows": self._map_control_flows(lines),
                "variables": self._map_variables(lines)
            },
        }
    
    def _map_functions(self, lines: List[str]) -> List[Dict[str, Any]]:
        """映射函数"""
        functions = []
        for i, line in enumerate(lines):
            if re.search(r'\bfunction\b', line) or re.search(r'^def\s+', line):
                functions.append({
                    "name": self._extract_function_name(line),
                    "line": i + 1,
                    "signature": line.strip()
                })
        return functions
    
    def _extract_function_name(self, line: str) -> str:
        """提取函数名"""
        patterns = [
            r'function\s+([a-zA-Z_]\w*)\s*\(',
            r'def\s+([a-zA-Z_]\w*)\s*\('
        ]
        
        for pattern in patterns:
            match = re.search(pattern, line)
            if match:
                return match.group(1)
        return "anonymous"
    
    def _map_classes(self, lines: List[str]) -> List[Dict[str, Any]]:
        """映射类"""
        classes = []
        for i, line in enumerate(lines):
            if re.search(r'\bclass\b', line):
                match = re.search(r'class\s+([a-zA-Z_]\w*)', line)
                if match:
                    classes.append({
                        "name": match.group(1),
                        "line": i + 1
                    })
        return classes
    
    def _map_control_flows(self, lines: List[str]) -> List[Dict[str, Any]]:
        """映射控制流"""
        controls = []
        for i, line in enumerate(lines):
            if any(re.search(rf'\b{keyword}\b', line) for keyword in ['if', 'for', 'while']):
                controls.append({
                    "type": self._get_control_type(line),
                    "line": i + 1,
                    "condition": line.strip()
                })
        return controls
    
    def _get_control_type(self, line: str) -> str:
        """获取控制类型"""
        if re.search(r'\bif\b', line):
            return "if_statement"
        elif re.search(r'\bfor\b', line):
            return "for_loop"
        elif re.search(r'\bwhile\b', line):
            return "while_loop"
        else:
            return "unknown"
    
    def _map_variables(self, lines: List[str]) -> List[Dict[str, Any]]:
        """映射变量"""
        variables = []
        for i, line in enumerate(lines):
            patterns = [
                r'(?:const|let|var)\s+([a-zA-Z_]\w*)',
                r'([a-zA-Z_]\w*)\s*=',
                r'self\.([a-zA-Z_]\w*)\s*='
            ]
            
            for pattern in patterns:
                matches = re.findall(pattern, line)
                for match in matches:
                    variables.append({
                        "name": match,
                        "line": i + 1,
                    })
        return variables
    
    def _calculate_repair_locations(self, structure: Dict[str, Any], position: QuantumPosition, lines: List[str], cursor_line: int) -> List[Dict[str, Any]]:
        """计算修复位置"""
        locations = []
        
        semantic_pos = position.semantic_position
        
        if "导入语句" in semantic_pos:
            locations.append({
                "type": "import_insert",
                "line": position.absolute_line,
                "description": "在导入区域添加新导入",
                "priority": 1,
            })
        elif "函数" in semantic_pos:
            locations.append({
                "type": "function_insert",
                "line": position.absolute_line,
                "description": "在函数内插入代码",
                "priority": 1,
            })
        else:
            locations.append({
                "type": "line_insert_before",
                "line": position.absolute_line,
                "description": "在当前行前插入",
                "priority": 1,
            })
        
        locations.append({
            "type": "line_insert_after", 
            "line": position.absolute_line + 1,
            "description": "在当前行后插入",
            "priority": 2,
        })
        
        locations.append({
            "type": "scope_start",
            "line": position.absolute_line - position.relative_line + 1,
            "description": "在作用域开始处插入",
            "priority": 3,
        })
        
        return locations
    
    def _calculate_confidence(self, lines: List[str], cursor_line: int, cursor_column: int) -> float:
        """计算置信度 - 修复置信度计算"""
        if cursor_line >= len(lines) or cursor_line < 0:
            return 0.1
            
        line = lines[cursor_line]
        
        base_confidence = 0.7
        
        if cursor_column < len(line):
            token = self._extract_current_token(line, cursor_column)
            if token:
                base_confidence += 0.2
                
        complexity = len(line.split()) / 15.0
        base_confidence += min(0.25, complexity)
        
        return min(0.95, base_confidence)

# ==================== Gradio 界面 ====================

# 创建视觉引擎实例
vision_engine = PrecisionCodeVision()

def analyze_code_with_cursor(code, cursor_line, cursor_column):
    """分析代码和光标位置"""
    if not code.strip():
        return "请输入代码", {}, {}
    
    try:
        cursor_line = int(cursor_line) - 1
        cursor_column = int(cursor_column) - 1
    except:
        return "请输入有效的行号和列号", {}, {}
    
    start_time = time.time()
    result = vision_engine.quantum_analyze(code, cursor_line, cursor_column)
    processing_time = (time.time() - start_time) * 1000
    
    quantum_pos = result['quantum_position']
    scope_info = result['scope_precision']
    
    report = f"""
## 🎯 量子级精准分析报告 (修复版)

### 📍 位置信息
- **绝对位置**: 第 {quantum_pos.absolute_line} 行, 第 {quantum_pos.column} 列
- **语义位置**: {quantum_pos.semantic_position}
- **作用域深度**: {quantum_pos.scope_depth}
- **位置哈希**: `{quantum_pos.context_hash[:8]}...`

### 🎪 作用域分析
- **作用域类型**: {scope_info['scope_type']}
- **作用域范围**: 第 {scope_info['scope_boundary']['start']} 行 - 第 {scope_info['scope_boundary']['end']} 行
- **可用变量**: {', '.join(scope_info['accessible_variables']) if scope_info['accessible_variables'] else '无'}

### 🔧 修复建议
"""
    
    for location in result['repair_targets'][:3]:
        report += f"- **{location['description']}** (优先级: {location['priority']})\n"
    
    report += f"\n**处理时间**: {processing_time:.2f}ms ⚡\n"
    report += f"**置信度**: {result['confidence_score']:.1%} ✅"
    
    visual_data = {
        "current_line": quantum_pos.absolute_line,
        "scope_start": scope_info['scope_boundary']['start'],
        "scope_end": scope_info['scope_boundary']['end'],
        "repair_suggestions": result['repair_targets'][:3]
    }
    
    return report, result, visual_data

# 示例代码和界面保持不变...
# [此处保留原有的Gradio界面代码]

if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=True
    )