"""
# Animation Library
A simple 2D frame animation tool that supports loading images, automatic playback, and displaying animations.
And if you need help in the editor, use the animation_help() function to get assistance. Go ahead and try it out!

Usage Example:
    add_file_path("Path", "Animation Name")
    animation_start("Animation Name", Total Frames)
    image_show("Animation Name", x, y, Window)
Exceptions:
    NotFoundFileError - Raised when the file does not exist
    UnboundWindowError - Raised when the window is not bound
    ImageNotShowError - Raised when the image is not displayed
    NotFoundAliasError - Raised when the alias does not exist
    ValueError - Raised when the input data is inconsistent with the facts
"""

import os
import pathlib
import json
from typing import Optional, Callable, Literal
import zipfile
import io

def _red(text):return f"\033[31m{text}\033[0m"   # 红色
def _green(text):return f"\033[32m{text}\033[0m" # 绿色
def _yellow(text):return f"\033[33m{text}\033[0m"# 黄色
def _blue(text):return f"\033[34m{text}\033[0m"  # 蓝色
def _purple(text):return f"\033[35m{text}\033[0m"# 紫色
def _cyan(text):return f"\033[36m{text}\033[0m"  # 青色
def _white(text):return f"\033[37m{text}\033[0m"# 白色
def _gray(text):return f"\033[90m{text}\033[0m"  # 灰色

class NotFoundFileError(Exception): __module__ = "builtins"
class UnboundWindowError(Exception): __module__ = "builtins"
class ImageNotShowError(Exception): __module__ = "builtins"
class NotFoundAliasError(Exception): __module__ = "builtins"

_number = 0

_file_path_and_alias = {}
_image_number = {}
_first_image_number = {}
_image_styling = {}
_first_image_wait = {}
_image_wait = {}
_already_started = {}
_is_reverse = {}
_animation_speed = {}

_image_is_show = []
_image_is_start = []
_image_type = []

__next_number = {}

# ----------------- 全局动画缓存机制 -----------------
_animation_data = {}      # 缓存解析后的 JSON 字典格式动画数据
_animation_surfaces = {}  # 缓存加载后的 Pygame Surface 纹理防止重复读取卡顿
_animation_zip_files = {}   # 存储 ZipFile 对象，alias -> ZipFile

print("\n\tHello, Animation. BiliBili @西瓜汁Code")

try:
    import pygame
except ImportError:
    print(_red("""Pygame is not installed. Please install it first. Download command:
               pip install pygame"""))
    exit()

__again = False
__help_language = ""

