import gradio as gr
import pandas as pd
import requests
import zipfile
import io
import re
from datetime import datetime
from pathlib import Path
import tempfile
from concurrent.futures import ThreadPoolExecutor, as_completed

TEMP_DIR = Path(tempfile.gettempdir()) / "img_downloader"
TEMP_DIR.mkdir(exist_ok=True)

def parse_published_time(time_str):
    if pd.isna(time_str):
        return datetime.min
    time_str = str(time_str).strip().lower()
    now = datetime.now()
    patterns = [
        (r'(\d+)\s*second', lambda m: now.replace(second=max(0, now.second - int(m.group(1))))),
        (r'(\d+)\s*minute', lambda m: now.replace(minute=max(0, now.minute - int(m.group(1))))),
        (r'(\d+)\s*hour',   lambda m: now.replace(hour=max(0, now.hour   - int(m.group(1))))),
        (r'(\d+)\s*day',    lambda m: now.replace(day=max(1,  now.day    - int(m.group(1))))),
        (r'(\d+)\s*week',   lambda m: now.replace(day=max(1,  now.day    - int(m.group(1))*7))),
        (r'(\d+)\s*month',  lambda m: now.replace(month=max(1,now.month  - int(m.group(1))))),
        (r'(\d+)\s*year',   lambda m: now.replace(year=max(2000, now.year- int(m.group(1))))),
    ]
    for pattern, handler in patterns:
        m = re.search(pattern, time_str)
        if m:
            try:
                return handler(m)
            except Exception:
                return datetime.min
    return datetime.min


def detect_ext_from_bytes(data):
    """Read magic bytes — the only 100% reliable way to know the true file type."""
    h = data[:16]
    # JPEG
    if h[:3] == b'\xff\xd8\xff':
        return ".jpg"
    # PNG
    if h[:8] == b'\x89PNG\r\n\x1a\n':
        return ".png"
    # GIF
    if h[:6] in (b'GIF87a', b'GIF89a'):
        return ".gif"
    # WEBP  (RIFF....WEBP)
    if h[:4] == b'RIFF' and h[8:12] == b'WEBP':
        return ".webp"
    # MP4 / MOV  (ftyp box)
    if h[4:8] == b'ftyp':
        return ".mp4"
    # WebM / MKV
    if h[:4] in (b'\x1aE\xdf\xa3', b'\x1aE\xbf\xa3'):
        return ".webm"
    # AVIF starts with ftyp avif
    if h[4:12] == b'ftypavif':
        return ".avif"
    # Default — YouTube thumbnails are almost always JPEG
    return ".jpg"


def load_csv(file_obj):
    if file_obj is None:
        return gr.update(choices=[], value=None), gr.update(visible=False), "No file uploaded.", None

    try:
        df = pd.read_csv(file_obj.name)
    except Exception as e:
        return gr.update(choices=[], value=None), gr.update(visible=False), f"❌ Error reading CSV: {e}", None

    image_cols = [c for c in df.columns if re.match(r'^images/\d+$', c)]
    other_cols = [c for c in df.columns if c not in image_cols]
    all_cols   = image_cols + other_cols

    n_rows   = len(df)
    n_img_c  = len(image_cols)
    total_possible = df[image_cols].notna().sum().sum() if image_cols else 0

    summary = (
        f"✅ **Loaded:** {n_rows} posts · "
        f"{n_img_c} image columns · "
        f"~{total_possible} image links found"
    )

    return (
        gr.update(choices=all_cols, value=image_cols, visible=True),
        gr.update(visible=True),
        summary,
        df.to_json(orient="split")
    )


