import asyncio
import gzip
import json
import sqlite3
import time
import os
from io import BytesIO
from typing import Callable, Any, List

from hdfs import InsecureClient


class HDFSDataProcessor:
    def __init__(
            self,
            hdfs_url="http://VIP-DC-MASTER-2:9870",
            hdfs_user="root",
            db_file="processed_files.db",
            batch_size=50,
            retry_limit=3,
    ):
        self.hdfs_url = hdfs_url
        self.hdfs_user = hdfs_user
        self.db_file = db_file
        self.batch_size = batch_size
        self.retry_limit = retry_limit
        self.client = InsecureClient(self.hdfs_url, user=self.hdfs_user)
        self.read_hdfs_fanc = {"all": self.all_read_gz, "batch": self.batch_read_gz}
        self.read_hdfs_model = "all"
        self.init_db()

    def init_db(self):
        """初始化 SQLite 数据库"""
        with sqlite3.connect(self.db_file) as conn:
            cursor = conn.cursor()
            cursor.execute("""
                    CREATE TABLE IF NOT EXISTS processed_files (
                        file_path TEXT PRIMARY KEY
                    )
                    """)
            conn.commit()

    def save_processed_file(self, file_path):
        """保存处理过的文件"""
        with sqlite3.connect(self.db_file) as conn:
            cursor = conn.cursor()
            cursor.execute("INSERT OR IGNORE INTO processed_files (file_path) VALUES (?)", (file_path,))
            conn.commit()

    def is_file_processed(self, file_path):
        """检查文件是否已处理"""
        with sqlite3.connect(self.db_file) as conn:
            cursor = conn.cursor()
            cursor.execute("SELECT file_path FROM processed_files WHERE file_path = ?", (file_path,))
            result = cursor.fetchone()
        return result is not None

    def list_gz_files(self, hdfs_dir):
        """列出 HDFS 目录中的所有 gzip 文件"""
        return [f"{hdfs_dir}/{file[0]}" for file in self.client.list(hdfs_dir, status=True) if file[0].endswith(".gz")]

    def count_total_lines(self, gz_file_path: str):
        with self.client.read(gz_file_path) as hdfs_file:
            with gzip.GzipFile(fileobj=hdfs_file) as gz:
                return sum(1 for _ in gz)

    def batch_read_gz(self, gz_file_path: str):
        """分批读取 gz 文件"""
        with self.client.read(gz_file_path) as hdfs_file:
            with gzip.GzipFile(fileobj=hdfs_file) as gz:
                while True:
                    lines = []
                    for _ in range(self.batch_size):
                        try:
                            line = next(gz)
                            if line.strip():  # 移除空行
                                lines.append(line.decode("utf-8"))  # 解码
                        except StopIteration:  # 文件已读完
                            break
                    if not lines:
                        break
                    yield lines

    def all_read_gz(self, gz_file_path: str, encoding='utf-8'):
        """
        读取 HDFS 上的 .gz 文件内容。
        :param hdfs_path: HDFS 文件路径（必须以 .gz 结尾）
        :param encoding: 文件编码格式（默认 utf-8）
        :return: 文件内容
        """
        with self.client.read(gz_file_path) as reader:  # 以二进制模式读取
            compressed_data = reader.read()  # 读取压缩数据
            with gzip.GzipFile(fileobj=BytesIO(compressed_data)) as gz_file:  # 解压缩
                content = gz_file.read().decode(encoding)  # 解码为字符串
                print(f"文件读取成功: {gz_file_path}")
                lines = [i for i in content.splitlines() if i.strip()]
                result = [lines[i:i + self.batch_size] for i in range(0, len(lines), self.batch_size)]
                return result

    async def process_data(self, data, process_func):
        """处理数据并执行处理函数"""
        retry_count = 0
        while retry_count < self.retry_limit:
            try:
                await process_func(data)
                return  # 成功处理后退出
            except Exception as e:
                retry_count += 1
                print(f"处理数据时发生错误: {e}, 正在重试 {retry_count}/{self.retry_limit}, data: {data}")
                await asyncio.sleep(2 ** retry_count)
        print(f"处理数据失败, 达到重试上限, data: {data}")

    async def process_file(self, hdfs_file_path, process_func):
        """处理单个 gz 文件"""
        total_lines = self.count_total_lines(hdfs_file_path)
        processed_lines = 0
        start_time = time.time()
        #   # 这里根据不同的配置选用不同的读取文件的方法
        for lines in self.read_hdfs_fanc[self.read_hdfs_model](hdfs_file_path):
            processing_start_time = time.time()  # 记录本批处理开始时间

            tasks = []
            for line in lines:
                try:
                    data = json.loads(line)
                    tasks.append(self.process_data(data, process_func))
                except json.JSONDecodeError as e:
                    print(f"解析JSON失败: {e}, 行内容: {line.strip()}")

            # await AsyncTaskPool(self.batch_size).run(tasks) # AsyncTaskPool 适用于一次提交所有任务, 限制并发数执行
            await asyncio.gather(*tasks)

            processed_lines += len(lines)

            elapsed_time = time.time() - start_time  # 已用时间
            processing_time = time.time() - processing_start_time  # 本次处理时间
            avg_processing_time = (
                (elapsed_time * 1000) / processed_lines if processed_lines > 0 else float("inf")
            )  # 平均每条数据的处理时间（毫秒）

            # 估算剩余时间
            remaining_time = (
                ((avg_processing_time / 1000) * (total_lines - processed_lines))
                if processed_lines > 0
                else float("inf")
            )

            # 显示总进度信息
            print(
                f"文件: {hdfs_file_path} 总进度: {processed_lines}/{total_lines} 行 | "
                f"已用时间: {elapsed_time:.2f}秒 | 本次处理时间: {processing_time:.2f}秒 | "
                f"预估剩余时间: {remaining_time:.2f}秒 | 平均每条处理时间: {avg_processing_time:.2f}毫秒"
            )

        # 最终进度显示
        final_elapsed_time = time.time() - start_time  # 最终已用时间
        print(
            f"文件: {hdfs_file_path} 处理完成 | 总进度: {processed_lines}/{total_lines} 行 | "
            f"总已用时间: {final_elapsed_time:.2f}秒 | "
            f"平均每条处理时间: {(final_elapsed_time * 1000) / processed_lines:.2f}毫秒"
            if processed_lines > 0
            else "处理无数据"
        )

        self.save_processed_file(hdfs_file_path)  # 保存处理过的文件

    async def retry_process_file(self, hdfs_file_path, process_func):
        """带重试机制的文件处理"""
        retry_count = 0
        while retry_count < self.retry_limit:
            try:
                await self.process_file(hdfs_file_path, process_func)
                return True  # 成功处理后退出
            except Exception as e:
                retry_count += 1
                print(f"处理文件 {hdfs_file_path} 时发生错误: {e}，正在重试 {retry_count}/{self.retry_limit}")
                await asyncio.sleep(2 ** retry_count)
        print(f"处理文件 {hdfs_file_path} 失败，达到重试上限")
        return False
        # raise

    async def batch_process_file(self, hdfs_dir: str, process_func: Callable[[dict], Any]):
        """批量更新所有 gz 文件"""
        gz_files = self.list_gz_files(hdfs_dir)
        all_succeed = True
        for hdfs_file_path in gz_files:
            if self.is_file_processed(hdfs_file_path):
                print(f"跳过已处理文件: {hdfs_file_path}")
                continue  # 如果文件已处理，跳过
            succeed = await self.retry_process_file(hdfs_file_path, process_func)  # 处理文件
            if succeed is False:
                all_succeed = False

        if all_succeed:
            # 处理完成后删除数据库文件
            try:
                if os.path.exists(self.db_file):
                    os.remove(self.db_file)
                    print(f"已删除断点重试文件: {self.db_file}")
            except Exception as e:
                print(f"删除断点重试文件失败: {e}")

    async def process_file_bulk(self, hdfs_file_path, process_func):
        """按批次处理单个文件，批量数据传递给处理函数"""
        total_lines = self.count_total_lines(hdfs_file_path)
        processed_lines = 0
        start_time = time.time()

        tasks = []
        # 这里根据不同的配置选用不同的读取文件的方法
        for lines in self.read_hdfs_fanc[self.read_hdfs_model](hdfs_file_path):
            processing_start_time = time.time()  # 记录本批处理开始时间

            batch_data = []
            for line in lines:
                try:
                    data = json.loads(line)
                    batch_data.append(data)
                except json.JSONDecodeError as e:
                    print(f"解析JSON失败: {e}, 行内容: {line.strip()}")

            # 处理读取到的批次数据
            if batch_data:
                tasks.append(process_func(batch_data))  # 将批次数据传递给处理函数并收集任务
                processed_lines += len(batch_data)  # 更新已处理行数

            # 当积累的任务数量达到 batch_size 时并发处理所有任务
            if len(tasks) >= self.batch_size:
                await asyncio.gather(*tasks)  # 同时处理多个批次

                elapsed_time = time.time() - start_time  # 已用时间
                processing_time = time.time() - processing_start_time  # 本次处理时间
                avg_processing_time = (
                    (elapsed_time * 1000) / processed_lines if processed_lines > 0 else float("inf")
                )  # 平均每条数据的处理时间（毫秒）

                # 估算剩余时间
                remaining_time = (
                    ((avg_processing_time / 1000) * (total_lines - processed_lines))
                    if processed_lines > 0
                    else float("inf")
                )

                # 显示总进度信息
                print(
                    f"文件: {hdfs_file_path} 总进度: {processed_lines}/{total_lines} 行 | "
                    f"已用时间: {elapsed_time:.2f}秒 | 本次处理时间: {processing_time:.2f}秒 | "
                    f"预估剩余时间: {remaining_time:.2f}秒 | 平均每条处理时间: {avg_processing_time:.2f}毫秒"
                )

                # 清空任务列表，准备下一批处理
                tasks.clear()
            # 处理剩余的任务
        if tasks:
            await asyncio.gather(*tasks)  # 处理未达到 batch_size 的剩余任务

        # 最终进度显示
        final_elapsed_time = time.time() - start_time  # 最终已用时间
        print(
            f"文件: {hdfs_file_path} 处理完成 | 总进度: {processed_lines}/{total_lines} 行 | "
            f"总已用时间: {final_elapsed_time:.2f}秒 | "
            f"平均每条处理时间: {(final_elapsed_time * 1000) / processed_lines:.2f}毫秒"
            if processed_lines > 0
            else "处理无数据"
        )

        self.save_processed_file(hdfs_file_path)

    async def retry_process_file_bulk(self, hdfs_file_path, process_func):
        """带重试机制的批量文件处理"""
        retry_count = 0
        while retry_count < self.retry_limit:
            try:
                await self.process_file_bulk(hdfs_file_path, process_func)
                return True  # 成功处理后退出
            except Exception as e:
                retry_count += 1
                print(f"处理文件 {hdfs_file_path} 时发生错误: {e}，正在重试 {retry_count}/{self.retry_limit}")
                await asyncio.sleep(2 ** retry_count)
        print(f"处理文件 {hdfs_file_path} 失败，达到重试上限")
        return False

    async def batch_process_file_bulk(self, hdfs_dir: str, process_func: Callable[[List[dict]], Any]):
        """批量处理 gz 文件中的数据"""
        gz_files = self.list_gz_files(hdfs_dir)
        all_succeed = True
        for hdfs_file_path in gz_files:
            if self.is_file_processed(hdfs_file_path):
                print(f"跳过已处理文件: {hdfs_file_path}")
                continue  # 跳过已处理文件
            succeed = await self.retry_process_file_bulk(hdfs_file_path, process_func)
            if succeed is False:
                all_succeed = False

        if all_succeed:
            # 处理完成后删除数据库文件
            try:
                if os.path.exists(self.db_file):
                    os.remove(self.db_file)
                    print(f"已删除断点重试文件: {self.db_file}")
            except Exception as e:
                print(f"删除断点重试文件失败: {e}")

# # 使用示例
# async def update_refer(data: dict):
#     ref_id = data["ref_id"]
#     url = f"http://192.168.98.79:8150/v1/fact_refer/update/{ref_id}"
#     update_data = data["update_data"]
#     if not update_data:
#         return
#
#     # 此处为实际处理逻辑
#     await ApiNetUtils.fetch_post(url=url, payload=update_data)
#
#
# if __name__ == "__main__":
#     processor = HDFSDataProcessor()  # 实例化数据处理类
#     asyncio.run(processor.batch_process_file("/user/libaiyun/output/confidence", update_refer))
