# app/utils/image_importer.py

from __future__ import annotations

from datetime import datetime
from pathlib import Path
import hashlib
import json
import mimetypes
import os
import re
import shutil
import subprocess
import tempfile
import time
import uuid

from PIL import Image
import exifread
import imagehash

from app import db
from app.models import ImageRecord, Tag, ArchiveRecord


# =============================================================================
# Configuration & Constants
# =============================================================================

THUMB_SIZE = (400, 400)

# Media and archive extensions we handle
ALLOWED_MEDIA_EXTS = (
    ".jpg", ".jpeg", ".jfif", ".png", ".gif", ".bmp", ".tiff", ".webp", ".mp4", ".mov"
)
ALLOWED_ARCHIVE_EXTS = (
    ".zip", ".rar", ".cbr", ".7z",
    ".tar", ".tar.gz", ".tgz", ".tar.bz2", ".tbz2", ".tar.xz", ".txz"
)

# Regex for multi-part detection
RAR_PART_RE    = re.compile(r"\.part(\d+)\.rar$", re.IGNORECASE)
SEVENZ_PART_RE = re.compile(r"\.7z\.(\d{3})$", re.IGNORECASE)
OLD_RAR_VOL_RE = re.compile(r"\.r\d{2}$", re.IGNORECASE)  # .r00, .r01, ...

# Environment configuration (optional)
ENV_TMPDIR = "ARCHIVE_TMPDIR"          # where to extract archives (default: system tmp)
ENV_MINFREE = "ARCHIVE_MIN_FREE_GB"    # min free GB required to start extraction (float; 0 = disabled)


# =============================================================================
# Utilities: Temp space management & free-space guard
# =============================================================================

def get_tmp_base() -> Path:
    """Return the base temp directory for archive extraction."""
    return Path(os.environ.get(ENV_TMPDIR, tempfile.gettempdir())).resolve()


def cleanup_stale_tempdirs(prefix: str = "extract_") -> None:
    """Delete any leftover extraction directories from previous crashes."""
    base = get_tmp_base()
    base.mkdir(parents=True, exist_ok=True)
    for entry in base.glob(f"{prefix}*"):
        if entry.is_dir():
            try:
                shutil.rmtree(entry, ignore_errors=True)
            except Exception:
                pass


def ensure_free_space_or_raise(path: Path, min_free_gb: float) -> None:
    """Raise if there is not enough free space at path."""
    if min_free_gb <= 0:
        return
    total, used, free = shutil.disk_usage(str(path))
    free_gb = free / (1024 ** 3)
    if free_gb < min_free_gb:
        raise RuntimeError(
            f"Insufficient free space at {path}: {free_gb:.1f} GB free "
            f"(required >= {min_free_gb:.1f} GB)."
        )


# =============================================================================
# Utilities: External extractors with resilient fallback
# =============================================================================

class ExtractError(Exception):
    """Raised when archive extraction fails with all strategies."""


def _run_cmd(cmd: list[str]) -> tuple[int, str, str]:
    p = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
    return p.returncode, p.stdout, p.stderr


def extract_with_unar(archive_path: str, out_dir: str) -> None:
    os.makedirs(out_dir, exist_ok=True)
    rc, out, err = _run_cmd(["unar", "-o", out_dir, archive_path])
    if rc != 0:
        raise ExtractError(f"unar failed (rc={rc})\nSTDOUT:\n{out}\nSTDERR:\n{err}")


def extract_with_7z(archive_path: str, out_dir: str) -> None:
    os.makedirs(out_dir, exist_ok=True)
    rc, out, err = _run_cmd(["7z", "x", "-y", f"-o{out_dir}", archive_path])
    if rc != 0:
        raise ExtractError(f"7z failed (rc={rc})\nSTDOUT:\n{out}\nSTDERR:\n{err}")


def looks_password_protected_with_lsar(archive_path: str) -> bool:
    """Best-effort pre-check for 'Encrypted/Password' in lsar output."""
    try:
        rc, out, err = _run_cmd(["lsar", archive_path])
        txt = (out + "\n" + err).lower()
        return ("encrypted" in txt) or ("password" in txt)
    except Exception:
        return False


