"""
EPUB 清理器模块 (EpubCleaner)

负责在翻译前对 EPUB 进行标准化清洗，解决兼容性问题。
核心功能:
1. Flatten Structure: 将 div 转换为 p，简化结构
2. Fix TOC: 修复目录中的死链和缺失 UID
3. CSS Restoration: 找回丢失的样式表
"""

from bs4 import BeautifulSoup, Tag
from loguru import logger
from ebooklib import epub
import ebooklib
import zipfile
import uuid
from ebooklib.epub import Link

class EpubCleaner:
    """标准 EPUB 清理器"""
    
    def clean_epub(self, input_path: str, output_path: str):
        """
        清理 EPUB 文件并保存到新路径
        """
        logger.info(f"开始清理: {input_path}")
        
        # 1. 尝试打开 Zip 以读取原始内容 (Ebooklib 回退机制)
        try:
            input_zip = zipfile.ZipFile(input_path, 'r')
            zip_files = set(input_zip.namelist())
        except Exception as e:
            logger.error(f"无法打开 Zip (样式回退功能将失效): {e}")
            input_zip = None
            zip_files = set()
        
        # 2. 读取 EPUB
        try:
            book = epub.read_epub(input_path)
        except Exception as e:
            logger.error(f"Ebooklib 读取失败: {e}")
            raise

        # 3. 遍历并清理文档
        count = 0
        for item in book.get_items():
            if item.get_type() == ebooklib.ITEM_DOCUMENT:
                try:
                    file_name = item.get_name()
                    content = None
                    
                    # 优先从 Zip 读取以保留 Head 信息 (CSS Links)
                    if input_zip and file_name in zip_files:
                        try:
                            content = input_zip.read(file_name).decode('utf-8')
                        except Exception:
                            pass
                    
                    # 回退到 ebooklib
                    if content is None:
                        raw_content = item.get_content()
                        if raw_content:
                            content = raw_content.decode('utf-8')
                    
                    if not content or not content.strip():
                        continue
                    
                    # 执行清理
                    cleaned = self._clean_content(content, item)
                    
                    # 安全检查
                    if not cleaned.strip():
                        logger.warning(f"警告: {file_name} 清理后为空，保留原始内容")
                        cleaned = content
                    
                    item.set_content(cleaned.encode('utf-8'))
                    count += 1
                    
                except Exception as e:
                    logger.warning(f"清理文档失败 {item.get_name()}: {e}")
        
        # 4. 修复 TOC (死链和 UID)
        try:
            book.toc = self._fix_and_clean_toc(book.toc, book)
        except Exception as e:
            logger.error(f"TOC 修复失败: {e}")

        # 5. 保存
        epub.write_epub(output_path, book)
        logger.info(f"清理完成: {output_path} (处理了 {count} 个文档)")

    def _clean_content(self, html_content: str, item=None) -> str:
        """
        执行具体的 HTML 清理逻辑
        """
        soup = BeautifulSoup(html_content, 'html.parser')
        
        # 恢复 CSS 链接
        if item:
            self._restore_css_links(soup, item)
        
        # div 转 p
        stats = {'divs_to_p': 0}
        inline_tags = {'span', 'em', 'strong', 'i', 'b', 'u', 'a', 'br', 'sub', 'sup'}
        
        for div in list(soup.find_all('div')):
            # 检查是否有块级子元素 (如果有，则保留 div 容器结构)
            has_block = any(
                isinstance(c, Tag) and c.name not in inline_tags 
                for c in div.children
            )
            
            if not has_block:
                div.name = 'p'
                stats['divs_to_p'] += 1
                
        # 可以在这里添加更多清理逻辑 (如移除无用的空的 span 等)
        
        return str(soup)

    def _restore_css_links(self, soup, item):
        """从原始 HTML 中提取并恢复 CSS 链接到 item 对象"""
        head = soup.find('head')
        if head:
            links = head.find_all('link', rel='stylesheet')
            for link in links:
                href = link.get('href')
                if href:
                    existing_links = list(item.get_links())
                    exists = False
                    for l in existing_links:
                        l_href = getattr(l, 'href', None)
                        if l_href is None and isinstance(l, dict):
                            l_href = l.get('href')
                        if l_href == href:
                            exists = True
                            break
                    
                    if not exists:
                        item.add_link(href=href, rel='stylesheet', type='text/css')

    def _fix_and_clean_toc(self, toc, book):
        """修复 TOC：补全 UID 并移除指向不存在文件的死链"""
        new_toc = []
        
        for item in toc:
            # Case 1: (Section, Children)
            if isinstance(item, (tuple, list)):
                section, children = item
                cleaned_children = self._fix_and_clean_toc(children, book)
                
                if isinstance(section, Link):
                    href = section.href.split('#')[0]
                    if book.get_item_with_href(href):
                        if section.uid is None:
                            section.uid = f'uuid-{uuid.uuid4()}'
                        new_toc.append((section, cleaned_children))
                    else:
                        logger.warning(f"移除无效 TOC 节点: {section.href}")
                        new_toc.extend(cleaned_children)
                else:
                    new_toc.append((section, cleaned_children))

            # Case 2: Link
            elif isinstance(item, Link):
                href = item.href.split('#')[0]
                if book.get_item_with_href(href):
                     if item.uid is None:
                         item.uid = f'uuid-{uuid.uuid4()}'
                     new_toc.append(item)
                else:
                    logger.warning(f"移除无效 TOC 节点: {item.href}")
            
            # Case 3: Other
            else:
                new_toc.append(item)
                
        return new_toc