def animation_help():
    global __again, __help_language, ic
    
    # 获取txt文件路径
    script_dir = os.path.dirname(os.path.abspath(__file__))
    help_dir = pathlib.Path(script_dir, "txt", "help")
    
    # 确保目录存在
    os.makedirs(help_dir, exist_ok=True)
    
    # 解析标记文件的函数
    def load_help_data(lang):
        data = {
            'title': '',
            'prompt': '',
            'sub_prompt': '',
            'class_prompt': '',
            'groups': {},
            'items': {},
            'funcs': {},
            'descs': {}
        }
        
        filename = f"help_{lang}.txt"
        filepath = pathlib.Path(help_dir, filename)
        
        if not os.path.exists(filepath):
            print(f"错误：找不到帮助文件 {filepath}")
            return None
        
        with open(filepath, 'r', encoding='utf-8') as f:
            for line in f:
                line = line.strip()
                if not line or line.startswith('#'):
                    continue
                
                parts = line.split('|')
                if not parts:
                    continue
                
                cmd = parts[0]
                
                if cmd == 'title':
                    data['title'] = parts[1]
                elif cmd == 'prompt':
                    data['prompt'] = parts[1]
                elif cmd == 'sub_prompt':
                    data['sub_prompt'] = parts[1]
                elif cmd == 'class_prompt':
                    data['class_prompt'] = parts[1]
                elif cmd == 'group':
                    data['groups'][parts[1]] = parts[2]
                elif cmd == 'item':
                    data['items'][(parts[1], parts[2])] = parts[3]
                elif cmd == 'func':
                    data['funcs'][parts[1]] = parts[2]
                elif cmd == 'desc':
                    if len(parts) == 4:
                        data['descs'][(parts[1], parts[2])] = parts[3].replace('\\n', '\n')
                    elif len(parts) == 3:
                        data['descs'][parts[1]] = parts[2].replace('\\n', '\n')
        
        return data
    
    # 显示菜单的函数
    def show_menu(data):
        lines = [data['title']]
        
        # 显示分组和项目
        for group_id, group_title in data['groups'].items():
            lines.append(f"{group_id}.{group_title}:")
            for (gid, item_id), item_text in data['items'].items():
                if gid == group_id:
                    lines.append(f"    -  {item_id}.{item_text}")
        
        # 显示独立函数
        for func_id, func_text in data['funcs'].items():
            lines.append(f"{func_id}.{func_text}")
        
        print('\n'.join(lines))
        return input(data['prompt'])
    
    # 显示子菜单
    def show_submenu(data, group_id):
        lines = [data['sub_prompt']]
        for (gid, item_id), item_text in data['items'].items():
            if gid == group_id:
                lines.append(f"{item_id}.{item_text}")
        
        print('\n'.join(lines))
        return input()
    
    # 显示类菜单
    def show_class_menu(data):
        lines = [data['class_prompt']]
        for (gid, item_id), item_text in data['items'].items():
            if gid == "11":
                lines.append(f"{item_id}.{item_text}")
        
        print('\n'.join(lines))
        return input()
    
    # 首次运行选择语言
    if __again == False:
        ic = input(f"Select language (cn/en): ")
        __help_language = ic
        if __help_language != "cn" and __help_language != "en":
            print("Please enter cn or en.")
            return
        help_data = load_help_data(__help_language)
        if help_data is None:
            return
        __again = True
    else:
        help_data = load_help_data(__help_language)
        if help_data is None:
            return
    
    # 使用循环代替递归
    while True:
        ic = show_menu(help_data)
        
        if ic == "exit":
            break
        
        # 处理返回值函数组
        if ic == "1":
            sub_ic = show_submenu(help_data, "1")
            desc_key = ("1", sub_ic)
            if desc_key in help_data['descs']:
                print(help_data['descs'][desc_key])
            else:
                print("帮助内容不存在")
            input("按任意键继续..." if __help_language == "cn" else "Press any key to continue...")
        
        # 处理类组
        elif ic == "11":
            sub_ic = show_class_menu(help_data)
            desc_key = ("11", sub_ic)
            if desc_key in help_data['descs']:
                print(help_data['descs'][desc_key])
            else:
                print("帮助内容不存在")
            input("按任意键继续..." if __help_language == "cn" else "Press any key to continue...")
        
        # 处理独立函数
        elif ic in help_data['funcs']:
            if ic in help_data['descs']:
                print(help_data['descs'][ic])
            else:
                print("帮助内容不存在")
            input("按任意键继续..." if __help_language == "cn" else "Press any key to continue...")
        
        else:
            print("无效的输入，请重新选择" if __help_language == "cn" else "Invalid input, please choose again")


def is_image_showing(alias : str):
    if alias in _image_number:
        return _image_is_show[_image_number[alias]]
    else:
        raise NotFoundAliasError(f"Not found alias: {alias}")
    
def is_animation_started(alias : str):
    if alias in _image_number:
        return _image_is_start[_image_number[alias]]
    else:
        raise NotFoundAliasError(f"Not found alias: {alias}")

def get_animation_frame(alias : str):
    if alias in _image_number:
        return _image_number[alias] + 1
    else:
        raise NotFoundAliasError(f"Not found alias: {alias}")

def _delete_number(string : str):
    number = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
    for num in number:
        string = string.replace(str(num), "")
    return string