def extract_archive_resilient(archive_path: str, out_dir: str) -> None:
    """
    Prefer unar for RAR/CBR (RAR5-friendly); fallback to 7z.
    For other formats, try 7z first then unar.
    Surfaces stderr so logs show CRC/password/permissions issues.
    """
    lower = archive_path.lower()
    errors = []

    # Optional password hint to skip early
    if lower.endswith((".rar", ".cbr")) and looks_password_protected_with_lsar(archive_path):
        raise ExtractError("Archive appears password-protected; skipping (no password provided).")

    if lower.endswith((".rar", ".cbr")):
        order = (extract_with_unar, extract_with_7z)
    else:
        order = (extract_with_7z, extract_with_unar)

    for fn in order:
        try:
            fn(archive_path, out_dir)
            return
        except ExtractError as e:
            errors.append(str(e))

    raise ExtractError("Both extractors failed:\n" + "\n---\n".join(errors))


# =============================================================================
# Core: Import task & archive processing
# =============================================================================

def import_images_task(source_dir: str, dest_dir: str) -> str:
    """
    Walk the import tree, process archives first (only first volumes),
    then media files. Batches DB commits for efficiency.
    """
    # Cleanup any stale extraction dirs from previous runs/crashes
    cleanup_stale_tempdirs()

    imported: list[str] = []
    batch_size = 10
    batch_counter = 0

    # Preload existing pHashes to avoid N+1 lookups
    existing_phashes = [
        (imagehash.hex_to_hash(img.perceptual_hash), img.width, img.height)
        for img in ImageRecord.query.filter(ImageRecord.perceptual_hash != None).all()
    ]

    # Ensure thumbs base exists
    (Path(dest_dir) / "thumbs").mkdir(parents=True, exist_ok=True)

    # Optional free-space guard
    min_free = float(os.environ.get(ENV_MINFREE, "0") or "0")
    try:
        ensure_free_space_or_raise(get_tmp_base(), min_free)
    except Exception as e:
        print(f"[WARN] Free-space check failed: {e}")

    for artist_dir in os.listdir(source_dir):
        artist_path = os.path.join(source_dir, artist_dir)
        if not os.path.isdir(artist_path):
            continue

        for root, _, files in os.walk(artist_path):
            for file in files:
                full_path = os.path.join(root, file)
                lower = file.lower()

                # Archives first
                if lower.endswith(ALLOWED_ARCHIVE_EXTS):
                    if not is_first_volume(full_path):
                        print(f"[SKIP] Not first archive volume: {file}")
                        continue
                    try:
                        count = process_archive(
                            archive_path=full_path,
                            dest_dir=dest_dir,
                            artist=artist_dir,
                            existing_phashes=existing_phashes
                        )
                        imported.append(f"[archive]{file}:{count}")
                    except Exception as e:
                        print(f"[WARN] Failed processing archive {file}: {e}")
                    continue

                # Non-archive media
                if not lower.endswith(ALLOWED_MEDIA_EXTS):
                    continue

                added, phash, _rec = import_single_file(
                    src_path=full_path,
                    dest_dir=dest_dir,
                    artist=artist_dir,
                    existing_phashes=existing_phashes
                )
                if added:
                    imported.append(file)

                if phash:
                    md = extract_metadata(full_path)
                    if md.get("width") and md.get("height"):
                        existing_phashes.append((phash, md["width"], md["height"]))

                batch_counter += 1
                if batch_counter >= batch_size:
                    db.session.commit()
                    print(f"[INFO] Committed batch of {batch_size} items.")
                    batch_counter = 0
                    if "sqlite" in db.engine.url.drivername:
                        time.sleep(1)

    if batch_counter > 0:
        db.session.commit()
        print(f"[INFO] Committed final batch of {batch_counter} items.")

    summary = f"Imported {len(imported)} items."
    print(f"[INFO] {summary}")
    return summary


def process_archive(archive_path: str, dest_dir: str, artist: str, existing_phashes) -> int:
    """
    Extract archive into a temp dir, import/tag media, and record ArchiveRecord for de-dup.
    Returns number of items imported or tagged.
    """
    archive_size = os.path.getsize(archive_path)
    archive_hash = calculate_hash(archive_path)

    # De-dup the archive itself
    existing_archive = ArchiveRecord.query.filter_by(hash=archive_hash).first()
    if existing_archive:
        print(f"[SKIP] Archive already imported: {archive_path}")
        return 0

    # Create/reuse tag for this archive
    tag_name = f"archive:{artist}/{os.path.basename(archive_path)}"
    tag = Tag.query.filter_by(name=tag_name).first()
    if not tag:
        tag = Tag(name=tag_name, kind="archive")
        db.session.add(tag)
        db.session.flush()  # ensure tag.id

    # Create an extraction dir under the configured temp base
    tmpdir = tempfile.mkdtemp(prefix="extract_", dir=str(get_tmp_base()))
    imported_or_tagged = 0

    try:
        # Resilient extraction (unar → fallback 7z for RAR; 7z → fallback unar for others)
        extract_archive_resilient(archive_path, tmpdir)

        # Walk extracted tree and process media files
        for root, _, files in os.walk(tmpdir):
            for file in files:
                if not file.lower().endswith(ALLOWED_MEDIA_EXTS):
                    continue
                src_path = os.path.join(root, file)

                added, _phash, rec = import_single_file(
                    src_path=src_path,
                    dest_dir=dest_dir,
                    artist=artist,
                    existing_phashes=existing_phashes,
                    commit=False  # batch after loop
                )
                if rec is not None and tag not in rec.tags:
                    rec.tags.append(tag)
                    imported_or_tagged += 1

        # Record this archive as processed
        arch = ArchiveRecord(
            filename=archive_path,
            file_size=archive_size,
            hash=archive_hash,
            artist=artist,
            tag=tag
        )
        db.session.add(arch)
        db.session.commit()
        print(f"[INFO] Archive processed: {archive_path} items={imported_or_tagged}")
        return imported_or_tagged

    finally:
        shutil.rmtree(tmpdir, ignore_errors=True)


