# /bin/env python
# -*- coding: utf-8 -*-
"""
@author: inklov3
@date: 2025-07-16
@description: molar function tools
"""

import datetime
import hashlib
import json
import sys
import uuid
from pathlib import Path, PurePath
from typing import Iterable, TextIO
from urllib.parse import unquote, urlparse

import numpy as np
from blake3 import blake3
from loguru import logger
from loguru._logger import Logger as LoguruLogger


def file_name(file: Path, suffix: str | None = None) -> str:
    """
    文件名处理
    """
    name = file.name

    if suffix is not None:
        name = name.rsplit(".", 1)[0] + suffix

    return name


def ensure_dir(path: str | Path) -> Path:
    """
    确保目录存在
    """
    if isinstance(path, str):
        path = Path(path)
    path.mkdir(parents=True, exist_ok=True)
    return path


def ensure_file(file: Path | str) -> Path:
    """
    确保父目录存在
    """
    if isinstance(file, str):
        file = Path(file)
    file.parent.mkdir(parents=True, exist_ok=True)
    return file


def init_logger(level: str = "WARNING", logs_path: Path = Path("logs")) -> LoguruLogger:
    """
    初始化日志
    """
    if getattr(logger, "_configured", False):
        return logger

    logger.remove()
    logger.add(
        sink=sys.stdout,
        format="<green>{time:YYYY-MM-DD HH:mm:ss}</green> | <level>{level: <7}</level> | "
        "<cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan> - "
        "<level>{message}</level>",
        level=level,
    )
    if logs_path:
        logs_path.mkdir(parents=True, exist_ok=True)
        log_path = logs_path / f"{datetime.datetime.now():%Y-%m-%d_%H-%M-%S}.log"
        logger.add(
            str(log_path),
            format="<green>{time:YYYY-MM-DD HH:mm:ss}</green> | <level>{level: <7}</level> | "
            "<cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan> - "
            "<level>{message}</level>",
            level=level,
            encoding="utf-8",
        )
    setattr(logger, "_configured", True)
    return logger


def read_jsonl(file: Path) -> Iterable[dict]:
    """
    读jsonl
    """
    with open(file, "r", encoding="utf8") as file:
        for line in file:
            yield json.loads(line.strip())


def read_json(file: Path) -> dict:
    """
    读json
    """
    return json.loads(file.read_text())


class NonSerializable:
    """标记字段为不可序列化：序列化时将跳过此字段"""

    def __repr__(self):
        return "<NonSerializable>"


NON_SERIALIZABLE = NonSerializable()


def serialize(obj):
    """
    序列化
    """
    result = None
    # 基础类型或 None
    if isinstance(obj, (str, int, float, bool)) or obj is None:
        result = obj
    elif isinstance(obj, (np.generic,)):
        result = obj.item()
    # NumPy 数组 -> Python 原生嵌套 list
    elif isinstance(obj, np.ndarray):
        result = obj.tolist()
    # 字典
    elif isinstance(obj, dict):
        result = {k: serialize(v) for k, v in obj.items() if v is not NON_SERIALIZABLE}
    # 列表/元组/集合
    elif isinstance(obj, (list, tuple, set)):
        result = [serialize(v) for v in obj if v is not NON_SERIALIZABLE]
    # 普通对象
    elif hasattr(obj, "__dict__"):
        filtered = {
            k: v for k, v in vars(obj).items() if not isinstance(v, NonSerializable)
        }
        result = serialize(filtered)
    else:
        raise TypeError(f"Unsupported serialization type: {type(obj)}")
    return result


def write_json(data: any, file: Path, indent=2):
    """
    写json
    """
    with open(file, "w", encoding="utf8") as file:
        json.dump(data, file, ensure_ascii=False, indent=indent, default=serialize)


def open_jsonl_writer(file: Path) -> TextIO:
    """
    打开jsonl写入器
    """
    return open(file, "a", encoding="utf8")


def send_jsonl(jsonl_writer: TextIO, data: dict, indent=2, default=serialize):
    line = json.dumps(data, ensure_ascii=False, indent=indent, default=default) + "\n"
    """
    发送jsonl
    """
    jsonl_writer.write(line)
    jsonl_writer.flush()


def hash_string(s: str, method="md5") -> str:
    """
    字符串hash
    """
    h = hashlib.new(method)
    h.update(s.encode("utf-8"))
    return h.hexdigest()


def partial_hash(file: Path, chunk_size=4 * 1024 * 1024) -> str:
    """
    分块hash
    """
    size = file.stat().st_size
    md5 = hashlib.md5()

    ranges = []
    ranges.append((0, chunk_size))  # 开头

    if size > chunk_size * 2:
        middle = max((size // 2) - (chunk_size // 2), chunk_size)
        ranges.append((middle, chunk_size))  # 中间

    if size > chunk_size:
        ranges.append((size - chunk_size, chunk_size))  # 结尾

    with file.open("rb") as f:
        for offset, length in ranges:
            f.seek(offset)
            md5.update(f.read(length))

    return md5.hexdigest()


def get_uuid():
    """
    获取uuid
    """
    random_uuid = uuid.uuid4()
    return str(random_uuid)


def get_timestamp():
    """
    获取时间戳
    """
    time = datetime.datetime.now()
    return time.strftime("%Y.%m.%d-%H:%M:%S.%f")


def search_file_with_suffix(path: Path, suffixes: list[str]) -> list[Path]:
    """递归查找目录后缀文件"""
    return [f for f in path.rglob(r"*") if f.suffix in suffixes]


def cut_url(
    url: str, start_level: int = 2, end_level: int = -1, with_file_name: bool = False
) -> PurePath:
    """切割url"""
    url_parse = urlparse(url)
    url_path = url_parse.path
    unquoted_url_path = Path(unquote(url_path))
    if with_file_name:
        end_level = len(unquoted_url_path.parts)
    return PurePath(*unquoted_url_path.parts[start_level:end_level])


def recursive_replace(s: str, deplace_map: dict[str, str]) -> str:
    """假递归替换"""
    for k, v in deplace_map.items():
        s = s.replace(k, v)
    return s


def recursive_clean(s: str, clean_words: Iterable[str]) -> str:
    """假递归清理"""
    return recursive_replace(s, {word: "" for word in clean_words})


def hash(data: bytes, scheme: str = "blake3") -> str:
    match scheme:
        case "blake3":
            return blake3(data).hexdigest()
        case "md5":
            return hashlib.md5(data).hexdigest()
        case _:
            raise NotImplementedError(f"Unsupported hash scheme: {scheme}")
