from contextlib import closing
from typing import Dict, List
from scripts.iib.db.datamodel import Image as DbImg, Tag, ImageTag, DataBase, Folder
import os
from scripts.iib.tool import (
    is_valid_media_path,
    get_modified_date,
    get_video_type,
    is_dev,
    get_modified_date,
    is_image_file,
    is_audio_file,
    case_insensitive_get,
    get_img_geninfo_txt_path,
    parse_generation_parameters
)
from scripts.iib.parsers.model import ImageGenerationInfo, ImageGenerationParams
from scripts.iib.logger import logger
from scripts.iib.parsers.index import parse_image_info
from scripts.iib.plugin import plugin_inst_map
from scripts.iib.auto_tag import AutoTagMatcher

# 定义一个函数来获取图片文件的EXIF数据
def get_exif_data(file_path):
    if get_video_type(file_path):
        # 对于视频文件，尝试读取对应的txt标签文件
        txt_path = get_img_geninfo_txt_path(file_path)
        if txt_path:
            try:
                with open(txt_path, 'r', encoding='utf-8') as f:
                    content = f.read().strip()
                    if content:
                        # 复用现有解析逻辑，添加视频标识
                        params = parse_generation_parameters(content + "\nSource Identifier: Video Tags")
                        return ImageGenerationInfo(
                            content,
                            ImageGenerationParams(
                                meta=params["meta"],
                                pos_prompt=params["pos_prompt"],
                                extra=params,
                            ),
                        )
            except Exception as e:
                if is_dev:
                    logger.error("Failed to read video txt file %s: %s", txt_path, e)
        return ImageGenerationInfo()
    try:
        return parse_image_info(file_path)
    except Exception as e:
        if is_dev:
            logger.error("get_exif_data %s", e)
    return ImageGenerationInfo()


def update_image_data(search_dirs: List[str], is_rebuild = False):
    conn = DataBase.get_conn()
    tag_incr_count_rec: Dict[int, int] = {}

    if is_rebuild:
        Folder.remove_all(conn)

    def safe_save_img_tag(img_tag: ImageTag):
        tag_incr_count_rec[img_tag.tag_id] = (
            tag_incr_count_rec.get(img_tag.tag_id, 0) + 1
        )
        img_tag.save_or_ignore(conn)  # 原先用来处理一些意外，但是写的正确完全没问题,去掉了try catch

    # 递归处理每个文件夹
    def process_folder(folder_path: str):
        if not Folder.check_need_update(conn, folder_path):
            return
        print(f"Processing folder: {folder_path}")
        for filename in os.listdir(folder_path):
            file_path = os.path.normpath(os.path.join(folder_path, filename))
            try:

                if os.path.isdir(file_path):
                    process_folder(file_path)
                elif is_valid_media_path(file_path):
                    build_single_img_idx(conn, file_path, is_rebuild, safe_save_img_tag)
                # neg暂时跳过感觉个没人会搜索这个
            except Exception as e:
                logger.error("Tag generation failed. Skipping this file. file:%s error: %s", file_path, e)
        # 提交对数据库的更改
        Folder.update_modified_date_or_create(conn, folder_path)
        conn.commit()

    for dir in search_dirs:
        process_folder(dir)
        conn.commit()
    for tag_id in tag_incr_count_rec:
        tag = Tag.get(conn, tag_id)
        tag.count += tag_incr_count_rec[tag_id]
        tag.save(conn)
    conn.commit()

def add_image_data_single(file_path):
    conn = DataBase.get_conn()
    tag_incr_count_rec: Dict[int, int] = {}

    def safe_save_img_tag(img_tag: ImageTag):
        tag_incr_count_rec[img_tag.tag_id] = (
            tag_incr_count_rec.get(img_tag.tag_id, 0) + 1
        )
        img_tag.save_or_ignore(conn)

    file_path = os.path.normpath(file_path)
    try:
        if not is_valid_media_path(file_path):
            return
        build_single_img_idx(conn, file_path, False, safe_save_img_tag)
        # neg暂时跳过感觉个没人会搜索这个
    except Exception as e:
        logger.error("Tag generation failed. Skipping this file. file:%s error: %s", file_path, e)
        conn.commit()

    for tag_id in tag_incr_count_rec:
        tag = Tag.get(conn, tag_id)
        tag.count += tag_incr_count_rec[tag_id]
        tag.save(conn)
    conn.commit()