# =============================================================================
# Core: Single-file import path
# =============================================================================

def import_single_file(src_path: str, dest_dir: str, artist: str, existing_phashes, commit: bool = True):
    """
    Import a single media file. If an equivalent file already exists, return that
    record so callers can tag it (no duplicate import).
    Returns (added: bool, phash_or_None, record_or_None).
    """
    filename = os.path.basename(src_path)
    print(f"[INFO] Processing: {src_path}")
    file_size = os.path.getsize(src_path)

    # Check by name + size
    existing = ImageRecord.query.filter_by(filename=filename, file_size=file_size).first()
    if existing:
        print(f"[SKIP] Duplicate by name+size: {filename}")
        return (False, None, existing)

    # Check by content hash
    file_hash = calculate_hash(src_path)
    existing = ImageRecord.query.filter_by(hash=file_hash).first()
    if existing:
        print(f"[SKIP] Duplicate by hash: {filename}")
        return (False, None, existing)

    # Metadata & pHash similarity
    metadata = extract_metadata(src_path)
    if not metadata["width"] or not metadata["height"]:
        print(f"[SKIP] Missing dimension data for {filename}")
        return (False, None, None)

    phash = calculate_perceptual_hash(src_path)
    if phash and is_similar_image(phash, metadata["width"], metadata["height"], existing_phashes):
        print(f"[SKIP] {filename} is visually similar to an existing larger image.")
        return (False, phash, None)

    # Copy into /images/<artist>/<filename>
    target_dir = os.path.join(dest_dir, artist)
    os.makedirs(target_dir, exist_ok=True)
    dest_path = os.path.join(target_dir, filename)
    shutil.copy2(src_path, dest_path)

    # Thumbnail
    try:
        if dest_path.lower().endswith((".mp4", ".mov")):
            thumb_filename = f"{uuid.uuid4().hex}.jpg"
            thumb_path = os.path.join(dest_dir, "thumbs", thumb_filename)
            thumb_path = generate_video_thumbnail(dest_path, thumb_path)
        else:
            thumb_path = generate_thumbnail(dest_path)
    except Exception as e:
        print(f"[WARN] Failed to generate thumbnail for {filename}: {e}")
        thumb_path = None

    # Create DB record
    record = ImageRecord(
        filename=filename,
        filepath=dest_path,
        thumb_path=thumb_path,
        hash=file_hash,
        perceptual_hash=str(phash) if phash else None,
        artist=artist,
        file_size=metadata["file_size"],
        width=metadata["width"],
        height=metadata["height"],
        format=metadata["format"],
        camera_model=metadata["camera_model"],
        taken_at=metadata["taken_at"],
        imported_at=datetime.utcnow()
    )
    db.session.add(record)
    if commit:
        db.session.commit()
    return (True, phash, record)


# =============================================================================
# Helpers: Thumbnails, metadata, hashing, similarity
# =============================================================================

def generate_thumbnail(image_path: str, size: tuple[int, int] = (400, 400), overwrite: bool = False) -> str:
    """
    Save thumbnails mirrored under /images/thumbs/<artist>/<filename>.
    """
    images_root = Path("/images").resolve()
    image_path = Path(image_path).resolve()
    rel_path = image_path.relative_to(images_root)  # raises if not under /images

    thumb_path = images_root / "thumbs" / rel_path
    thumb_path.parent.mkdir(parents=True, exist_ok=True)

    if not overwrite and thumb_path.exists():
        return str(thumb_path)

    with Image.open(image_path) as img:
        img.thumbnail(size)
        img.save(thumb_path)

    return str(thumb_path)


