import json
import os
import shutil
import subprocess
import tempfile
import uuid
from functools import wraps
import piexif
import base64
from PIL import Image
from PIL.PngImagePlugin import PngInfo
from io import BytesIO
from threading import Lock, Timer
from nanoid import generate
from datetime import datetime
from awe.util.logging import logger
from awe.oss import upload_file
from awe.util.http import get_content, get_json, post_multipart

WATERMARK_SUBJECT = "91330110MA7DR7497J"
MEDIA_WATERMARK_TAG = "comment"
PNG_WATERMARK_TAG = "Comment"
MEDIA_FILE_EXTENSIONS = {"mp4", "webm", "mp3"}

def get_remote_json(url):
  return get_json(url)

def get_remote_base64_image(url):
  if url is None:
    logger.error("URL is None, cannot get image.")
    return None
  if url.startswith("http"):
    image_data = get_content(url)
    image_bytes = BytesIO(image_data)
    image = Image.open(image_bytes)
    buff = BytesIO()
    image.save(buff, format=image.format)
    byte_data = buff.getvalue()
    return base64.b64encode(byte_data).decode('utf-8')
  
  # 判断是否是 base64 编码的图片
  try:
    # 尝试解码base64字符串
    base64.b64decode(url)
    # 如果没有抛出异常，则是base64编码的图片
    return url
  except Exception as e:
    logger.error(f"URL is not valid, cannot get image. {e}")
    return None

def generate_filename_with_date(file_format):
  file_id = generate('1234567890abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ', 10)
  dir_path = datetime.now().strftime("%Y/%m/%d")
  return f"{dir_path}/{file_id}.{file_format}"

def get_file_extension(file_name):
  return file_name.split('.')[-1]

def create_dir_if_not_exists(directory):
  try:
    os.makedirs(directory, exist_ok=True)
    logger.info(f"确保目录 '{directory}' 存在。")
  except OSError as error:
    logger.error(f"创建目录时发生错误：{error}")
  return directory

def is_animated_webp(image):
  return getattr(image, 'is_animated', False)

def compress_png(image_path, output_path):
  try:
    with Image.open(image_path) as img:
      if img.mode == 'RGBA':
        try:
          img.save(output_path, format='WEBP', lossless=True, transparency=img.info.get('transparency'))
        except Exception:
          img.save(output_path, format='WEBP')
      else:
        img.save(output_path, format='WEBP')
    logger.info(f"compress_png succeeded: {output_path}")
  except Exception as e:
    logger.error(f"compress_png failed: {e}")

def add_img_source(source_path, target_path, meta: dict):
  now = datetime.now().strftime("%Y:%m:%d %H:%M:%S")
  watermark = meta.get("watermark")
  img_description = json.dumps(watermark, ensure_ascii=False) if watermark is not None else meta.get("description", "")
  img_artist = meta.get("artist", "")
  img_software = meta.get("software", "")
  img_copyright = meta.get("copyright", "")
  img_datetime = meta.get("DateTime", now)
  img_datetime_original = meta.get("DateTimeOriginal", now)
  comment = f"{img_description}\n{img_artist}\n{img_software}\n{img_copyright}"
  with Image.open(source_path) as image:
    image_format = (image.format or "").upper()

  if image_format == "PNG" or str(target_path).lower().endswith(".png"):
    return add_png_source(source_path, target_path, comment)
  
  exif_info = {
    "0th": {
      piexif.ImageIFD.ImageDescription: comment,
      piexif.ImageIFD.Artist: img_artist,
      piexif.ImageIFD.Software: img_software,
      piexif.ImageIFD.Copyright: img_copyright,
      piexif.ImageIFD.DateTime: img_datetime,
    },
    "Exif": {
      piexif.ExifIFD.DateTimeOriginal: img_datetime_original,
    }
  }

  # 试图把EXIF信息转为字节流
  try:
      exif_bytes = piexif.dump(exif_info)
      if exif_bytes is not None:
        # 读取图片
        img = Image.open(source_path)
        # 把EXIF信息写入图片
        img.save(target_path, exif=exif_bytes)
  except ValueError as e:
      logger.error(f"add_img_source failed: {e}")
      exif_bytes = b""

