"""
Dify文件上传工具 - 优化版

主要功能：
- 从URL下载文件并自动上传到Dify API
- 支持常见文件类型：JPG、PNG、GIF、PDF、DOCX、TXT等
- 自动处理文件大小检查、MIME类型识别、临时文件清理

使用方法：
    from upload_file import upload_file_from_url

    result = upload_file_from_url(
        file_url="http://example.com/image.jpg",
        base_url="http://192.168.2.236:3001/v1",
        api_key="app-QdfDKqHAI3dlB6tvnibuh6rv"
    )

    file_id = result['id']  # 获取上传后的文件ID
"""

import os
import tempfile
import requests
import logging
from urllib.parse import urlparse, unquote
from typing import Optional

# 获取模块级别的logger，避免影响全局日志配置
logger = logging.getLogger(__name__)

# 常用MIME类型映射
MIME_TYPE_MAP = {
    '.jpg': 'image/jpeg',
    '.jpeg': 'image/jpeg',
    '.png': 'image/png',
    '.gif': 'image/gif',
    '.webp': 'image/webp',
    '.svg': 'image/svg+xml',
    '.pdf': 'application/pdf',
    '.txt': 'text/plain',
    '.docx': 'application/vnd.openxmlformats-officedocument.wordprocessingml.document'
}


def download_file_from_url(
    url: str,
    download_dir: Optional[str] = None,
    filename: Optional[str] = None,
    timeout: int = 30,
    max_retries: int = 3
) -> str:
    """
    从URL下载文件到本地并返回文件路径

    Args:
        url: 文件的URL地址
        download_dir: 下载目录，如果为None则使用系统临时目录
        filename: 指定文件名，如果为None则从URL中提取
        timeout: 请求超时时间（秒）
        max_retries: 最大重试次数

    Returns:
        str: 下载后的本地文件路径

    Raises:
        Exception: 下载失败时抛出异常
    """

    # 设置下载目录
    if download_dir is None:
        download_dir = tempfile.gettempdir()

    # 确保下载目录存在
    os.makedirs(download_dir, exist_ok=True)

    # 提取文件名
    if filename is None:
        filename = _extract_filename_from_url(url)

    # 构建完整的文件路径
    file_path = os.path.join(download_dir, filename)

    # 下载文件（带重试机制）
    for attempt in range(max_retries):
        try:
            logger.info(f"正在下载文件: {url} (尝试 {attempt + 1}/{max_retries})")

            # 发送GET请求下载文件
            response = requests.get(url, timeout=timeout, stream=True)
            response.raise_for_status()

            # 写入文件
            with open(file_path, 'wb') as f:
                for chunk in response.iter_content(chunk_size=8192):
                    if chunk:
                        f.write(chunk)

            # 验证文件是否下载成功
            if os.path.exists(file_path) and os.path.getsize(file_path) > 0:
                logger.info(f"文件下载成功: {file_path} (大小: {os.path.getsize(file_path)} 字节)")
                return file_path
            else:
                raise Exception("下载的文件为空或不存在")

        except requests.exceptions.Timeout:
            error_msg = f"请求超时 (超过 {timeout} 秒)"
            logger.warning(f"{error_msg}, 尝试 {attempt + 1}/{max_retries}")
            if attempt == max_retries - 1:
                raise Exception(f"下载失败：{error_msg}")

        except requests.exceptions.RequestException as e:
            error_msg = f"请求异常: {str(e)}"
            logger.warning(f"{error_msg}, 尝试 {attempt + 1}/{max_retries}")
            if attempt == max_retries - 1:
                raise Exception(f"下载失败：{error_msg}")

        except Exception as e:
            error_msg = f"下载过程中发生错误: {str(e)}"
            logger.warning(f"{error_msg}, 尝试 {attempt + 1}/{max_retries}")
            if attempt == max_retries - 1:
                raise Exception(f"下载失败：{error_msg}")

    raise Exception("下载失败：已达到最大重试次数")


def _extract_filename_from_url(url: str) -> str:
    """
    从URL中提取文件名

    Args:
        url: 文件URL

    Returns:
        str: 提取的文件名
    """
    try:
        # 解析URL
        parsed_url = urlparse(url)
        path = unquote(parsed_url.path)

        # 从路径中提取文件名
        filename = os.path.basename(path)

        # 如果没有找到文件名或文件名为空，使用默认名称
        if not filename or filename == '/':
            filename = "downloaded_file"

        # 移除查询参数（如果文件名中包含）
        if '?' in filename:
            filename = filename.split('?')[0]

        return filename

    except Exception as e:
        logger.warning(f"无法从URL提取文件名: {str(e)}, 使用默认文件名")
        return "downloaded_file"