def generate_video_thumbnail(video_path: str, output_path: str, time_position: str = "00:00:01") -> str | None:
    os.makedirs(os.path.dirname(output_path), exist_ok=True)
    command = [
        "ffmpeg", "-i", video_path, "-ss", time_position,
        "-vframes", "1", "-vf", f"scale={THUMB_SIZE[0]}:-1", output_path
    ]
    try:
        subprocess.run(command, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        return output_path
    except subprocess.CalledProcessError as e:
        print(f"[WARN] Failed to extract video thumbnail for {video_path}: {e}")
        return None


def calculate_hash(file_path: str) -> str:
    hash_func = hashlib.sha256()
    with open(file_path, "rb") as f:
        for chunk in iter(lambda: f.read(1024 * 1024), b""):  # 1MB chunks work well for big files
            hash_func.update(chunk)
    return hash_func.hexdigest()


def calculate_perceptual_hash(file_path: str):
    try:
        with Image.open(file_path) as img:
            return imagehash.phash(img)
    except Exception:
        # videos or unreadable images land here (phash not applicable)
        return None


def is_similar_image(phash, width: int, height: int, existing_phashes, threshold: int = 2) -> bool:
    for existing, ew, eh in existing_phashes:
        try:
            distance = phash - existing
            if distance <= threshold and ew >= width and eh >= height:
                return True
        except Exception as e:
            print(f"[WARN] Failed pHash comparison: {e}")
    return False


def extract_metadata(file_path: str) -> dict:
    metadata = {
        "file_size": None,
        "width": None,
        "height": None,
        "format": None,
        "camera_model": None,
        "taken_at": None
    }

    metadata["file_size"] = os.path.getsize(file_path)
    ext = os.path.splitext(file_path)[1].lower()

    # Video: use ffprobe
    if ext in (".mp4", ".mov"):
        try:
            cmd = [
                "ffprobe", "-v", "error",
                "-select_streams", "v:0",
                "-show_entries", "stream=width,height",
                "-of", "json", file_path
            ]
            result = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, text=True, check=False)
            data = json.loads(result.stdout or "{}")
            stream = data.get("streams", [{}])[0]
            metadata["width"] = stream.get("width")
            metadata["height"] = stream.get("height")
            metadata["format"] = ext.lstrip(".")
        except Exception as e:
            print(f"[WARN] Failed to extract video metadata: {file_path} – {e}")
        return metadata

    # Image: PIL + EXIF (best-effort)
    try:
        with Image.open(file_path) as img:
            metadata["width"], metadata["height"] = img.size
            metadata["format"] = (img.format or "").lower()
    except Exception as e:
        mime_type, _ = mimetypes.guess_type(file_path)
        print(f"[WARN] Failed to read image metadata: {file_path} ({e}) ext={ext} mime={mime_type}")

    try:
        with open(file_path, "rb") as f:
            tags = exifread.process_file(f, stop_tag="UNDEF", details=False)
        if "EXIF DateTimeOriginal" in tags:
            metadata["taken_at"] = datetime.strptime(str(tags["EXIF DateTimeOriginal"]), "%Y:%m:%d %H:%M:%S")
        if "Image Model" in tags:
            metadata["camera_model"] = str(tags["Image Model"])
    except Exception:
        pass

    return metadata


# =============================================================================
# Helpers: Multi-part detection
# =============================================================================

def is_first_volume(path: str) -> bool:
    """
    Determine if 'path' is the first volume of a multi-part archive.
    We only attempt extraction from the first volume.
    """
    name = os.path.basename(path).lower()

    # RAR: .partN.rar (only N==1 is first)
    m = RAR_PART_RE.search(name)
    if m:
        n = m.group(1)
        return n in ("1", "01", "001")

    # RAR old-style: first is .rar, subsequent are .r00, .r01...
    if name.endswith(".rar"):
        return True
    if OLD_RAR_VOL_RE.search(name):
        return False

    # CBR is typically a single-volume RAR
    if name.endswith(".cbr"):
        return True

    # 7z multi-part: .7z.001 is first
    m = SEVENZ_PART_RE.search(name)
    if m:
        n = m.group(1)
        return n in ("1", "01", "001")
    if name.endswith(".7z"):
        return True  # single-volume 7z

    # ZIP and tarballs: treat as single-volume (we don't include .z01 in our allowlist)
    if name.endswith((".zip", ".tar", ".tar.gz", ".tgz", ".tar.bz2", ".tbz2", ".tar.xz", ".txz")):
        return True

    return True