#!/usr/bin/env python3
# 修复知识库功能：预览、下载、生成知识库、文件分类存储
import re

PORTAL = r"D:\AI-Ad-Agent\portal\proxy_server.py"

with open(PORTAL, 'r', encoding='utf-8') as f:
    content = f.read()

# ===== 1. 在 do_GET() 中添加预览和下载路由 =====
# 在知识库列表路由后添加
old_get_route = "            # Knowledge Management API (GET)\n            if self.path == '/api/knowledge/list':\n                return self._handle_knowledge_list()"
new_get_route = """            # Knowledge Management API (GET)
            if self.path == '/api/knowledge/list':
                return self._handle_knowledge_list()
            if self.path.startswith('/api/knowledge/preview'):
                return self._handle_knowledge_preview()
            if self.path.startswith('/api/knowledge/download'):
                return self._handle_knowledge_download()"""
if old_get_route in content:
    content = content.replace(old_get_route, new_get_route)
    print("[✓] 已添加预览和下载路由到 do_GET()")
else:
    print("[✗] 未找到 do_GET() 中的知识库列表路由")

# ===== 2. 修改预览API，支持通过文件ID预览 =====
old_preview = '''    def _handle_knowledge_preview(self):
        """GET /api/knowledge/preview?path=xxx — 预览上传的文件"""
        try:
            from urllib.parse import urlparse, parse_qs
            parsed = urlparse(self.path)
            params = parse_qs(parsed.query)
            filepath = params.get('path', [None])[0]
            if not filepath:
                self._json_response(400, {"error": "Missing path"})
                return
            # 安全校验：只允许预览 knowledge_uploads 目录下的文件
            allowed_dir = "D:/AI-Ad-Agent/portal/data/knowledge_uploads/"
            if not os.path.normpath(filepath).startswith(os.path.normpath(allowed_dir)):
                self._json_response(403, {"error": "Access denied"})
                return
            if not os.path.exists(filepath):
                self._json_response(404, {"error": "File not found"})
                return
            # 根据文件类型设置 Content-Type
            ext = os.path.splitext(filepath)[1].lower()
            mime_types = {
                '.jpg': 'image/jpeg', '.jpeg': 'image/jpeg', '.png': 'image/png',
                '.gif': 'image/gif', '.bmp': 'image/bmp',
                '.mp4': 'video/mp4', '.avi': 'video/x-msvideo', '.mov': 'video/quicktime',
                '.pdf': 'application/pdf',
                '.doc': 'application/msword', '.docx': 'application/vnd.openxmlformats-officedocument.wordprocessingml.document',
                '.xls': 'application/vnd.ms-excel', '.xlsx': 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet',
                '.ppt': 'application/vnd.ms-powerpoint', '.pptx': 'application/vnd.openxmlformats-officedocument.presentationml.presentation',
                '.txt': 'text/plain', '.md': 'text/markdown',
            }
            content_type = mime_types.get(ext, 'application/octet-stream')
            with open(filepath, 'rb') as f:
                file_data = f.read()
            self.send_response(200)
            self.send_header('Content-Type', content_type)
            self.send_header('Content-Length', len(file_data))
            self.send_header('Cache-Control', 'public, max-age=3600')
            self.end_headers()
            self.wfile.write(file_data)
        except Exception as e:
            self._json_response(500, {"error": str(e})'''

new_preview = '''    def _handle_knowledge_preview(self):
        """GET /api/knowledge/preview?id=xxx — 预览上传的文件（通过ID）"""
        try:
            from urllib.parse import urlparse, parse_qs
            parsed = urlparse(self.path)
            params = parse_qs(parsed.query)
            file_id = params.get('id', [None])[0]
            if not file_id:
                self._json_response(400, {"error": "Missing file id"})
                return
            
            # 从数据库读取文件路径
            db_path = "D:/AI-Ad-Agent/portal/data/knowledge.db"
            conn = sqlite3.connect(db_path)
            c = conn.cursor()
            c.execute('SELECT filepath, filename FROM files WHERE id = ?', (file_id,))
            row = c.fetchone()
            conn.close()
            
            if not row:
                self._json_response(404, {"error": "File not found in database"})
                return
            
            filepath = row[0]
            original_filename = row[1]
            
            # 安全校验：只允许预览 knowledge_uploads 目录下的文件
            allowed_dir = os.path.normpath("D:/AI-Ad-Agent/portal/data/knowledge_uploads")
            if not os.path.normpath(filepath).startswith(allowed_dir):
                self._json_response(403, {"error": "Access denied"})
                return
            
            if not os.path.exists(filepath):
                self._json_response(404, {"error": f"File not found on disk: {original_filename}"})
                return
            
            # 根据文件类型设置 Content-Type
            ext = os.path.splitext(filepath)[1].lower()
            mime_types = {
                '.jpg': 'image/jpeg', '.jpeg': 'image/jpeg', '.png': 'image/png',
                '.gif': 'image/gif', '.bmp': 'image/bmp',
                '.mp4': 'video/mp4', '.avi': 'video/x-msvideo', '.mov': 'video/quicktime',
                '.pdf': 'application/pdf',
                '.doc': 'application/msword', '.docx': 'application/vnd.openxmlformats-officedocument.wordprocessingml.document',
                '.xls': 'application/vnd.ms-excel', '.xlsx': 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet',
                '.ppt': 'application/vnd.ms-powerpoint', '.pptx': 'application/vnd.openxmlformats-officedocument.presentationml.presentation',
                '.txt': 'text/plain', '.md': 'text/markdown',
            }
            content_type = mime_types.get(ext, 'application/octet-stream')
            
            with open(filepath, 'rb') as f:
                file_data = f.read()
            
            self.send_response(200)
            self.send_header('Content-Type', content_type)
            self.send_header('Content-Length', len(file_data))
            self.send_header('Cache-Control', 'public, max-age=3600')
            self.end_headers()
            self.wfile.write(file_data)
        except Exception as e:
            self._json_response(500, {"error": str(e)})'''