def upload_file_to_dify(
    file_path: str,
    base_url: str,
    api_key: str,
    user: str = "default_user",
    verify_ssl: bool = False
) -> dict:
    """
    上传文件到Dify API

    Args:
        file_path: 本地文件路径
        base_url: Dify API基础URL (例如: http://192.168.2.236:3001/v1)
        api_key: API密钥
        user: 用户标识
        verify_ssl: 是否验证SSL证书

    Returns:
        dict: 上传响应结果

    Raises:
        Exception: 上传失败时抛出异常
    """
    # 检查文件是否存在
    if not os.path.exists(file_path):
        raise FileNotFoundError(f"文件不存在: {file_path}")

    # 检查文件大小
    file_size = os.path.getsize(file_path)
    file_size_mb = file_size / (1024 * 1024)
    logger.info(f"准备上传文件: {file_path} (大小: {file_size_mb:.2f} MB)")

    # 检查文件大小是否超过限制
    if file_size_mb > 10:
        logger.warning(f"文件大小 {file_size_mb:.2f} MB 可能超过服务器限制 (10 MB)")

    upload_url = f"{base_url}/files/upload"
    headers = {"Authorization": f"Bearer {api_key}"}

    try:
        with open(file_path, "rb") as f:
            # 获取文件扩展名和MIME类型
            file_ext = os.path.splitext(file_path)[1].lower()
            mime_type = MIME_TYPE_MAP.get(file_ext, "application/octet-stream")

            # 构建文件上传数据
            files = {"file": (os.path.basename(file_path), f, mime_type)}
            data = {"user": user}

            # 发送上传请求
            response = requests.post(
                upload_url,
                headers=headers,
                files=files,
                data=data,
                verify=verify_ssl
            )

            # 检查响应
            if response.status_code == 201:
                result = response.json()
                logger.info(f"文件上传成功: {result.get('name', 'unknown')} (ID: {result.get('id', 'unknown')})")
                return result
            else:
                error_msg = f"上传失败 (状态码: {response.status_code}): {response.text}"
                logger.error(error_msg)
                raise requests.exceptions.HTTPError(error_msg)

    except requests.exceptions.RequestException as e:
        logger.error(f"上传文件请求失败: {str(e)}")
        raise
    except Exception as e:
        logger.error(f"上传文件失败: {str(e)}")
        raise



def check_app_config(base_url: str, api_key: str):
    """
    检查Dify应用的文件上传配置

    Args:
        base_url: Dify API基础URL
        api_key: API密钥
    """
    try:
        # 获取应用参数配置
        config_url = f"{base_url}/parameters"
        headers = {"Authorization": f"Bearer {api_key}"}

        response = requests.get(config_url, headers=headers, verify=False)
        response.raise_for_status()

        config = response.json()
        logger.info(f"应用配置获取成功")

        # 检查文件上传配置
        file_upload = config.get("file_upload", {})
        if file_upload.get("enabled", False):
            logger.info("✓ 文件上传功能已启用")
            logger.info(f"  - 允许的文件类型: {file_upload.get('allowed_file_types', [])}")
            logger.info(f"  - 允许的文件扩展名: {file_upload.get('allowed_file_extensions', [])}")
            logger.info(f"  - 文件大小限制: {file_upload.get('fileUploadConfig', {}).get('image_file_size_limit', 'N/A')} MB")
        else:
            logger.warning("✗ 文件上传功能未启用")

        return config

    except Exception as e:
        logger.error(f"检查应用配置失败: {str(e)}")
        return None


def test_download_and_upload():
    """
    测试下载和上传功能的示例
    """
    # 测试用的文件URL
    file_url = "http://192.168.2.236:9000/lzwcai/upload/2025-07-29/34b28da03f3c43b0921ba1b76857bbc0/34b28da03f3c43b0921ba1b76857bbc0.JPG?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=minio%2F20250729%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20250729T075242Z&X-Amz-Expires=604800&X-Amz-SignedHeaders=host&X-Amz-Signature=80f58d37c3bd52fb2b25efa36b5df0d73c8da7c773e7a7066dd6563710c619d6"

    # 上传API配置
    base_url = "http://192.168.2.236:3001/v1"
    upload_url = f"{base_url}/files/upload"
    api_key = "app-QdfDKqHAI3dlB6tvnibuh6rv"
    headers = {"Authorization": f"Bearer {api_key}"}
    user = "abc-123"

    try:
        logger.info("开始测试下载和上传流程...")

        # 0. 检查应用配置
        logger.info("检查Dify应用配置...")
        config = check_app_config(base_url, api_key)
        if not config:
            logger.error("无法获取应用配置，终止测试")
            return

        # 1. 下载文件
        logger.info(f"正在下载文件: {file_url}")
        file_path = download_file_from_url(file_url)
        logger.info(f"文件下载成功，保存路径: {file_path}")

        # 2. 上传文件
        logger.info(f"正在上传文件到: {upload_url}")

        # 检查文件大小
        file_size = os.path.getsize(file_path)
        file_size_mb = file_size / (1024 * 1024)
        logger.info(f"文件大小: {file_size_mb:.2f} MB")

        # 检查文件大小是否超过限制（通常为10MB对于图片）
        if file_size_mb > 10:
            logger.warning(f"文件大小 {file_size_mb:.2f} MB 可能超过服务器限制")

        with open(file_path, "rb") as f:
            # 获取文件扩展名来确定MIME类型
            file_ext = os.path.splitext(file_path)[1].lower()
            mime_type = "image/jpeg" if file_ext in ['.jpg', '.jpeg'] else "application/octet-stream"

            # 构建文件上传数据，指定正确的MIME类型
            files = {"file": (os.path.basename(file_path), f, mime_type)}
            data = {"user": user}

            # 不要在headers中设置Content-Type，让requests自动处理multipart/form-data
            upload_headers = headers.copy()
            if "Content-Type" in upload_headers:
                del upload_headers["Content-Type"]

            # 添加SSL验证跳过选项，用于自签名证书
            response = requests.post(upload_url, headers=upload_headers, files=files, data=data, verify=False)

            # 打印详细的响应信息用于调试
            logger.info(f"响应状态码: {response.status_code}")
            logger.info(f"响应头: {response.headers}")
            logger.info(f"响应内容: {response.text}")

            response.raise_for_status()

        result = response.json()
        logger.info(f"文件上传成功，响应: {result}")

        # 3. 清理临时文件
        try:
            os.remove(file_path)
            logger.info(f"已清理临时文件: {file_path}")
        except Exception as e:
            logger.warning(f"清理临时文件失败: {str(e)}")

        return result

    except Exception as e:
        logger.error(f"测试失败: {str(e)}")
        raise


