"""
从 192.168.1.100:4500 API 拉取指定数量相册作为测试数据（多线程并行版）
导入到 Rust imager 项目的 SQLite 数据库和 albums 目录
"""
import requests
import sqlite3
import os
import sys
import time
import uuid
import threading
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime

API_BASE = "http://192.168.1.100:4500"
DB_PATH = os.path.join(os.path.dirname(os.path.dirname(__file__)), "data", "imager.db")
ALBUMS_DIR = os.path.join(os.path.dirname(os.path.dirname(__file__)), "albums")

# 全局 DB 写锁
db_lock = threading.Lock()
# 进度计数器
progress_lock = threading.Lock()
completed_count = [0]
total_count = [0]

# 线程本地 Session（复用连接）
thread_local = threading.local()


def get_session():
    if not hasattr(thread_local, "session"):
        thread_local.session = requests.Session()
        thread_local.session.timeout = 60
    return thread_local.session


def fetch_albums(limit=100):
    """获取相册列表"""
    url = f"{API_BASE}/api/albums"
    print(f"[1/4] 获取相册列表: {url}")
    resp = requests.get(url, timeout=30)
    data = resp.json()

    if isinstance(data, dict) and "albums" in data:
        albums = data["albums"]
    elif isinstance(data, list):
        albums = data
    else:
        print(f"未知响应格式: {type(data)}")
        return []

    albums = albums[:limit]
    print(f"  获取到 {len(albums)} 个相册")
    return albums


def fetch_album_images(album_id):
    """获取相册的图片列表"""
    url = f"{API_BASE}/api/album/{album_id}"
    try:
        resp = requests.get(url, timeout=30)
        data = resp.json()
        if isinstance(data, dict) and "images" in data:
            return data["images"]
        elif isinstance(data, list):
            return data
    except Exception as e:
        print(f"  获取图片列表失败 {album_id}: {e}")
    return []


def download_single_image(img):
    """下载单张图片（线程安全）"""
    img_id = img.get("id", str(uuid.uuid4()))
    mime = img.get("mime_type", "image/jpeg")
    ext = "jpg"
    if mime and "/" in mime:
        ext = mime.split("/")[-1]

    album_id = str(img.get("_album_id", ""))
    filename = f"{img_id}.{ext}"
    db_filename = f"{album_id}/{filename}"  # 数据库路径：{album_id}/{img_id}.ext
    save_path = os.path.join(ALBUMS_DIR, album_id, filename)

    # 跳过已下载
    if os.path.exists(save_path) and os.path.getsize(save_path) > 0:
        file_size = os.path.getsize(save_path)
        width, height = get_image_dimensions(save_path)
        return (img_id, db_filename, file_size, width, height, mime)

    url = f"{API_BASE}/api/image/{img_id}"
    try:
        session = get_session()
        resp = session.get(url, timeout=60)
        if resp.status_code == 200:
            os.makedirs(os.path.dirname(save_path), exist_ok=True)
            with open(save_path, "wb") as f:
                f.write(resp.content)
            file_size = len(resp.content)
            width, height = get_image_dimensions(save_path)
            return (img_id, db_filename, file_size, width, height, resp.headers.get("content-type", mime))
    except Exception as e:
        pass
    return None


def get_image_dimensions(path):
    """获取图片尺寸（二进制头解析）"""
    try:
        with open(path, "rb") as f:
            header = f.read(32)
            if header[:2] == b"\xff\xd8":
                f.seek(2)
                while True:
                    chunk = f.read(4)
                    if len(chunk) < 4:
                        break
                    if chunk[0] != 0xFF:
                        break
                    seg_type = chunk[1]
                    if seg_type == 0xDA:
                        break
                    seg_len = int.from_bytes(chunk[2:4], "big")
                    if seg_type in (0xC0, 0xC1, 0xC2):
                        data = f.read(seg_len - 2)
                        if len(data) >= 5:
                            return int.from_bytes(data[1:3], "big"), int.from_bytes(data[3:5], "big")
                    f.seek(seg_len - 2, 1)
                return 0, 0
            elif header[:8] == b"\x89PNG\r\n\x1a\n":
                return int.from_bytes(header[16:20], "big"), int.from_bytes(header[20:24], "big")
            elif header[:4] == b"RIFF" and header[8:12] == b"WEBP":
                return int.from_bytes(header[24:26], "little"), int.from_bytes(header[26:28], "little")
            elif header[:6] in (b"GIF87a", b"GIF89a"):
                return int.from_bytes(header[6:8], "little"), int.from_bytes(header[8:10], "little")
            return 0, 0
    except Exception:
        return 0, 0


