# image_import_worker.py
import os
import time
import logging
from dataclasses import dataclass
from typing import Callable, Optional, Tuple

from sqlalchemy.exc import OperationalError

from app import create_app
from app.utils.image_importer import import_images_task

# -------------------------------------------------------------------
# Paths & basic config
# -------------------------------------------------------------------
THUMBNAIL_FLAG = "/import/thumbnail.flag"
TRIGGER_FLAG   = "/import/trigger.flag"
SOURCE_DIR     = "/import"
DEST_DIR       = "/images"

CHECK_INTERVAL = int(os.getenv("CHECK_INTERVAL", "10"))  # seconds

# Periodic import interval (default 8 hours)
IMPORT_EVERY_SECONDS = int(os.getenv("IMPORT_EVERY_SECONDS", str(8 * 60 * 60)))

# Thumbnail verbosity controls
THUMBS_VERBOSE     = os.getenv("THUMBS_VERBOSE", "0").lower() in ("1", "true", "yes")
THUMBS_LOG_EVERY   = int(os.getenv("THUMBS_LOG_EVERY", "50"))
THUMBS_BATCH_SIZE  = int(os.getenv("THUMBS_BATCH_SIZE", "500"))

MAX_RETRIES = int(os.getenv("MAX_RETRIES", "3"))

# -------------------------------------------------------------------
# Logging
# -------------------------------------------------------------------
logging.basicConfig(
    level=os.getenv("LOG_LEVEL", "INFO").upper(),
    format="%(asctime)s %(levelname)s %(message)s",
)
log = logging.getLogger("image-import-worker")

app = create_app()


# -------------------------------------------------------------------
# Job / flag helpers
# -------------------------------------------------------------------
@dataclass(frozen=True)
class FlagJob:
    name: str
    flag_path: str

    @property
    def run_path(self) -> str:
        return self.flag_path + ".run"

    @property
    def err_path(self) -> str:
        return self.flag_path + ".err"


def claim_job(job: FlagJob) -> bool:
    """
    Atomically claim a job by renaming <flag> -> <flag>.run.
    Returns True if claimed, False otherwise.
    """
    try:
        os.replace(job.flag_path, job.run_path)
        log.info("[%s] Claimed job -> %s", job.name, job.run_path)
        return True
    except FileNotFoundError:
        # Flag disappeared between exists() check and replace()
        return False
    except Exception as e:
        log.warning("[%s] Could not claim job: %s", job.name, e)
        return False


def complete_job(job: FlagJob) -> None:
    """Remove the .run file on success."""
    try:
        os.remove(job.run_path)
        log.info("[%s] Cleared run flag.", job.name)
    except FileNotFoundError:
        log.info("[%s] Run flag already cleared.", job.name)
    except Exception as e:
        log.warning("[%s] Could not remove run flag: %s", job.name, e)


def fail_job(job: FlagJob, error_text: Optional[str] = None) -> None:
    """
    Move .run -> .err and write error details into .err.
    """
    import datetime

    try:
        os.replace(job.run_path, job.err_path)
        log.error("[%s] Marked failure -> %s", job.name, job.err_path)

        if error_text:
            try:
                with open(job.err_path, "a", encoding="utf-8") as f:
                    f.write("\n")
                    f.write(datetime.datetime.utcnow().isoformat() + "Z\n")
                    f.write(error_text)
                    f.write("\n")
            except Exception as e:
                log.warning("[%s] Could not write error details: %s", job.name, e)

    except FileNotFoundError:
        log.error("[%s] Could not mark failure: run flag missing.", job.name)
    except Exception as e:
        log.error("[%s] Could not mark failure: %s", job.name, e)


# -------------------------------------------------------------------
# DB helpers
# -------------------------------------------------------------------
def db_ping() -> None:
    """Validate DB session connection quickly."""
    from sqlalchemy import text
    from app import db

    db.session.execute(text("SELECT 1"))


def recover_db_session(dispose_engine: bool = False) -> None:
    """
    Rollback session and optionally dispose engine pool to force fresh connections.
    """
    try:
        from app import db
        db.session.rollback()
        if dispose_engine:
            db.engine.dispose()
    except Exception:
        pass