def rebuild_image_index(search_dirs: List[str]):
    conn = DataBase.get_conn()
    with closing(conn.cursor()) as cur:
        cur.execute(
            """DELETE FROM image_tag
            WHERE image_tag.tag_id IN (
                SELECT tag.id FROM tag WHERE tag.type <> 'custom'
            )
            """
        )
        cur.execute("""DELETE FROM tag WHERE tag.type <> 'custom'""")
        conn.commit()
        update_image_data(search_dirs=search_dirs, is_rebuild=True)


def get_extra_meta_keys_from_plugins(source_identifier: str):
    try:
        plugin = plugin_inst_map.get(source_identifier)
        if plugin:
            return plugin.extra_convert_to_tag_meta_keys
    except Exception as e:
        logger.error("get_extra_meta_keys_from_plugins %s", e)
    return []

def build_single_img_idx(conn, file_path, is_rebuild, safe_save_img_tag):
    img = DbImg.get(conn, file_path)

    if img and is_rebuild and img.exif_edited:
        logger.info(f"Image {file_path} has been manually edited, skipping rebuild.")
        return

    parsed_params = None
    if is_rebuild:
        info = get_exif_data(file_path)
        parsed_params = info.params
        if not img:
            img = DbImg(
                file_path,
                info.raw_info,
                os.path.getsize(file_path),
                get_modified_date(file_path),
            )
            img.save(conn)
    else:
        if img:  # 已存在的跳过
            if img.date == get_modified_date(img.path):
                return
            else:
                DbImg.safe_batch_remove(conn=conn, image_ids=[img.id])
        info = get_exif_data(file_path)
        parsed_params = info.params
        img = DbImg(
            file_path,
            info.raw_info,
            os.path.getsize(file_path),
            get_modified_date(file_path),
        )
        img.save(conn)

    if not parsed_params:
        return
    meta = parsed_params.meta
    lora = parsed_params.extra.get("lora", [])
    lyco = parsed_params.extra.get("lyco", [])
    if "final_width" in meta and "final_height" in meta:
        size_str = str(meta["final_width"]) + " × " + str(meta["final_height"])
    else:
        size_str = "Unknown Size"
    pos = parsed_params.pos_prompt
    size_tag = Tag.get_or_create(
        conn,
        size_str,
        type="size",
    )
    if size_tag:
        safe_save_img_tag(ImageTag(img.id, size_tag.id))
    # 确定媒体类型：Image / Video / Audio / Unknown
    if is_image_file(file_path):
        media_type_name = "Image"
    elif is_audio_file(file_path):
        media_type_name = "Audio"
    elif get_video_type(file_path):
        media_type_name = "Video"
    else:
        media_type_name = "Unknown"
    media_type_tag = Tag.get_or_create(conn, media_type_name, 'Media Type')
    if media_type_tag:
        safe_save_img_tag(ImageTag(img.id, media_type_tag.id))
    keys = [
        "Model",
        "Sampler",
        "Source Identifier",
        "Postprocess upscale by",
        "Postprocess upscaler",
        "Size",
        "Refiner",
        "Hires upscaler"
    ]
    keys += get_extra_meta_keys_from_plugins(meta.get("Source Identifier", ""))
    for k in keys:
        v = case_insensitive_get(meta, k)
        if not v:
            continue
        
        tag = Tag.get_or_create(conn, str(v), k)
        if tag:
            safe_save_img_tag(ImageTag(img.id, tag.id))
            if "Hires upscaler" == k:
                tag = Tag.get_or_create(conn, 'Hires All', k)
                if tag:
                    safe_save_img_tag(ImageTag(img.id, tag.id))
            elif "Refiner" == k:
                tag = Tag.get_or_create(conn, 'Refiner All', k)
                if tag:
                    safe_save_img_tag(ImageTag(img.id, tag.id))
    for i in lora:
        tag = Tag.get_or_create(conn, i["name"], "lora")
        if tag:
            safe_save_img_tag(ImageTag(img.id, tag.id))
    for i in lyco:
        tag = Tag.get_or_create(conn, i["name"], "lyco")
        if tag:
            safe_save_img_tag(ImageTag(img.id, tag.id))
    for k in pos:
        tag = Tag.get_or_create(conn, k, "pos")
        if tag:
            safe_save_img_tag(ImageTag(img.id, tag.id))
    
    AutoTagMatcher.get_instance(conn).apply(img.id, parsed_params)