if old_preview in content:
    content = content.replace(old_preview, new_preview)
    print("[✓] 已修改预览API，支持通过文件ID预览")
else:
    print("[!] 未找到旧的预览API，尝试搜索...")
    if '_handle_knowledge_preview' in content:
        print("[✓] 预览API已存在，可能需要手动修改")

# ===== 3. 添加下载API =====
download_api = '''
    def _handle_knowledge_download(self):
        """GET /api/knowledge/download?id=xxx — 下载上传的文件"""
        try:
            from urllib.parse import urlparse, parse_qs
            parsed = urlparse(self.path)
            params = parse_qs(parsed.query)
            file_id = params.get('id', [None])[0]
            if not file_id:
                self._json_response(400, {"error": "Missing file id"})
                return
            
            # 从数据库读取文件路径
            db_path = "D:/AI-Ad-Agent/portal/data/knowledge.db"
            conn = sqlite3.connect(db_path)
            c = conn.cursor()
            c.execute('SELECT filepath, filename FROM files WHERE id = ?', (file_id,))
            row = c.fetchone()
            conn.close()
            
            if not row:
                self._json_response(404, {"error": "File not found in database"})
                return
            
            filepath = row[0]
            original_filename = row[1]
            
            if not os.path.exists(filepath):
                self._json_response(404, {"error": f"File not found on disk: {original_filename}"})
                return
            
            # 读取文件并返回（设置 Content-Disposition 为 attachment）
            with open(filepath, 'rb') as f:
                file_data = f.read()
            
            self.send_response(200)
            self.send_header('Content-Type', 'application/octet-stream')
            self.send_header('Content-Disposition', f'attachment; filename="{original_filename}"')
            self.send_header('Content-Length', len(file_data))
            self.end_headers()
            self.wfile.write(file_data)
        except Exception as e:
            self._json_response(500, {"error": str(e)})
'''

if '_handle_knowledge_download' not in content:
    # 在预览API后插入下载API
    insert_pos = content.find('def _handle_knowledge_preview(self):')
    if insert_pos > 0:
        # 找到预览API的结束位置
        preview_end = content.find('\n    \n    def ', insert_pos)
        if preview_end > 0:
            content = content[:preview_end] + download_api + '\n' + content[preview_end:]
            print("[✓] 已添加下载API")
        else:
            print("[✗] 找不到预览API结束位置")
    else:
        print("[✗] 找不到预览API")
else:
    print("[!] 下载API已存在")

# ===== 4. 修改上传API，按文件类型分类存储 =====
old_upload = '''    def _handle_knowledge_upload(self):
        """POST /api/knowledge/upload"""
        try:
            fields, files = self._parse_multipart_form()
            if not files or 'file' not in files:
                self._json_response(400, {"error": "No file"})
                return
            file_info = files['file']
            filename = file_info[0]  # 文件名
            file_data = file_info[1]  # 文件内容 (bytes)
            
            import time
            timestamp = int(time.time())
            safe_filename = f"{timestamp}_{filename}"
            filepath = f"D:/AI-Ad-Agent/portal/data/knowledge_uploads/{safe_filename}"
            
            import os
            os.makedirs(os.path.dirname(filepath), exist_ok=True)
            with open(filepath, 'wb') as f:
                f.write(file_data)'''