def add_png_source(source_path, target_path, comment):
  png_info = PngInfo()
  png_info.add_text(PNG_WATERMARK_TAG, comment)
  with Image.open(source_path) as image:
    image.save(target_path, format='PNG', pnginfo=png_info)
  return target_path

def parse_hidden_watermark(raw_value):
  if raw_value is None:
    return None
  if isinstance(raw_value, bytes):
    raw_value = raw_value.decode("utf-8", errors="ignore")
  raw_value = str(raw_value).strip()
  if raw_value == "":
    return None

  watermark = safe_json_loads(raw_value)
  if watermark is not None:
    return watermark

  first_line = raw_value.splitlines()[0].strip()
  watermark = safe_json_loads(first_line)
  return watermark if watermark is not None else raw_value

def read_img_source(image_path):
  try:
    with Image.open(image_path) as image:
      image_format = (image.format or "").upper()
      if image_format == "PNG" or str(image_path).lower().endswith(".png"):
        return read_png_source(image)

    exif_dict = piexif.load(image_path)
    description = exif_dict.get("0th", {}).get(piexif.ImageIFD.ImageDescription)
    if description is None:
      return None
    return parse_hidden_watermark(description)
  except Exception as e:
    logger.error(f'read_img_source failed: {e}')
    return None

def read_png_source(image_or_path):
  if isinstance(image_or_path, Image.Image):
    image = image_or_path
    info = image.info or {}
  else:
    with Image.open(image_or_path) as image:
      info = image.info or {}

  for key in [PNG_WATERMARK_TAG, PNG_WATERMARK_TAG.lower(), "Description", "description"]:
    watermark = parse_hidden_watermark(info.get(key))
    if watermark is not None:
      return watermark
  return None

def _check_ffmpeg_command(command_name):
  command_path = shutil.which(command_name)
  if command_path is None:
    raise FileNotFoundError(f"{command_name} not found in PATH")
  return command_path

def add_media_source(source_path, target_path, meta: dict):
  watermark = meta.get("watermark")
  if watermark is None:
    raise ValueError("watermark is required for media metadata")

  ffmpeg_path = _check_ffmpeg_command("ffmpeg")
  metadata_value = json.dumps(watermark, ensure_ascii=False)
  command = [
    ffmpeg_path,
    "-y",
    "-i",
    source_path,
    "-metadata",
    f"{MEDIA_WATERMARK_TAG}={metadata_value}",
    "-codec",
    "copy",
    target_path,
  ]
  subprocess.run(command, check=True, capture_output=True, text=True)
  return target_path

def read_media_source(media_path):
  try:
    ffprobe_path = _check_ffmpeg_command("ffprobe")
    command = [
      ffprobe_path,
      "-v",
      "quiet",
      "-print_format",
      "json",
      "-show_format",
      media_path,
    ]
    result = subprocess.run(command, check=True, capture_output=True, text=True)
    output = safe_json_loads(result.stdout)
    tags = (output or {}).get("format", {}).get("tags", {})
    for key in [MEDIA_WATERMARK_TAG, MEDIA_WATERMARK_TAG.upper(), "DESCRIPTION", "description"]:
      watermark = parse_hidden_watermark(tags.get(key))
      if watermark is not None:
        return watermark
    return None
  except Exception as e:
    logger.error(f'read_media_source failed: {e}')
    return None

def _build_reserved_code(task_id):
  if task_id is None:
    return uuid.uuid4().hex
  task_id_str = str(task_id)
  try:
    return uuid.UUID(task_id_str).hex
  except Exception:
    return uuid.uuid5(uuid.NAMESPACE_URL, task_id_str).hex