def _get_file_name(list : list, _ : bool = True):
    if len(list) <= 1:
        list.append("")
    string = ""
    if _:
        for i in range(len(list)-1):
            if i != len(list)-2:
                string = string + list[i] + "_"
            else:
                string = string + list[i]
    else:
        return _delete_number(list[0])
    return string

def add_file_path(file_path : str, alias : str, wait : int = 10):
    global _number
    if os.path.exists(file_path):
        file_ext = os.path.splitext(os.path.basename(file_path))[1].lower()
        
        # 处理传统 JSON
        if file_ext == ".json":
            _image_type.append("Animation")
            with open(file_path, "r", encoding="utf-8") as f:
                _animation_data[alias] = json.load(f)
                
        # 处理 ZIP 格式动画包
        elif file_ext == ".animdata":
            _image_type.append("Animation")
            zf = zipfile.ZipFile(file_path, 'r')
            _animation_zip_files[alias] = zf
            with zf.open("data.json") as f:
                _animation_data[alias] = json.load(f)
        else:
            _image_type.append("Image")

        if not wait < 0:
            _file_path_and_alias[alias] = file_path
            _image_styling[alias] = file_path
            _image_number[alias] = _number
            _first_image_number[alias] = _number
            _first_image_wait[alias] = wait
            _image_wait[alias] = wait
            _already_started[alias] = False
            _is_reverse[alias] = False
            _animation_speed[alias] = 1
            __next_number[alias] = 0
            _number += 1
            _image_is_show.append(False)
            _image_is_start.append(False)
        else:
            raise ValueError("Wait time cannot be negative.")
    else:
        raise NotFoundFileError(f"Not found file: {file_path}")

# ----------------- 加载 Pygame 纹理缓存函数 (适配 ZIP) -----------------
def _get_cached_surface(original_path: str, alias: str):
    global _animation_surfaces
    cache_key = f"{alias}_{original_path}"
    if cache_key in _animation_surfaces:
        return _animation_surfaces[cache_key]
    
    # 策略 1: 尝试从 ZIP 压缩包中直接读取内存流
    if alias in _animation_zip_files:
        try:
            zf = _animation_zip_files[alias]
            img_data = zf.read(original_path)
            surf = pygame.image.load(io.BytesIO(img_data)).convert_alpha()
            _animation_surfaces[cache_key] = surf
            return surf
        except KeyError:
            pass # 可能路径不在 zip 中

    # 策略 2: 尝试从本地磁盘路径读取
    if os.path.exists(original_path):
        try:
            surf = pygame.image.load(original_path).convert_alpha()
            _animation_surfaces[cache_key] = surf
            return surf
        except Exception:
            pass
            
    return None

def _render_composite_animation(alias: str, x: int, y: int, window: pygame.Surface):
    data = _animation_data.get(alias)
    if not data:
        return
        
    frame_idx = _image_number.get(alias, 0) + 1
    frame_str = str(frame_idx)
    
    editor_w = data.get("editor_width", 1261)
    editor_h = data.get("editor_height", 815)
    
    obj_keys = [k for k in data.keys() if k.startswith("Object")]
    try:
        obj_keys.sort(key=lambda k: int(k.replace("Object", "")))
    except ValueError:
        obj_keys.sort()
        
    for obj_key in obj_keys:
        obj_val = data[obj_key]
        show_frames = obj_val.get("show_frame", "").split(";")
        
        if frame_str not in show_frames:
            continue
            
        frame_details = obj_val.get("frame", {}).get(frame_str)
        if not frame_details:
            continue
            
        orig_path = obj_val.get("path", "")
        if not orig_path:
            continue
            
        surf = _get_cached_surface(orig_path, alias)
        if not surf:
            continue
            
        pos = frame_details.get("position", {})
        rel_x = pos.get("rel_x", 0.0)
        rel_y = pos.get("rel_y", 0.0)
        direction = frame_details.get("direction", 0.0)
        size_info = frame_details.get("size", {})
        size_x = size_info.get("size_x", surf.get_width())
        size_y = size_info.get("size_y", surf.get_height())
        
        draw_x = x + rel_x * editor_w
        draw_y = y + rel_y * editor_h
        
        try:
            scaled_surf = pygame.transform.scale(surf, (int(size_x), int(size_y)))
            if direction != 0.0:
                rotated_surf = pygame.transform.rotate(scaled_surf, direction)
            else:
                rotated_surf = scaled_surf
            rect = rotated_surf.get_rect(center=(int(draw_x + size_x / 2), int(draw_y + size_y / 2)))
            window.blit(rotated_surf, rect)
        except Exception as e:
            print(f"Error rendering {obj_key} for animation {alias}: {e}")