new_upload = '''    def _handle_knowledge_upload(self):
        """POST /api/knowledge/upload — 按文件类型分类存储"""
        try:
            fields, files = self._parse_multipart_form()
            if not files or 'file' not in files:
                self._json_response(400, {"error": "No file"})
                return
            file_info = files['file']
            filename = file_info[0]  # 文件名
            file_data = file_info[1]  # 文件内容 (bytes)
            
            import time, os
            timestamp = int(time.time())
            safe_filename = f"{timestamp}_{filename}"
            
            # 根据文件扩展名分类存储
            ext = os.path.splitext(filename)[1].lower()
            type_dirs = {
                '.jpg': 'images', '.jpeg': 'images', '.png': 'images', '.gif': 'images', '.bmp': 'images', '.webp': 'images',
                '.mp4': 'videos', '.avi': 'videos', '.mov': 'videos', '.wmv': 'videos', '.flv': 'videos',
                '.doc': 'docs', '.docx': 'docs', '.pdf': 'docs', '.txt': 'docs', '.md': 'docs',
                '.xls': 'sheets', '.xlsx': 'sheets', '.csv': 'sheets',
                '.ppt': 'slides', '.pptx': 'slides',
            }
            subdir = type_dirs.get(ext, 'misc')
            base_dir = f"D:/AI-Ad-Agent/portal/data/knowledge_uploads/{subdir}"
            os.makedirs(base_dir, exist_ok=True)
            filepath = f"{base_dir}/{safe_filename}"
            
            with open(filepath, 'wb') as f:
                f.write(file_data)'''

if old_upload in content:
    content = content.replace(old_upload, new_upload)
    print("[✓] 已修改上传API，按文件类型分类存储")
else:
    print("[!] 未找到旧的上传API，可能已修改")

# ===== 5. 修复生成知识库API =====
old_generate = '''    def _handle_knowledge_generate(self):
        """POST /api/knowledge/generate — 基于上传文件生成知识"""
        try:
            fields, files = self._parse_multipart_form()
            prompt = fields.get('prompt', [''])[0]
            category = fields.get('category', ['misc'])[0]
            # TODO: 调用 Ollama 生成知识
            self._json_response(200, {"success": True, "result": "知识生成功能开发中"})
        except Exception as e:
            self._json_response(500, {"error": str(e)})'''

new_generate = '''    def _handle_knowledge_generate(self):
        """POST /api/knowledge/generate — 基于上传文件生成知识"""
        try:
            from urllib.parse import urlparse, parse_qs
            parsed = urlparse(self.path)
            params = parse_qs(parsed.query)
            file_id = params.get('id', [None])[0]
            
            if not file_id:
                self._json_response(400, {"error": "Missing file id"})
                return
            
            # 从数据库读取文件
            db_path = "D:/AI-Ad-Agent/portal/data/knowledge.db"
            conn = sqlite3.connect(db_path)
            c = conn.cursor()
            c.execute('SELECT filepath, filename, filetype FROM files WHERE id = ?', (file_id,))
            row = c.fetchone()
            conn.close()
            
            if not row:
                self._json_response(404, {"error": "File not found"})
                return
            
            filepath = row[0]
            filename = row[1]
            filetype = row[2]
            
            if not os.path.exists(filepath):
                self._json_response(404, {"error": f"File not found on disk: {filename}"})
                return
            
            # 读取文件内容（根据类型）
            content_text = ""
            if filetype in ['txt', 'md', 'csv']:
                with open(filepath, 'r', encoding='utf-8', errors='ignore') as f:
                    content_text = f.read()[:5000]  # 前5000字符
            elif filetype in ['doc', 'docx', 'pdf']:
                content_text = f"[二进制文件: {filename}] 需要安装 python-docx/PyPDF2 来提取内容"
            else:
                content_text = f"[文件类型: {filetype}] {filename}"
            
            # 调用 Ollama 生成知识摘要
            import requests
            ollama_url = "http://127.0.0.1:11434/api/generate"
            prompt = f"请对以下内容进行知识提取和摘要（200字以内）：\\n\\n{content_text[:2000]}"
            
            try:
                resp = requests.post(ollama_url, json={
                    "model": "qwen3:8b",
                    "prompt": prompt,
                    "stream": False,
                    "think": False,
                }, timeout=60)
                if resp.status_code == 200:
                    result = resp.json().get('response', '生成失败')
                else:
                    result = content_text[:500]  # fallback: 使用前500字符
            except Exception as e:
                print(f"[Knowledge Generate] Ollama error: {e}")
                result = content_text[:500]
            
            # 更新数据库 metadata 字段
            conn = sqlite3.connect(db_path)
            c = conn.cursor()
            c.execute('UPDATE files SET metadata = ? WHERE id = ?', (result, file_id))
            conn.commit()
            conn.close()
            
            self._json_response(200, {"success": True, "result": result})
        except Exception as e:
            self._json_response(500, {"error": str(e)})'''

if old_generate in content:
    content = content.replace(old_generate, new_generate)
    print("[✓] 已修复生成知识库API")
else:
    print("[!] 未找到旧的生产API，可能已修改")

# ===== 保存修改 =====
with open(PORTAL, 'w', encoding='utf-8') as f:
    f.write(content)

print("\n✅ 后端修复完成！需要重启代理服务器。")