def upload_file_from_url(
    file_url: str,
    base_url: str,
    api_key: str,
    user: str = "default_user",
    verify_ssl: bool = False
) -> dict:
    """
    从URL下载文件并上传到Dify API - 一站式解决方案

    这是一个优化后的方法，只需要提供URL地址、base_url和api_key，就能自动完成下载和上传。
    支持常见的文件类型：JPG、PNG、GIF、PDF、DOCX、TXT等。
    自动处理文件大小检查、MIME类型识别、临时文件清理等。

    Args:
        file_url: 要下载的文件URL
        base_url: Dify API基础URL (例如: http://192.168.2.236:3001/v1)
        api_key: API密钥 (例如: app-QdfDKqHAI3dlB6tvnibuh6rv)
        user: 用户标识 (可选，默认为 default_user)
        verify_ssl: 是否验证SSL证书 (可选，默认为 False)

    Returns:
        dict: 上传成功后的结果，包含文件ID、名称、大小等信息
        示例: {
            'id': 'a239b623-40a8-482c-859f-bb8368d5b1fe',
            'name': 'example.jpg',
            'size': 495240,
            'extension': 'jpg',
            'mime_type': 'image/jpeg',
            'created_by': '92c4b250-e0e7-4123-900d-f5c2187679a2',
            'created_at': 1753777420
        }

    使用示例:
        result = upload_file_from_url(
            file_url="http://example.com/image.jpg",
            base_url="http://192.168.2.236:3001/v1",
            api_key="app-QdfDKqHAI3dlB6tvnibuh6rv"
        )
        file_id = result['id']  # 获取上传后的文件ID

    Raises:
        Exception: 下载或上传失败时抛出异常
    """
    temp_file_path = None
    try:
        logger.info(f"开始处理文件: {file_url}")

        # 1. 下载文件到临时目录
        temp_file_path = download_file_from_url(file_url)

        # 2. 上传文件到Dify (复用已有的上传函数)
        result = upload_file_to_dify(temp_file_path, base_url, api_key, user, verify_ssl)
        return result

    except Exception as e:
        logger.error(f"处理文件失败: {str(e)}")
        raise
    finally:
        # 清理临时文件
        if temp_file_path and os.path.exists(temp_file_path):
            try:
                os.remove(temp_file_path)
                logger.info(f"已清理临时文件: {temp_file_path}")
            except Exception as e:
                logger.warning(f"清理临时文件失败: {str(e)}")


def test_upload_functionality():
    """
    测试文件上传功能的示例
    注意：这个函数包含示例配置，实际使用时请替换为真实的配置
    """
    # 示例配置 - 实际使用时请替换为真实的配置
    file_url = "https://example.com/test-image.jpg"  # 替换为实际的文件URL
    base_url = "http://localhost:3001/v1"  # 替换为实际的Dify API地址
    api_key = "your-api-key-here"  # 替换为实际的API密钥

    try:
        logger.info("=== 开始测试文件上传功能 ===")

        # 检查应用配置
        logger.info("检查Dify应用配置...")
        config = check_app_config(base_url, api_key)
        if not config:
            logger.warning("无法获取应用配置，但继续测试上传功能")

        # 调用上传方法
        result = upload_file_from_url(file_url, base_url, api_key)

        logger.info("=== 上传成功！结果如下 ===")
        logger.info(f"文件ID: {result.get('id')}")
        logger.info(f"文件名: {result.get('name')}")
        logger.info(f"文件大小: {result.get('size')} 字节")
        logger.info(f"文件类型: {result.get('mime_type')}")
        logger.info(f"扩展名: {result.get('extension')}")

        return result

    except Exception as e:
        logger.error(f"测试失败: {str(e)}")
        raise


if __name__ == "__main__":
    # 运行测试 - 注意：需要先配置正确的URL和API密钥
    print("警告：测试函数包含示例配置，请先修改为实际配置后再运行")
    # test_upload_functionality()  # 取消注释并配置后运行