def build_hidden_watermark(task_id):
  produce_id = str(task_id) if task_id is not None else uuid.uuid4().hex
  reserved_code = _build_reserved_code(task_id)
  return {
    "Label": "1",
    "ContentProducer": WATERMARK_SUBJECT,
    "ProduceID": produce_id,
    "ReservedCode1": reserved_code,
    "ContentPropagator": WATERMARK_SUBJECT,
    "PropagateID": produce_id,
    "ReservedCode2": reserved_code
  }

def _write_media_watermark(content, path, task_id):
  output_dir = "output/" + "/".join(path.split("/")[:-1])
  create_dir_if_not_exists(output_dir)
  source_path = "output/" + path
  file_name, file_ext = os.path.splitext(source_path)
  target_path = f"{file_name}.watermarked{file_ext}"
  with open(source_path, "wb") as file_obj:
    file_obj.write(content)

  watermark = build_hidden_watermark(task_id)
  add_media_source(source_path, target_path, {"watermark": watermark})
  with open(target_path, "rb") as file_obj:
    return file_obj.read()

def write_hidden_watermark_file(source_path, task_id, target_path=None, compress=True):
  source_path = os.path.abspath(source_path)
  file_extension = get_file_extension(source_path).lower()

  if target_path is None:
    file_root, file_ext = os.path.splitext(source_path)
    target_ext = ".webp" if compress and file_extension not in MEDIA_FILE_EXTENSIONS else file_ext
    target_path = f"{file_root}.watermarked{target_ext}"

  target_path = os.path.abspath(target_path)
  create_dir_if_not_exists(os.path.dirname(target_path) or ".")
  watermark = build_hidden_watermark(task_id)

  if file_extension in MEDIA_FILE_EXTENSIONS:
    add_media_source(source_path, target_path, {"watermark": watermark})
    return target_path

  with Image.open(source_path) as image:
    if is_animated_webp(image):
      shutil.copyfile(source_path, target_path)
      logger.info(f"write_hidden_watermark_file: {source_path} is animated webp, skip watermark")
      return target_path

  with tempfile.TemporaryDirectory() as temp_dir:
    temp_file_path = os.path.join(temp_dir, "source.png")
    temp_output_path = os.path.join(temp_dir, "output.webp" if compress else "output.png")

    with Image.open(source_path) as image:
      if image.mode == 'RGBA':
        try:
          image.save(temp_file_path, format='PNG', lossless=True, transparency=image.info.get('transparency'))
        except Exception as e:
          logger.error(f'write_hidden_watermark_file failed: {e}')
          image.save(temp_file_path, format='PNG')
      else:
        image.save(temp_file_path, format='PNG')

    if compress:
      compress_png(temp_file_path, temp_output_path)
    else:
      shutil.copyfile(temp_file_path, temp_output_path)

    now = datetime.now().strftime("%Y:%m:%d %H:%M:%S")
    add_img_source(temp_output_path, temp_output_path, {
      "watermark": watermark,
      "description": "",
      "artist": "",
      "software": "",
      "copyright": "",
      "DateTimeOriginal": now,
      "DateTime": now
    })
    shutil.copyfile(temp_output_path, target_path)

  return target_path

