import os
import sqlite3
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Dict, List, Optional, Sequence, Tuple

import numpy as np
from PIL import Image
from tqdm import tqdm

ROOT_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
IMAGES_DIR = os.path.join(ROOT_DIR, "images")
STASH_DIR = os.path.join(IMAGES_DIR, "Stash")
DB_PATH = os.path.join(ROOT_DIR, "db.sqlite")
MAX_WORKERS = min(16, os.cpu_count() or 8)
EXIF_METADATA_MAX_BYTES = 512
EXIF_TYPE_ORDER = ("novelai", "sd", "comfy", "mj", "celsys", "photoshop", "stealth")
EXIF_TYPE_TO_CODE = {name: idx + 1 for idx, name in enumerate(EXIF_TYPE_ORDER)}
PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n"

def open_db(path: str) -> sqlite3.Connection:
    conn = sqlite3.connect(path)
    conn.execute(
        """
        CREATE TABLE IF NOT EXISTS pixif_cache (
            post_id TEXT PRIMARY KEY,
            url TEXT,
            exif_type INTEGER
        )
        """
    )
    conn.commit()
    ensure_db_schema(conn)
    return conn

def ensure_db_schema(conn: sqlite3.Connection) -> None:
    columns = [row[1] for row in conn.execute("PRAGMA table_info(pixif_cache)")]
    if "exif_type" not in columns:
        conn.execute("ALTER TABLE pixif_cache ADD COLUMN exif_type INTEGER")
        conn.commit()

def determine_exif_type(metadata: Optional[bytes]) -> Optional[str]:
    if metadata is None:
        return None
    if metadata == b"TitleAI generated image":
        return "novelai"
    if metadata.startswith(b"parameter"):
        return "sd"
    if b'{"' in metadata:
        return "comfy"
    if metadata.startswith(b"SoftwareCelsys"):
        return "celsys"
    return "photoshop"

def exif_type_to_code(exif_type: Optional[str]) -> Optional[int]:
    if not exif_type:
        return None
    return EXIF_TYPE_TO_CODE.get(exif_type)

def parse_png_metadata(data: bytes) -> Optional[bytes]:
    index = 8
    while index < len(data):
        if index + 8 > len(data):
            break
        chunk_len = int.from_bytes(data[index:index + 4], "big")
        chunk_type = data[index + 4:index + 8]
        index += 8
        if chunk_type == b"tEXt":
            content = data[index:index + chunk_len]
            return content.replace(b"\0", b"")
        if chunk_type == b"iTXt":
            content = data[index:index + chunk_len]
            return content.strip()
        index += chunk_len + 4
    return None

def parse_png_metadata_file(path: str) -> Optional[bytes]:
    try:
        with open(path, "rb") as handle:
            head = handle.read(EXIF_METADATA_MAX_BYTES)
            if not head.startswith(PNG_SIGNATURE):
                return None
            return parse_png_metadata(head)
    except Exception:
        return None

def byteize(alpha: np.ndarray) -> np.ndarray:
    alpha = alpha.T.reshape((-1,))
    alpha = alpha[:(alpha.shape[0] // 8) * 8]
    alpha = np.bitwise_and(alpha, 1)
    alpha = alpha.reshape((-1, 8))
    alpha = np.packbits(alpha, axis=1)
    return alpha

class LSBExtractor:
    def __init__(self, alpha: np.ndarray) -> None:
        self.data = byteize(alpha)
        self.pos = 0

    def get_next_n_bytes(self, n: int) -> bytearray:
        n_bytes = self.data[self.pos:self.pos + n]
        self.pos += n
        return bytearray(n_bytes)

    def read_32bit_integer(self) -> Optional[int]:
        bytes_list = self.get_next_n_bytes(4)
        if len(bytes_list) == 4:
            return int.from_bytes(bytes_list, byteorder="big")
        return None

def extract_stealth_metadata(image: Image.Image) -> bool:
    if "A" not in image.getbands():
        raise AssertionError("image format")
    alpha = np.array(image.getchannel("A"))
    reader = LSBExtractor(alpha)
    magic = "stealth_pngcomp"
    read_magic = reader.get_next_n_bytes(len(magic)).decode("utf-8")
    if magic != read_magic:
        raise AssertionError("magic number")
    read_len = reader.read_32bit_integer()
    if read_len is None:
        raise AssertionError("length missing")
    return True

def has_stealth_png_path(path: str) -> bool:
    try:
        with Image.open(path) as image:
            return extract_stealth_metadata(image)
    except Exception:
        return False

def detect_exif_code_from_path(path: str) -> Optional[int]:
    metadata = parse_png_metadata_file(path)
    exif_type = determine_exif_type(metadata)
    code = exif_type_to_code(exif_type)
    if code is not None:
        return code
    if has_stealth_png_path(path):
        return EXIF_TYPE_TO_CODE.get("stealth")
    return None

def fetch_pending_post_ids(conn: sqlite3.Connection) -> List[str]:
    rows = conn.execute(
        """
        SELECT post_id
        FROM pixif_cache
        WHERE exif_type IS NULL
          AND COALESCE(url, '') != ''
        """
    ).fetchall()
    return [str(row[0]) for row in rows]

def update_exif_types(conn: sqlite3.Connection, rows: Sequence[Tuple[int, str]]) -> None:
    if not rows:
        return
    conn.executemany(
        """
        UPDATE pixif_cache SET exif_type = ?
        WHERE post_id = ?
        """,
        rows,
    )

def detect_exif_codes_from_files(
    post_ids: Sequence[str],
    stash_dir: str,
    max_workers: int = MAX_WORKERS,
) -> Dict[str, Optional[int]]:
    if not post_ids:
        return {}
    results: Dict[str, Optional[int]] = {}
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        futures = {
            executor.submit(
                detect_exif_code_from_path,
                os.path.join(stash_dir, f"{post_id}.png"),
            ): post_id
            for post_id in post_ids
        }
        with tqdm(total=len(futures), unit="image", desc="Scanning exif") as pbar:
            for future in as_completed(futures):
                post_id = futures[future]
                try:
                    code = future.result()
                except Exception:
                    code = None
                results[post_id] = code
                pbar.update(1)
    return results

def main() -> int:
    os.makedirs(STASH_DIR, exist_ok=True)
    conn = open_db(DB_PATH)
    try:
        post_ids = fetch_pending_post_ids(conn)
        if not post_ids:
            print("No pending rows.")
            return 0
        existing = [post_id for post_id in post_ids if os.path.exists(os.path.join(STASH_DIR, f"{post_id}.png"))]
        if not existing:
            print("No matching images in stash.")
            return 0
        results = detect_exif_codes_from_files(existing, STASH_DIR)
        rows = [
            (exif_type, post_id)
            for post_id, exif_type in results.items()
            if exif_type is not None
        ]
        if rows:
            with conn:
                update_exif_types(conn, rows)
        print(f"Updated {len(rows)} rows.")
        return 0
    finally:
        conn.close()

if __name__ == "__main__":
    raise SystemExit(main())