def image_show(alias : str, x : int, y : int, window : pygame.Surface):
    if _image_number[alias] >= len(_image_is_show):
        _image_is_show.extend([False] * (_image_number[alias] - len(_image_is_show) + 1))
    if alias in _file_path_and_alias:
        _image_is_show[_first_image_number[alias]] = True
        actual_type = _image_type[_first_image_number[alias]]
        if actual_type == "Animation":
            _render_composite_animation(alias, x, y, window)
        else:
            image = pygame.image.load(_image_styling[alias])
            window.blit(image, (x, y))
    else:
        raise NotFoundAliasError(f"Not found alias: {alias}")

def animation_start(alias : str, total_frames : int = 1, callback : Optional[Callable[[], None]] = None,  
                    _ : bool = True, file_type : str = "png"):
    if alias not in _image_number:
        raise ImageNotShowError(f"Image not show: {alias} (alias not initialized)")

    if _image_number[alias] >= len(_image_is_show):
        _image_is_show.extend([False] * (_image_number[alias] - len(_image_is_show) + 1))

    if _image_is_show[_first_image_number[alias]] == True:
        if alias in _file_path_and_alias:
            _image_is_start[_first_image_number[alias]] = True
            _already_started[alias] = True
            if _image_wait[alias] <= 0:
                _image_wait[alias] = _first_image_wait[alias] // _animation_speed[alias]
                file_base_name : str = os.path.basename(_file_path_and_alias[alias])
                file_name  = file_base_name.split(".")[0].split("_")
                file_name = _delete_number(_get_file_name(file_name))

                # === 自动根据文件类型获取安全、正确的 total_frames ===
                actual_type = _image_type[_first_image_number[alias]]
                if actual_type == "Animation":
                    # 从已缓存的 JSON 数据中自动获取总帧数 total_frame
                    if alias in _animation_data:
                        total_frames = int(_animation_data[alias].get("total_frame", total_frames))
                    else:
                        # 兜底：如果还未加载则在此重新热加载一次
                        json_path = _file_path_and_alias.get(alias)
                        if json_path and os.path.exists(json_path):
                            try:
                                with open(json_path, "r", encoding="utf-8") as f:
                                    _animation_data[alias] = json.load(f)
                                total_frames = int(_animation_data[alias].get("total_frame", total_frames))
                            except Exception as e:
                                print(f"Error loading target animation file: {e}")
                
                # 获取当前的image_number
                current_number = _image_number[alias]

                if _is_reverse[alias] == False:
                    # 计算下一个image_number，如果超过count则重置为0
                    __next_number[alias] = (current_number + 1) % total_frames
                    if __next_number[alias] == 0:
                        if callback:
                            callback()
                elif _is_reverse[alias] == True:
                    # 计算上一个image_number，如果小于1则重置为count-1
                    __next_number[alias] = (current_number - 1) % total_frames
                    if __next_number[alias] == total_frames-1:
                        if callback:
                            callback()

                # 修正原开发版中的 string / integer 混用列表寻址 Bug
                for i in range(1):  # 只更新一次
                    if actual_type == "Image":
                        string = ""
                        if _ == True:
                            string = file_name + "_" + str(__next_number[alias] + 1) + "." + file_type
                        else:
                            string = file_name + str(__next_number[alias] + 1) + "." + file_type

                        _image_styling[alias] = os.path.join(os.path.dirname(_file_path_and_alias[alias]), string)

                        if not os.path.exists(_image_styling[alias]):
                            raise NotFoundFileError(f"Not found file: {_image_styling[alias]}")
                    elif actual_type == "Animation":
                        pass

                # 更新_image_number[alias]为next_number
                _image_number[alias] = __next_number[alias]

            else:
                _image_wait[alias] -= 1
        else:
            raise NotFoundAliasError(f"Not found alias: {alias}")
    elif _already_started[alias] == True:
        pass
    else:
        raise ImageNotShowError(f"Image not show: {_file_path_and_alias[alias]}")