def upload_remote_image(url, params, path, file_extension='png', compress=True, task_id=None):
  content = b""
  try:
    logger.info(f"upload_remote_image: {url}")
    content = get_content(url, params=params)
    # 视频或音频不压缩
    if file_extension in MEDIA_FILE_EXTENSIONS:
      logger.info(f"upload_remote_image: {path} is media, writing metadata watermark")
      try:
        media_content = _write_media_watermark(content, path, task_id)
        return path, upload_file(BytesIO(media_content), path)
      except Exception as e:
        logger.error(f'upload_remote_image media watermark failed: {e}')
        return path, upload_file(BytesIO(content), path)
    image = Image.open(BytesIO(content))
    logger.info(f"image.mode: {image.mode}")
    # 动图不压缩
    if is_animated_webp(image):
      logger.info(f"upload_remote_image: {path} is animated webp, no compression")
      return path, upload_file(BytesIO(content), path)
    else:
      logger.info(f"upload_remote_image: {path} is image, compressing")
      dir_path = "output/" + "/".join(path.split("/")[:-1])
      create_dir_if_not_exists(dir_path)
      temp_file_path = "output/" + path
      if image.mode == 'RGBA':
        logger.info(f"upload_remote_image: {path} is RGBA, converting to webp")
        try:
          image.save(temp_file_path, format='PNG', lossless=True, transparency=image.info.get('transparency'))
        except Exception as e:
          logger.error(f'upload_remote_image failed: {e}')
          image.save(temp_file_path, format='PNG')
      else:
        image.save(temp_file_path, format='PNG')
      # image.save(temp_file_path)
      temp_webp_path = temp_file_path.split(".")[0] + ".webp"
      if compress:
        compress_png(temp_file_path, temp_webp_path)
      else:
        temp_webp_path = temp_file_path
      now = datetime.now().strftime("%Y:%m:%d %H:%M:%S")
      watermark = build_hidden_watermark(task_id)
      add_img_source(temp_webp_path, temp_webp_path, {
        "watermark": watermark,
        "description": "",
        "artist": "",
        "software": "",
        "copyright": "",
        "DateTimeOriginal": now,
        "DateTime": now
      })
      with open(temp_webp_path, "rb") as file_obj:
        new_content = file_obj.read()
      if compress:
        new_path = path.split(".")[0] + ".webp"
      else:
        new_path = path
    return new_path, upload_file(BytesIO(new_content), new_path)
  except Exception as e:
    logger.error(f'upload_remote_image failed: {e}')
    return path, upload_file(BytesIO(content), path)
  
def debounce(wait_time):
  def decorator(func):
    timer = None
    timer_lock = Lock()

    @wraps(func)
    def debounced_function(*args, **kwargs):
      nonlocal timer

      def call_it():
        nonlocal timer
        with timer_lock:
          timer = None
        return func(*args, **kwargs)

      with timer_lock:
        if timer is not None:
          timer.cancel()
        timer = Timer(wait_time, call_it)
        timer.start()

    return debounced_function

  return decorator
  
def upload_image_to_comfyui(upload_url, filename, image_url):
  image_data = get_content(image_url)

  dir_path = datetime.now().strftime("%Y-%m-%d")
  file_id = generate('1234567890abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ', 10)
  filename = f'{dir_path}-{file_id}-{filename}'

  # 准备 multipart/form-data 请求体
  files = {
      'image': (filename, image_data, 'image/png'),
      'subfolder': (None, ''),
      'type': (None, 'input')
  }

  resp = post_multipart(upload_url, files=files)

  # 检查响应
  if resp.ok:
      res = resp.json()
      logger.info(f'Image uploaded successfully: {res}')
      return res
  else:
      raise Exception(f'Failed to upload image: {resp.text}')
    
def upload_file_to_comfyui(upload_url, filename, file_url):
  file_data = get_content(file_url)

  dir_path = datetime.now().strftime("%Y-%m-%d")
  file_id = generate('1234567890abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ', 10)
  filename = f'{dir_path}-{file_id}-{filename}'

  # 准备 multipart/form-data 请求体
  files = {
      'image': (filename, file_data, 'application/octet-stream'),
      'subfolder': (None, ''),
      'type': (None, 'input')
  }

  resp = post_multipart(upload_url, files=files)

  # 检查响应
  if resp.ok:
      res = resp.json()
      logger.info(f'File uploaded successfully: {res}')
      return res
  else:
      raise Exception(f'Failed to upload file: {resp.text}')
      
def safe_json_loads(json_str):
  try:
    return json.loads(json_str)
  except Exception as e:
    logger.error(f'safe_json_loads failed: {e}')
    return None