def run_with_retries(
    job: FlagJob,
    fn: Callable[[], None],
    *,
    dispose_engine_on_operational_error: bool = False,
    max_retries: int = MAX_RETRIES,
) -> Tuple[bool, Optional[str]]:
    """
    Run fn() with retry/backoff for OperationalError.
    Returns (success, error_text). error_text contains the last traceback on failure.
    """
    import traceback

    attempt = 0
    last_error_text: Optional[str] = None

    while attempt < max_retries:
        try:
            # Catch stale pool conns early
            db_ping()

            fn()
            return True, None

        except OperationalError as e:
            attempt += 1
            last_error_text = "".join(traceback.format_exception(type(e), e, e.__traceback__))
            log.warning(
                "[%s] DB error (attempt %d/%d): %s",
                job.name,
                attempt,
                max_retries,
                e,
            )

            recover_db_session(dispose_engine=dispose_engine_on_operational_error)

            if attempt >= max_retries:
                return False, last_error_text

            sleep_time = 2 ** attempt
            log.info("[%s] Retrying in %ss...", job.name, sleep_time)
            time.sleep(sleep_time)

        except Exception as e:
            last_error_text = "".join(traceback.format_exception(type(e), e, e.__traceback__))
            log.exception("[%s] Unexpected error: %s", job.name, e)
            recover_db_session(dispose_engine=True)
            return False, last_error_text

    return False, last_error_text


# -------------------------------------------------------------------
# Actual job bodies
# -------------------------------------------------------------------
def do_import() -> None:
    result = import_images_task(SOURCE_DIR, DEST_DIR)
    log.info("[Import] %s", result)


def do_thumbnail_regen() -> None:
    """
    Your thumbnail regeneration logic, moved into a function.
    Behavior is unchanged: chunked, progress logging, commit per batch.
    """
    from pathlib import Path
    from app import db
    from app.models import ImageRecord
    from app.utils.image_importer import generate_thumbnail, generate_video_thumbnail

    total = db.session.query(ImageRecord).count()
    log.info("[Thumbs] Starting regen; ImageRecord count = %d", total)

    images_root = Path(DEST_DIR).resolve()
    batch_size = THUMBS_BATCH_SIZE

    processed = 0
    writes = 0
    path_updates = 0
    failures = 0
    last_id = 0
    t_start = time.time()

    while True:
        rows = (
            ImageRecord.query
            .filter(ImageRecord.id > last_id)
            .order_by(ImageRecord.id.asc())
            .limit(batch_size)
            .all()
        )
        if not rows:
            break

        batch_start_id = rows[0].id
        batch_end_id   = rows[-1].id
        batch_processed = 0
        batch_writes = 0
        batch_updates = 0
        batch_failures = 0
        bt0 = time.time()

        for img in rows:
            try:
                processed += 1
                batch_processed += 1

                ext = (img.format or "").lower()
                is_video = ext in ("mp4", "mov") or img.filepath.lower().endswith((".mp4", ".mov"))

                if is_video:
                    out_path = img.thumb_path
                    if not out_path:
                        try:
                            rel = Path(img.filepath).resolve().relative_to(images_root)
                        except Exception:
                            batch_failures += 1
                            failures += 1
                            if THUMBS_VERBOSE:
                                log.info(
                                    "[Thumbs] SKIP (not under %s): id=%s %s",
                                    images_root,
                                    img.id,
                                    img.filepath,
                                )
                            continue
                        out_path = str((images_root / "thumbs" / rel).with_suffix(".jpg"))

                    new_path = generate_video_thumbnail(img.filepath, out_path)
                else:
                    new_path = generate_thumbnail(img.filepath, overwrite=True)

                if new_path:
                    writes += 1
                    batch_writes += 1

                if new_path and new_path != img.thumb_path:
                    img.thumb_path = new_path
                    path_updates += 1
                    batch_updates += 1

                if THUMBS_VERBOSE and (processed <= 10 or processed % THUMBS_LOG_EVERY == 0):
                    kind = "video" if is_video else "image"
                    log.info(
                        "[Thumbs] #%d (%s) id=%s name='%s' -> %s",
                        processed,
                        kind,
                        img.id,
                        img.filename,
                        new_path or "None",
                    )

            except Exception as e:
                failures += 1
                batch_failures += 1
                log.warning(
                    "[Thumbs] Failed for id=%s name='%s': %s",
                    img.id,
                    img.filename,
                    e,
                )

        db.session.commit()
        last_id = rows[-1].id
        db.session.expunge_all()

        bt = time.time() - bt0
        rate = batch_processed / bt if bt > 0 else 0
        log.info(
            "[Thumbs] Batch %s-%s: processed=%d writes=%d db_updates=%d failures=%d (%.1f items/s) "
            "— totals: processed=%d writes=%d db_updates=%d failures=%d (last_id=%d)",
            batch_start_id,
            batch_end_id,
            batch_processed,
            batch_writes,
            batch_updates,
            batch_failures,
            rate,
            processed,
            writes,
            path_updates,
            failures,
            last_id,
        )

    elapsed = time.time() - t_start
    overall_rate = processed / elapsed if elapsed > 0 else 0
    log.info(
        "[Thumbs] COMPLETE in %.1fs — processed=%d writes=%d db_updates=%d failures=%d (%.1f items/s).",
        elapsed,
        processed,
        writes,
        path_updates,
        failures,
        overall_rate,
    )


