"""
数据导出器
支持多种数据格式导出
"""
import json
from typing import Optional, List, Dict, Any
from pathlib import Path
from datetime import datetime
from loguru import logger


class DataExporter:
    """
    数据导出器
    支持多种数据格式和文件格式
    """

    def __init__(self):
        self._custom_serializers = {}

    def register_serializer(self, name: str, serializer: callable):
        """注册自定义序列化器"""
        self._custom_serializers[name] = serializer

    # ==================== JSON 格式 ====================

    async def save_to_json(
        self,
        data: Any,
        filepath: str,
        indent: int = 2,
        ensure_ascii: bool = False,
        include_metadata: bool = False,
    ) -> bool:
        """保存为 JSON 文件"""
        try:
            Path(filepath).parent.mkdir(parents=True, exist_ok=True)

            if include_metadata:
                output = {
                    "_metadata": {
                        "exported_at": datetime.now().isoformat(),
                        "count": len(data) if isinstance(data, list) else 1,
                        "type": type(data).__name__,
                    },
                    "data": data,
                }
            else:
                output = data

            with open(filepath, 'w', encoding='utf-8') as f:
                json.dump(output, f, ensure_ascii=ensure_ascii, indent=indent, default=str)

            logger.info(f"数据已保存为 JSON: {filepath}")
            return True
        except Exception as e:
            logger.error(f"保存 JSON 失败: {e}")
            return False

    async def save_to_jsonlines(self, data: List[Dict], filepath: str) -> bool:
        """保存为 JSON Lines 格式（每行一个 JSON 对象）"""
        try:
            Path(filepath).parent.mkdir(parents=True, exist_ok=True)

            with open(filepath, 'w', encoding='utf-8') as f:
                for item in data:
                    f.write(json.dumps(item, ensure_ascii=False, default=str) + '\n')

            logger.info(f"数据已保存为 JSON Lines: {filepath}")
            return True
        except Exception as e:
            logger.error(f"保存 JSON Lines 失败: {e}")
            return False

    # ==================== CSV 格式 ====================

    async def save_to_csv(
        self,
        data: List[Dict],
        filepath: str,
        fieldnames: Optional[List[str]] = None,
        encoding: str = 'utf-8',
        delimiter: str = ',',
        include_header: bool = True,
    ) -> bool:
        """保存为 CSV 文件"""
        import csv

        if not data:
            logger.warning("没有数据可保存")
            return False

        try:
            Path(filepath).parent.mkdir(parents=True, exist_ok=True)

            # 自动推断字段名
            if not fieldnames:
                fieldnames = []
                for item in data:
                    for key in item.keys():
                        if key not in fieldnames:
                            fieldnames.append(key)

            with open(filepath, 'w', newline='', encoding=encoding) as f:
                writer = csv.DictWriter(f, fieldnames=fieldnames, delimiter=delimiter, extrasaction='ignore')
                if include_header:
                    writer.writeheader()
                writer.writerows(data)

            logger.info(f"数据已保存为 CSV: {filepath}")
            return True
        except Exception as e:
            logger.error(f"保存 CSV 失败: {e}")
            return False

    async def save_to_tsv(self, data: List[Dict], filepath: str, fieldnames: Optional[List[str]] = None) -> bool:
        """保存为 TSV 文件（制表符分隔）"""
        return await self.save_to_csv(data, filepath, fieldnames, delimiter='\t')

    # ==================== Excel 格式 ====================

    async def save_to_excel(
        self,
        data: List[Dict],
        filepath: str,
        sheet_name: str = "Sheet1",
        fieldnames: Optional[List[str]] = None,
        include_header: bool = True,
        auto_adjust_width: bool = True,
    ) -> bool:
        """保存为 Excel 文件 (.xlsx)"""
        try:
            from openpyxl import Workbook
            from openpyxl.utils import get_column_letter
        except ImportError:
            logger.error("需要安装 openpyxl: pip install openpyxl")
            return False

        if not data:
            logger.warning("没有数据可保存")
            return False

        try:
            Path(filepath).parent.mkdir(parents=True, exist_ok=True)

            wb = Workbook()
            ws = wb.active
            ws.title = sheet_name

            # 自动推断字段名
            if not fieldnames:
                fieldnames = []
                for item in data:
                    for key in item.keys():
                        if key not in fieldnames:
                            fieldnames.append(key)

            # 写入表头
            if include_header:
                for col, field in enumerate(fieldnames, 1):
                    ws.cell(row=1, column=col, value=field)

            # 写入数据
            start_row = 2 if include_header else 1
            for row_idx, item in enumerate(data, start_row):
                for col, field in enumerate(fieldnames, 1):
                    value = item.get(field, '')
                    if isinstance(value, (list, dict)):
                        value = json.dumps(value, ensure_ascii=False)
                    ws.cell(row=row_idx, column=col, value=value)

            # 自动调整列宽
            if auto_adjust_width:
                for col in range(1, len(fieldnames) + 1):
                    max_length = 0
                    column = get_column_letter(col)
                    for cell in ws[column]:
                        try:
                            if cell.value:
                                max_length = max(max_length, len(str(cell.value)))
                        except:
                            pass
                    ws.column_dimensions[column].width = min(max_length + 2, 50)

            wb.save(filepath)
            logger.info(f"数据已保存为 Excel: {filepath}")
            return True
        except Exception as e:
            logger.error(f"保存 Excel 失败: {e}")
            return False

    # ==================== XML 格式 ====================

    async def save_to_xml(
        self,
        data: List[Dict],
        filepath: str,
        root_name: str = "data",
        item_name: str = "item",
    ) -> bool:
        """保存为 XML 文件"""
        import xml.etree.ElementTree as ET

        try:
            Path(filepath).parent.mkdir(parents=True, exist_ok=True)

            root = ET.Element(root_name)

            for item in data:
                item_el = ET.SubElement(root, item_name)
                for key, value in item.items():
                    field_el = ET.SubElement(item_el, key)
                    if isinstance(value, (list, dict)):
                        field_el.text = json.dumps(value, ensure_ascii=False)
                    else:
                        field_el.text = str(value) if value is not None else ""

            tree = ET.ElementTree(root)
            ET.indent(tree, space="  ")
            tree.write(filepath, encoding='utf-8', xml_declaration=True)

            logger.info(f"数据已保存为 XML: {filepath}")
            return True
        except Exception as e:
            logger.error(f"保存 XML 失败: {e}")
            return False

    # ==================== YAML 格式 ====================

    async def save_to_yaml(self, data: Any, filepath: str) -> bool:
        """保存为 YAML 文件"""
        try:
            import yaml
        except ImportError:
            logger.error("需要安装 PyYAML: pip install pyyaml")
            return False

        try:
            Path(filepath).parent.mkdir(parents=True, exist_ok=True)

            with open(filepath, 'w', encoding='utf-8') as f:
                yaml.dump(data, f, allow_unicode=True, default_flow_style=False)

            logger.info(f"数据已保存为 YAML: {filepath}")
            return True
        except Exception as e:
            logger.error(f"保存 YAML 失败: {e}")
            return False

    # ==================== Markdown 格式 ====================

    async def save_to_markdown(
        self,
        data: List[Dict],
        filepath: str,
        title: str = "数据导出",
        fieldnames: Optional[List[str]] = None,
    ) -> bool:
        """保存为 Markdown 表格"""
        try:
            Path(filepath).parent.mkdir(parents=True, exist_ok=True)

            if not data:
                logger.warning("没有数据可保存")
                return False

            # 自动推断字段名
            if not fieldnames:
                fieldnames = []
                for item in data:
                    for key in item.keys():
                        if key not in fieldnames:
                            fieldnames.append(key)

            lines = [f"# {title}\n"]

            # 表头
            lines.append("| " + " | ".join(fieldnames) + " |")
            lines.append("| " + " | ".join(["---"] * len(fieldnames)) + " |")

            # 数据行
            for item in data:
                row = []
                for field in fieldnames:
                    value = item.get(field, '')
                    if isinstance(value, (list, dict)):
                        value = json.dumps(value, ensure_ascii=False)
                    value = str(value).replace('|', '\\|').replace('\n', '<br>')
                    row.append(value)
                lines.append("| " + " | ".join(row) + " |")

            with open(filepath, 'w', encoding='utf-8') as f:
                f.write('\n'.join(lines))

            logger.info(f"数据已保存为 Markdown: {filepath}")
            return True
        except Exception as e:
            logger.error(f"保存 Markdown 失败: {e}")
            return False

    # ==================== HTML 格式 ====================

    async def save_to_html(
        self,
        data: List[Dict],
        filepath: str,
        title: str = "数据导出",
        fieldnames: Optional[List[str]] = None,
        include_css: bool = True,
    ) -> bool:
        """保存为 HTML 表格"""
        try:
            Path(filepath).parent.mkdir(parents=True, exist_ok=True)

            if not data:
                logger.warning("没有数据可保存")
                return False

            # 自动推断字段名
            if not fieldnames:
                fieldnames = []
                for item in data:
                    for key in item.keys():
                        if key not in fieldnames:
                            fieldnames.append(key)

            css = """
            <style>
                body { font-family: Arial, sans-serif; margin: 20px; }
                h1 { color: #333; }
                table { border-collapse: collapse; width: 100%; }
                th, td { border: 1px solid #ddd; padding: 8px; text-align: left; }
                th { background-color: #4CAF50; color: white; }
                tr:nth-child(even) { background-color: #f2f2f2; }
            </style>
            """ if include_css else ""

            html_lines = [
                "<!DOCTYPE html>", "<html>", "<head>",
                f"<meta charset='utf-8'>", f"<title>{title}</title>", css,
                "</head>", "<body>", f"<h1>{title}</h1>", "<table>", "<thead><tr>",
            ]

            for field in fieldnames:
                html_lines.append(f"<th>{field}</th>")
            html_lines.append("</tr></thead><tbody>")

            for item in data:
                html_lines.append("<tr>")
                for field in fieldnames:
                    value = item.get(field, '')
                    if isinstance(value, (list, dict)):
                        value = json.dumps(value, ensure_ascii=False)
                    value = str(value).replace('<', '&lt;').replace('>', '&gt;')
                    html_lines.append(f"<td>{value}</td>")
                html_lines.append("</tr>")

            html_lines.extend(["</tbody></table>", "</body>", "</html>"])

            with open(filepath, 'w', encoding='utf-8') as f:
                f.write('\n'.join(html_lines))

            logger.info(f"数据已保存为 HTML: {filepath}")
            return True
        except Exception as e:
            logger.error(f"保存 HTML 失败: {e}")
            return False

    # ==================== SQL 格式 ====================

    async def save_to_sql(
        self,
        data: List[Dict],
        filepath: str,
        table_name: str = "extracted_data",
        create_table: bool = True,
        fieldnames: Optional[List[str]] = None,
    ) -> bool:
        """保存为 SQL INSERT 语句"""
        try:
            Path(filepath).parent.mkdir(parents=True, exist_ok=True)

            if not data:
                logger.warning("没有数据可保存")
                return False

            # 自动推断字段名
            if not fieldnames:
                fieldnames = []
                for item in data:
                    for key in item.keys():
                        if key not in fieldnames:
                            fieldnames.append(key)

            lines = []

            if create_table:
                lines.append(f"-- 自动生成的表结构")
                lines.append(f"CREATE TABLE IF NOT EXISTS {table_name} (")
                lines.append(f"    id INTEGER PRIMARY KEY AUTOINCREMENT,")
                for field in fieldnames:
                    lines.append(f"    {field} TEXT,")
                lines.append(");")
                lines.append("")

            lines.append(f"-- 数据插入")
            for item in data:
                columns = []
                values = []
                for field in fieldnames:
                    columns.append(field)
                    value = item.get(field, '')
                    if isinstance(value, (list, dict)):
                        value = json.dumps(value, ensure_ascii=False)
                    value = str(value).replace("'", "''") if value is not None else ""
                    values.append(f"'{value}'")
                lines.append(f"INSERT INTO {table_name} ({', '.join(columns)}) VALUES ({', '.join(values)});")

            with open(filepath, 'w', encoding='utf-8') as f:
                f.write('\n'.join(lines))

            logger.info(f"数据已保存为 SQL: {filepath}")
            return True
        except Exception as e:
            logger.error(f"保存 SQL 失败: {e}")
            return False

    # ==================== SQLite 数据库 ====================

    async def save_to_sqlite(
        self,
        data: List[Dict],
        filepath: str,
        table_name: str = "extracted_data",
        fieldnames: Optional[List[str]] = None,
    ) -> bool:
        """保存到 SQLite 数据库"""
        import sqlite3

        try:
            Path(filepath).parent.mkdir(parents=True, exist_ok=True)

            if not data:
                logger.warning("没有数据可保存")
                return False

            # 自动推断字段名
            if not fieldnames:
                fieldnames = []
                for item in data:
                    for key in item.keys():
                        if key not in fieldnames:
                            fieldnames.append(key)

            conn = sqlite3.connect(filepath)
            cursor = conn.cursor()

            columns = ["id INTEGER PRIMARY KEY AUTOINCREMENT"]
            for field in fieldnames:
                columns.append(f"{field} TEXT")

            cursor.execute(f"CREATE TABLE IF NOT EXISTS {table_name} ({', '.join(columns)})")

            placeholders = ", ".join(["?" for _ in fieldnames])
            for item in data:
                values = []
                for field in fieldnames:
                    value = item.get(field, '')
                    if isinstance(value, (list, dict)):
                        value = json.dumps(value, ensure_ascii=False)
                    values.append(str(value) if value is not None else None)
                cursor.execute(
                    f"INSERT INTO {table_name} ({', '.join(fieldnames)}) VALUES ({placeholders})",
                    values
                )

            conn.commit()
            conn.close()

            logger.info(f"数据已保存到 SQLite: {filepath}")
            return True
        except Exception as e:
            logger.error(f"保存 SQLite 失败: {e}")
            return False

    # ==================== 通用保存方法 ====================

    async def save(
        self,
        data: Any,
        filepath: str,
        format: Optional[str] = None,
        **kwargs,
    ) -> bool:
        """通用保存方法，自动根据文件扩展名选择格式"""
        # 确定格式
        if format:
            ext = format.lower()
        else:
            ext = Path(filepath).suffix.lower().lstrip('.')

        # 格式映射
        format_map = {
            'json': self.save_to_json,
            'jsonl': self.save_to_jsonlines,
            'csv': self.save_to_csv,
            'tsv': self.save_to_tsv,
            'xlsx': self.save_to_excel,
            'excel': self.save_to_excel,
            'xml': self.save_to_xml,
            'yaml': self.save_to_yaml,
            'yml': self.save_to_yaml,
            'md': self.save_to_markdown,
            'markdown': self.save_to_markdown,
            'html': self.save_to_html,
            'htm': self.save_to_html,
            'sql': self.save_to_sql,
            'sqlite': self.save_to_sqlite,
            'db': self.save_to_sqlite,
        }

        if ext not in format_map:
            logger.error(f"不支持的格式: {ext}")
            logger.info(f"支持的格式: {', '.join(format_map.keys())}")
            return False

        # 确保数据是列表（除了 JSON/YAML 可以是任意类型）
        if ext not in ['json', 'yaml', 'yml', 'jsonl'] and not isinstance(data, list):
            data = [data] if isinstance(data, dict) else list(data)

        return await format_map[ext](data, filepath, **kwargs)

    def get_supported_formats(self) -> List[str]:
        """获取支持的格式列表"""
        return [
            'json', 'jsonl', 'csv', 'tsv',
            'xlsx', 'xml', 'yaml', 'yml',
            'md', 'html', 'sql', 'sqlite', 'db',
        ]


# 快捷函数
async def save_to_json(data: Any, filepath: str):
    """保存数据到JSON文件"""
    exporter = DataExporter()
    return await exporter.save_to_json(data, filepath)


async def save_to_csv(data: List[Dict], filepath: str):
    """保存数据到CSV文件"""
    exporter = DataExporter()
    return await exporter.save_to_csv(data, filepath)