def download_images(selected_cols, df_json, progress=gr.Progress()):
    if not selected_cols:
        return None, "⚠️ Please select at least one column."
    if not df_json:
        return None, "⚠️ Please upload a CSV first."

    df = pd.read_json(io.StringIO(df_json), orient="split")

    if "publishedTime" in df.columns:
        df["_sort_dt"] = df["publishedTime"].apply(parse_published_time)
        df = df.sort_values("_sort_dt", ascending=True).reset_index(drop=True)

    image_cols = [c for c in selected_cols if re.match(r'^images/\d+$', c)]
    if not image_cols:
        return None, "⚠️ No image columns selected (columns like images/0, images/1…)."

    tasks = []
    for idx, row in df.iterrows():
        date_label = str(row.get("publishedTime", f"post_{idx}")).strip().replace(" ", "_").replace("/","_").replace(":","")
        post_id    = str(row.get("postId", f"post_{idx}"))[:20]
        for col in image_cols:
            url = row.get(col)
            if pd.notna(url) and isinstance(url, str) and url.startswith("http"):
                tasks.append((date_label, post_id, col, url))

    if not tasks:
        return None, "⚠️ No valid image URLs found in the selected columns."

    def fetch(args):
        i, date_label, post_id, col, url = args
        col_safe = col.replace("/", "_")
        try:
            r = requests.get(url, timeout=15, headers={"User-Agent": "Mozilla/5.0"})
            r.raise_for_status()
            ext = detect_ext_from_bytes(r.content)   # magic bytes — never lies
            filename = f"{i+1:04d}_{date_label}_{post_id}_{col_safe}{ext}"
            return i, filename, r.content, None
        except Exception as e:
            filename = f"{i+1:04d}_{date_label}_{post_id}_{col_safe}.unknown"
            return i, filename, None, str(e)

    zip_path  = TEMP_DIR / "images_download.zip"
    results   = {}
    completed = 0

    with ThreadPoolExecutor(max_workers=16) as pool:
        futures = {pool.submit(fetch, (i,) + t): i for i, t in enumerate(tasks)}
        for future in as_completed(futures):
            i, filename, content, err = future.result()
            results[i] = (filename, content, err)
            completed += 1
            progress(completed / len(tasks), desc=f"Downloaded {completed}/{len(tasks)}")

    downloaded = 0
    failed     = 0
    log_lines  = []

    with zipfile.ZipFile(zip_path, "w", compression=zipfile.ZIP_DEFLATED) as zf:
        for i in range(len(tasks)):
            filename, content, err = results[i]
            if content:
                zf.writestr(f"images/{filename}", content)
                downloaded += 1
                log_lines.append(f"✅ {filename}")
            else:
                failed += 1
                log_lines.append(f"❌ {filename} — {err}")

        overview = "\n".join([
            "IMAGE DOWNLOAD OVERVIEW",
            "=" * 50,
            f"Generated : {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}",
            f"Total     : {len(tasks)} attempted",
            f"Downloaded: {downloaded}",
            f"Failed    : {failed}",
            f"Columns   : {', '.join(image_cols)}",
            "",
            "FILES (sorted oldest → newest post):",
            "-" * 50,
        ] + log_lines)
        zf.writestr("overview.txt", overview)

    status = f"✅ Done! **{downloaded}** images downloaded, {failed} failed. ZIP is ready."
    return str(zip_path), status


# ---------- UI ----------
with gr.Blocks(title="YouTube Post Image Downloader", theme=gr.themes.Soft()) as demo:
    df_state = gr.State(None)

    gr.Markdown(
        """
        # 🖼️ YouTube Post Image Downloader
        Upload a YouTube-posts-scraper CSV, choose which image columns to download, and get a ZIP.
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            file_input  = gr.File(label="📂 Upload CSV", file_types=[".csv"])
            status_box  = gr.Markdown("Upload a CSV to begin.")
            col_select  = gr.CheckboxGroup(label="Select columns to download", visible=False)
            dl_btn      = gr.Button("⬇️ Download Images → ZIP", variant="primary", visible=False)

        with gr.Column(scale=1):
            dl_status = gr.Markdown("")
            dl_file   = gr.File(label="📦 Download ZIP", visible=True)

    file_input.change(
        load_csv,
        inputs=[file_input],
        outputs=[col_select, dl_btn, status_box, df_state]
    )

    dl_btn.click(
        download_images,
        inputs=[col_select, df_state],
        outputs=[dl_file, dl_status]
    )

demo.launch(server_name="0.0.0.0", server_port=7860)