# -------------------------------------------------------------------
# Orchestration helpers
# -------------------------------------------------------------------
def handle_flag_job(
    job: FlagJob,
    fn: Callable[[], None],
    *,
    dispose_engine_on_operational_error: bool = False,
) -> bool:
    """
    If <flag> exists, claim it and run fn with retries.
    Returns True if it did anything (including another worker claiming it).
    """
    if not os.path.exists(job.flag_path):
        return False

    log.info("[%s] Flag detected: %s", job.name, job.flag_path)

    if not claim_job(job):
        # It existed, but someone else may have claimed it.
        return True

    ok, err_text = run_with_retries(
        job,
        fn,
        dispose_engine_on_operational_error=dispose_engine_on_operational_error,
    )
    if ok:
        complete_job(job)
    else:
        fail_job(job, error_text=err_text)

    return True


# -------------------------------------------------------------------
# Main loop
# -------------------------------------------------------------------
def monitor_and_import() -> None:
    thumbnail_job = FlagJob(name="Thumbs", flag_path=THUMBNAIL_FLAG)
    import_job    = FlagJob(name="Import", flag_path=TRIGGER_FLAG)

    # next time we should run periodic import (monotonic avoids wall-clock jumps)
    next_periodic_import = time.monotonic() + IMPORT_EVERY_SECONDS

    with app.app_context():
        log.info(
            "[Worker] Started. Poll=%ss, periodic import every=%ss (%.2fh).",
            CHECK_INTERVAL,
            IMPORT_EVERY_SECONDS,
            IMPORT_EVERY_SECONDS / 3600.0,
        )

        while True:
            did_work = False

            # 1) Thumbnail regen (priority)
            did_work |= handle_flag_job(thumbnail_job, do_thumbnail_regen)

            # 2) Flag-triggered import
            did_work |= handle_flag_job(
                import_job,
                do_import,
                dispose_engine_on_operational_error=True,
            )

            # 3) Periodic import (every IMPORT_EVERY_SECONDS)
            now = time.monotonic()
            if IMPORT_EVERY_SECONDS > 0 and now >= next_periodic_import:
                next_periodic_import = now + IMPORT_EVERY_SECONDS

                # Skip if an import is already in progress or in error
                if os.path.exists(import_job.run_path):
                    log.info("[Import] Periodic run skipped (import already running).")
                elif os.path.exists(import_job.err_path):
                    log.info(
                        "[Import] Periodic run skipped (import in error state: %s).",
                        import_job.err_path,
                    )
                else:
                    # Create a .run lock for periodic run (no trigger.flag required)
                    tmp = import_job.run_path + ".tmp"
                    try:
                        with open(tmp, "w", encoding="utf-8") as f:
                            f.write("periodic\n")
                        os.replace(tmp, import_job.run_path)
                        log.info("[Import] Periodic run claimed -> %s", import_job.run_path)

                        ok, err_text = run_with_retries(
                            import_job,
                            do_import,
                            dispose_engine_on_operational_error=True,
                        )
                        if ok:
                            complete_job(import_job)
                        else:
                            fail_job(import_job, error_text=err_text)

                        did_work = True
                    except Exception as e:
                        log.warning("[Import] Could not start periodic run: %s", e)
                        try:
                            if os.path.exists(tmp):
                                os.remove(tmp)
                        except Exception:
                            pass

            if not did_work:
                time.sleep(CHECK_INTERVAL)


if __name__ == "__main__":
    monitor_and_import()