import logging
import os
from pathlib import Path
from typing import Any, List

from django.core.management.base import BaseCommand
from django.db import models, transaction

# 基本用法（从上次中断处继续）
# # 使用进度文件自动断点续传
# python manage.py migrate_cutwords --chunk-size 1000 --resume-file progress.json
# # 标签从第150000条开始，文章从第4000条开始
# python manage.py migrate_cutwords --model all --chunk-size 1000 --tag-start-id 150000 --article-start-id 4000
# 指定起始位置
# # 只处理文章，从第4000条开始
# python manage.py migrate_cutwords --model article --chunk-size 1500 --article-start-id 4000

# # 只处理标签，从第150000条开始
# python manage.py migrate_cutwords --model tag --chunk-size 2000 --tag-start-id 150000
# 结合进度文件使用
# # 第一次运行（会创建进度文件）
# python manage.py migrate_cutwords --chunk-size 1000 --resume-file migration_progress.json

# # 如果中途中断，再次运行相同的命令即可从断点继续
# python manage.py migrate_cutwords --chunk-size 1000 --resume-file migration_progress.json

logger = logging.getLogger(__name__)


class Command(BaseCommand):
    help = "海量数据安全迁移：支持断点续传的分批处理标签和文章的分词关联"

    def add_arguments(self, parser):
        # 创建项目专用的logs子目录
        MIGRATION_LOGS_DIR = Path("logs/migrations")
        MIGRATION_LOGS_DIR.mkdir(parents=True, exist_ok=True)
        parser.add_argument(
            "--chunk-size", type=int, default=1000, help="每批处理的数据量（默认: 1000）"
        )
        parser.add_argument(
            "--model",
            type=str,
            choices=["tag", "article", "all"],
            default="all",
            help="指定要处理的模型（默认: all）",
        )
        parser.add_argument(
            "--tag-start-id", type=int, default=1, help="标签数据起始ID（断点续传）"
        )
        parser.add_argument(
            "--article-start-id", type=int, default=1, help="文章数据起始ID（断点续传）"
        )
        parser.add_argument(
            "--resume-file",
            type=str,
            default=str(MIGRATION_LOGS_DIR / "cutwords_progress.json"),
            help="断点续传进度记录文件路径（默认: logs/migrations/cutwords_progress.json）",
        )

    def handle(self, *args, **options):
        chunk_size = options["chunk_size"]
        model_type = options["model"]
        resume_file = options["resume_file"]

        # 加载断点续传进度
        progress = self._load_progress(resume_file, options)

        try:
            self.stdout.write(f"开始海量数据迁移，批次大小: {chunk_size}")
            self.stdout.write(
                f'断点续传进度: 标签从ID {progress["tag_current_id"]} 开始, 文章从ID {progress["article_current_id"]} 开始'
            )

            if model_type in ["tag", "all"]:
                self.stdout.write("开始处理标签数据...")
                self._batch_process_tags_safe(chunk_size, progress)

            if model_type in ["article", "all"]:
                self.stdout.write("开始处理文章数据...")
                self._batch_process_articles_safe(chunk_size, progress)

            # 清理进度记录文件
            if resume_file:
                self._cleanup_progress_file(resume_file)

            self.stdout.write(self.style.SUCCESS("海量数据迁移完成！"))

        except Exception as e:
            self.stdout.write(self.style.ERROR(f"迁移失败: {str(e)}"))
            # 保存当前进度以便续传
            if resume_file:
                self._save_progress(resume_file, progress)
            raise

    def _load_progress(self, resume_file, options):
        """加载断点续传进度"""
        progress = {
            "tag_current_id": options["tag_start_id"],
            "article_current_id": options["article_start_id"],
            "tag_processed": 0,
            "article_processed": 0,
        }

        if resume_file:
            try:
                import json

                with open(resume_file, "r", encoding="utf-8") as f:
                    saved_progress = json.load(f)
                    progress.update(saved_progress)
                    self.stdout.write(f"从文件加载进度: {saved_progress}")
            except FileNotFoundError:
                self.stdout.write("未找到进度文件，使用命令行参数")
            except Exception as e:
                self.stdout.write(f"加载进度文件失败: {e}，使用默认值")

        return progress

    def _save_progress(self, resume_file, progress):
        """保存当前进度到文件"""
        try:
            import json

            with open(resume_file, "w", encoding="utf-8") as f:
                json.dump(progress, f, ensure_ascii=False, indent=2)
            self.stdout.write(f"进度已保存: {progress}")
        except Exception as e:
            self.stdout.write(f"保存进度失败: {e}")

    def _cleanup_progress_file(self, resume_file):
        """清理进度记录文件"""
        try:
            import os

            if os.path.exists(resume_file):
                os.remove(resume_file)
                self.stdout.write("进度文件已清理")
        except Exception as e:
            self.stdout.write(f"清理进度文件失败: {e}")

    def _batch_process_tags_safe(self, chunk_size: int, progress: dict):
        """安全处理海量标签数据（支持断点续传）"""
        from apps.content.models import CutWord, Tag

        tag_current_id = progress["tag_current_id"]
        total_processed = progress["tag_processed"]
        batch_count = 0

        # 使用ID范围查询，支持断点续传[1,2](@ref)
        while True:
            # 分批获取数据，避免内存溢出[4](@ref)
            tags_batch = list(
                Tag.objects.filter(id__gte=tag_current_id)
                .only("id", "title", "segmented_words")
                .order_by("id")[:chunk_size]
            )

            if not tags_batch:
                break  # 没有更多数据

            with transaction.atomic():
                for tag in tags_batch:
                    try:
                        # 处理当前标签
                        keywords = self.extract_keywords(tag.title)
                        tag.segmented_words = keywords

                        # 立即保存当前标签的更新
                        tag.save(update_fields=["segmented_words"])

                        # 处理分词关联
                        if keywords:
                            cutwords = []
                            for kw in keywords:
                                cutword, created = CutWord.objects.get_or_create(title=kw)
                                cutwords.append(cutword)

                            # 批量关联分词
                            if cutwords:
                                tag.cut_keywords.add(*cutwords)

                        total_processed += 1
                        batch_count += 1
                        tag_current_id = tag.id + 1  # 更新当前处理的ID

                        # 更新进度
                        progress["tag_current_id"] = tag_current_id
                        progress["tag_processed"] = total_processed

                    except Exception as e:
                        logger.error(f"处理标签 {tag.id} 时出错: {str(e)}")
                        continue

            # 每批处理完成后输出进度
            self.stdout.write(
                f"标签进度: 已处理 {total_processed} 个标签，当前ID: {tag_current_id}"
            )

            # 每处理10批后保存进度（可根据数据量调整）
            if batch_count >= chunk_size * 10:
                if progress.get("resume_file"):
                    self._save_progress(progress["resume_file"], progress)
                batch_count = 0

        self.stdout.write(f"标签处理完成，共处理 {total_processed} 个标签")

    def _batch_process_articles_safe(self, chunk_size: int, progress: dict):
        """修复变量名拼写错误的安全批处理文章数据"""
        from apps.content.models import Article, CutWord
        from django.db import transaction

        article_current_id = progress["article_current_id"]
        total_processed = progress["article_processed"]
        batch_count = 0

        while True:
            articles_batch = list(
                Article.objects.filter(id__gte=article_current_id)
                .only("id", "title")
                .order_by("id")[:chunk_size]
            )

            if not articles_batch:
                break

            try:
                with transaction.atomic():
                    # 在事务内部重新获取分词映射
                    existing_cutwords = CutWord.objects.in_bulk(field_name="title")
                    cutword_dict = {cw.title: cw for cw in existing_cutwords.values()}

                    cutwords_to_create = []
                    articles_to_process = []

                    # 第一步：收集所有需要处理的数据
                    for article in articles_batch:
                        keywords = self.extract_keywords(article.title)
                        if keywords:
                            for kw in keywords:
                                if kw not in cutword_dict:
                                    # 避免重复创建
                                    if kw not in [cw.title for cw in cutwords_to_create]:
                                        cutword = CutWord(title=kw)
                                        cutwords_to_create.append(cutword)

                        articles_to_process.append((article, keywords))

                    # 第二步：批量创建新分词（忽略冲突）
                    if cutwords_to_create:
                        try:
                            CutWord.objects.bulk_create(cutwords_to_create, ignore_conflicts=True)
                            # 修复：正确的变量名应该是 existing_cutwords，不是 exful_cutwords
                            # 重新查询获取所有分词（包括新创建的和已存在的）
                            all_keywords = set()
                            for article, keywords in articles_to_process:
                                if keywords:
                                    all_keywords.update(keywords)

                            if all_keywords:
                                # 修复拼写错误：使用正确的变量名
                                updated_cutwords = CutWord.objects.filter(
                                    title__in=list(all_keywords)
                                ).in_bulk(field_name="title")
                                cutword_dict.update(updated_cutwords)  # 正确的变量名
                        except Exception as e:
                            logger.warning(f"批量创建分词时忽略冲突: {e}")
                            # 即使有冲突也继续执行，使用已存在的分词

                    # 第三步：处理文章关联
                    for article, keywords in articles_to_process:
                        if keywords:
                            cutwords_to_add = []
                            for kw in keywords:
                                if kw in cutword_dict:
                                    cutwords_to_add.append(cutword_dict[kw])

                            if cutwords_to_add:
                                try:
                                    article.cut_keywords.add(*cutwords_to_add)
                                except Exception as e:
                                    logger.warning(f"文章 {article.id} 添加分词关联失败: {e}")
                                    continue

                        total_processed += 1
                        article_current_id = article.id + 1
                        progress.update(
                            {
                                "article_current_id": article_current_id,
                                "article_processed": total_processed,
                            }
                        )

            except Exception as e:
                logger.error(
                    f"处理文章批次时出错（ID范围: {articles_batch[0].id}-{articles_batch[-1].id}）: {str(e)}"
                )
                article_current_id = (
                    articles_batch[-1].id + 1 if articles_batch else article_current_id
                )
                continue

            # 进度输出
            self.stdout.write(
                f"文章进度: 已处理 {total_processed} 篇文章，当前ID: {article_current_id}"
            )

            # 定期保存进度
            if batch_count >= chunk_size * 10:
                if progress.get("resume_file"):
                    self._save_progress(progress["resume_file"], progress)
                batch_count = 0

        self.stdout.write(f"文章处理完成，共处理 {total_processed} 篇文章")

    def extract_keywords(self, text):
        """提取关键词（你的现有逻辑）"""
        from apps.content.utils.extract_keywords import extract_keywords

        return extract_keywords(text)