def animation_stop(alias : str):
    if alias in _file_path_and_alias:
        _image_is_start[_first_image_number[alias]] = False
        _image_wait[alias] = _first_image_wait[alias]
    else:
        raise NotFoundAliasError(f"Not found alias: {alias}")
    
def set_animation_frame(alias : str, frame : int = 1):
    if alias in _file_path_and_alias:
        _image_number[alias] = frame-2
    else:
        raise NotFoundAliasError(f"Not found alias: {alias}")

def set_animation_speed(alias : str, speed : int = 2):
    if alias in _file_path_and_alias:
        if speed > 0:
            if isinstance(speed, int):
                _animation_speed[alias] = speed
            else:
                raise ValueError(f"Speed cannot be {type(speed)}. It int of type")
        else:
            raise ValueError("Speed cannot be negative.")
    else:
        raise NotFoundAliasError(f"Not found alias: {alias}")

def animation_reverse(alias : str):
    if alias in _file_path_and_alias:
        _is_reverse[alias] = True
    else:
        raise NotFoundAliasError(f"Not found alias: {alias}")
    
def animation_reverse_back(alias : str):

    if alias in _file_path_and_alias:
        _is_reverse[alias] = False
    else:
        raise NotFoundAliasError(f"Not found alias: {alias}")

class AnimationStatus:
    def __init__(self, alias : str):
        if alias in _file_path_and_alias:
            self.alias = alias
            self.is_showing = is_image_showing(alias)
            self.is_started = is_animation_started(alias)
            self.frame = get_animation_frame(alias)
            self.speed = _animation_speed[alias]
            self.is_reverse = _is_reverse[alias]
            self.file_path = _file_path_and_alias[alias]
            self.file_name = os.path.basename(self.file_path).split(".")[0].split("_")
            self.file_name = _delete_number(_get_file_name(self.file_name))
            self.total_frames = len(os.listdir(os.path.dirname(self.file_path))) - 1
            self.current_frame = self.frame - 1
            self.image_type = _image_type[_first_image_number[alias]]
        else:
            raise NotFoundAliasError(f"Not found alias: {alias}")

    def get_settings(self):
        return json.dumps({
            "alias": self.alias,
            "is_showing": self.is_showing,
            "is_started": self.is_started,
            "frame": self.frame,
            "speed": self.speed,
            "is_reverse": self.is_reverse,
            "file_path": str(self.file_path),
            "file_name": self.file_name,
            "total_frames": self.total_frames,
            "current_frame": self.current_frame, 
            "image_type": self.image_type
        })

    def get_status(self, status : Literal["animation_started", 
                                          "image_showing", 
                                          "animation_start", 
                                          "animation_end", 
                                          "animation_reverse"]):
        
        if status == "animation_started":
            return self.is_started
        elif status == "image_showing":
            return self.is_showing
        elif status == "animation_start":
            return self.current_frame == 0
        elif status == "animation_end":
            return self.current_frame == self.total_frames - 1
        elif status == "animation_reverse":
            return self.is_reverse
        else:
            raise ValueError(f"Not found status: {status}")

__all__ = [
    'animation_help', 
    'add_file_path', 
    'image_show', 
    'animation_start', 
    'animation_stop', 
    'set_animation_frame', 
    'set_animation_speed', 
    'animation_reverse', 
    'animation_reverse_back',
    'AnimationStatus'
]