def sync_album_to_db(conn, album, images_info):
    """将相册和图片写入数据库（调用方持有 db_lock）"""
    cursor = conn.cursor()
    now = datetime.utcnow().isoformat() + "Z"
    album_id = album["id"]
    title = album.get("title", "Untitled")

    cursor.execute(
        """INSERT OR REPLACE INTO albums
           (id, title, document, description, cover_image_id, image_count, created_at, updated_at)
           VALUES (?, ?, ?, ?, ?, ?, ?, ?)""",
        (
            album_id, title, album_id,
            album.get("description"),
            images_info[0][0] if images_info else None,
            len(images_info),
            album.get("created_at", now),
            now,
        ),
    )

    for img_id, filename, file_size, width, height, mime_type in images_info:
        cursor.execute(
            """INSERT OR REPLACE INTO images
               (id, album_id, filename, original_url, title, description, file_size, width, height, mime_type, created_at)
               VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
            (img_id, album_id, filename, f"{API_BASE}/api/image/{img_id}",
             f"图片 {img_id[:8]}", None, file_size, width, height, mime_type or "image/jpeg", now),
        )

    conn.commit()


def process_album(album, conn, dl_workers):
    """处理单个相册：下载所有图片并写入数据库"""
    aid = album["id"]
    title = album.get("title", "Untitled")[:60]

    # 获取图片列表
    images = fetch_album_images(aid)
    if not images:
        with progress_lock:
            completed_count[0] += 1
            print(f"[{completed_count[0]}/{total_count[0]}] {title} — 无图片，跳过")
        return 0

    # 为每个图片打上 album_id 标记
    for img in images:
        img["_album_id"] = aid

    # 并行下载
    downloaded = []
    with ThreadPoolExecutor(max_workers=dl_workers) as pool:
        futures = {pool.submit(download_single_image, img): img for img in images}
        for future in as_completed(futures):
            result = future.result()
            if result:
                downloaded.append(result)

    if not downloaded:
        with progress_lock:
            completed_count[0] += 1
            print(f"[{completed_count[0]}/{total_count[0]}] {title} — 0 张下载成功")
        return 0

    # 写入数据库
    with db_lock:
        sync_album_to_db(conn, album, downloaded)

    n = len(downloaded)
    with progress_lock:
        completed_count[0] += 1
        print(f"[{completed_count[0]}/{total_count[0]}] {title} — {n} 张图片")

    return n


def main():
    total = int(sys.argv[1]) if len(sys.argv) > 1 else 100
    dl_workers = int(sys.argv[2]) if len(sys.argv) > 2 else 8
    album_workers = int(sys.argv[3]) if len(sys.argv) > 3 else 4

    print(f"=== 从 API 导入 {total} 个相册（并行下载: {dl_workers}线程/相册, 并行相册: {album_workers}）===\n")

    os.makedirs(os.path.dirname(DB_PATH), exist_ok=True)
    os.makedirs(ALBUMS_DIR, exist_ok=True)

    conn = sqlite3.connect(DB_PATH, check_same_thread=False)
    conn.execute("PRAGMA journal_mode=WAL")
    conn.execute("PRAGMA foreign_keys=ON")
    conn.execute("PRAGMA synchronous=NORMAL")

    # 获取已有相册
    existing = set()
    for row in conn.execute("SELECT id FROM albums"):
        existing.add(row[0])
    print(f"  数据库中已有 {len(existing)} 个相册")

    # 获取相册列表
    albums = fetch_albums(total)
    if not albums:
        print("未获取到相册，退出")
        return

    new_albums = [a for a in albums if a["id"] not in existing]
    skipped = len(albums) - len(new_albums)
    if skipped > 0:
        print(f"  跳过 {skipped} 个已存在的相册")
    print(f"  需要导入 {len(new_albums)} 个新相册\n")

    if not new_albums:
        print("没有新相册需要导入")
        conn.close()
        return

    total_count[0] = len(new_albums)
    total_images = 0
    t0 = time.time()

    # 并行处理相册
    with ThreadPoolExecutor(max_workers=album_workers) as pool:
        futures = {pool.submit(process_album, album, conn, dl_workers): album for album in new_albums}
        for future in as_completed(futures):
            n = future.result()
            total_images += n

    elapsed = time.time() - t0

    # 重建 image_count
    print(f"\n[4/4] 更新相册图片计数...")
    conn.execute("""UPDATE albums SET image_count = (SELECT COUNT(*) FROM images WHERE images.album_id = albums.id)""")
    conn.commit()
    conn.close()

    print(f"\n=== 导入完成 ===")
    print(f"  成功导入: {len(new_albums)} 个相册")
    print(f"  总图片数: {total_images} 张")
    print(f"  耗时: {elapsed:.1f} 秒")
    print(f"  数据库: {DB_PATH}")
    print(f"  图片目录: {ALBUMS_DIR}")


if __name__ == "__main__":
    main()