import os
import io
import zipfile
import threading
import uuid
import json
from concurrent.futures import ThreadPoolExecutor, as_completed
from functools import wraps
from flask import (
    Flask, request, render_template, redirect,
    url_for, session, jsonify, send_file, flash
)
import openpyxl
import requests
from urllib.parse import urlparse
from werkzeug.utils import secure_filename

app = Flask(__name__)
app.secret_key = os.environ.get("SECRET_KEY", "change-this-in-production-please")

# ── CONFIG ────────────────────────────────────────────────────────────────────
UPLOAD_FOLDER    = "uploads"
HYPERLINK_COLUMN = 11
MAX_CONTENT_LENGTH = 50 * 1024 * 1024
app.config["UPLOAD_FOLDER"]      = UPLOAD_FOLDER
app.config["MAX_CONTENT_LENGTH"] = MAX_CONTENT_LENGTH

USERS = {
    "admin": "SecurePass123!",
    "user1": "Password456!",
}

jobs      = {}
jobs_lock = threading.Lock()

os.makedirs(UPLOAD_FOLDER, exist_ok=True)

# ── AUTH ──────────────────────────────────────────────────────────────────────
def login_required(f):
    @wraps(f)
    def decorated(*args, **kwargs):
        if "username" not in session:
            return redirect(url_for("login"))
        return f(*args, **kwargs)
    return decorated

# ── HELPERS ───────────────────────────────────────────────────────────────────
def get_filename_from_url(url, index, base_name, total):
    path = urlparse(url).path
    name = os.path.basename(path)
    pad  = len(str(total))
    if base_name:
        return f"{base_name}_{str(index).zfill(pad)}.mp3"
    if name.lower().endswith(".mp3") and name:
        return name
    return f"recording_{str(index).zfill(pad)}.mp3"

def download_mp3_bytes(url):
    response = requests.get(url, stream=True, timeout=30)
    response.raise_for_status()
    buf = io.BytesIO()
    for chunk in response.iter_content(chunk_size=8192):
        buf.write(chunk)
    buf.seek(0)
    return buf

def extract_urls_from_sheet(ws):
    urls = []
    for row in ws.iter_rows(min_row=2):
        cell = row[HYPERLINK_COLUMN - 1]
        url  = None
        if cell.hyperlink:
            url = cell.hyperlink.target
        elif cell.value and str(cell.value).startswith("http"):
            url = str(cell.value).strip()
        if url:
            urls.append(url)
    return urls

def safe_folder_name(name):
    return "".join(c for c in name if c not in r'\/:*?"<>|').strip() or "sheet"

# ── CORE MULTI-SHEET JOB ──────────────────────────────────────────────────────
def process_multi_sheet_job(job_id, filepath, base_name, sheet_names):
    job = jobs[job_id]
    try:
        wb = openpyxl.load_workbook(filepath)

        # Collect URLs per sheet
        sheet_url_map = {}
        for sname in sheet_names:
            if sname not in wb.sheetnames:
                job["sheet_status"][sname]["status"] = "error"
                job["sheet_status"][sname]["error"]  = f"Sheet '{sname}' not found."
                continue
            urls = extract_urls_from_sheet(wb[sname])
            sheet_url_map[sname] = urls
            job["sheet_status"][sname]["total"]  = len(urls)
            job["sheet_status"][sname]["status"] = "running"

        grand_total = sum(len(v) for v in sheet_url_map.values())
        job["total"] = grand_total

        if grand_total == 0:
            job["status"] = "error"
            job["error"]  = "No URLs found in Column K across the selected sheets."
            return

        # Each sheet is processed by its own thread; results collected here
        all_results   = []
        results_lock  = threading.Lock()
        progress_done = [0]

        def process_sheet(sheet_name, urls):
            ss    = job["sheet_status"][sheet_name]
            total = len(urls)
            local = []
            for i, url in enumerate(urls, start=1):
                fname = get_filename_from_url(url, i, base_name, total)
                try:
                    data = download_mp3_bytes(url)
                    local.append((sheet_name, fname, "ok", data, None))
                    ss["files"].append({"name": fname, "status": "ok", "reason": None})
                except Exception as e:
                    local.append((sheet_name, fname, "failed", None, str(e)))
                    ss["files"].append({"name": fname, "status": "failed", "reason": str(e)})
                ss["progress"] += 1
                progress_done[0] += 1
                job["progress"] = progress_done[0]
            ss["status"] = "done"
            with results_lock:
                all_results.extend(local)

        with ThreadPoolExecutor(max_workers=min(3, len(sheet_url_map))) as ex:
            futures = [ex.submit(process_sheet, sn, urls) for sn, urls in sheet_url_map.items()]
            for f in as_completed(futures):
                f.result()  # re-raise any unexpected exception

        # Pack everything into one ZIP with per-sheet subfolders
        zip_buf = io.BytesIO()
        with zipfile.ZipFile(zip_buf, "w", zipfile.ZIP_DEFLATED) as zf:
            for (sheet_name, fname, status, data, _) in all_results:
                if status == "ok" and data:
                    data.seek(0)
                    zf.writestr(f"{safe_folder_name(sheet_name)}/{fname}", data.read())

        zip_buf.seek(0)
        zip_path = os.path.join(UPLOAD_FOLDER, f"{job_id}.zip")
        with open(zip_path, "wb") as fh:
            fh.write(zip_buf.read())

        job["zip_path"] = zip_path
        job["status"]   = "done"

    except Exception as e:
        job["status"] = "error"
        job["error"]  = str(e)
    finally:
        try:
            os.remove(filepath)
        except Exception:
            pass

# ── ROUTES ────────────────────────────────────────────────────────────────────
@app.route("/", methods=["GET"])
@login_required
def index():
    return render_template("index.html", username=session["username"])

@app.route("/login", methods=["GET", "POST"])
def login():
    if request.method == "POST":
        username = request.form.get("username", "").strip()
        password = request.form.get("password", "")
        if USERS.get(username) == password:
            session["username"] = username
            return redirect(url_for("index"))
        flash("Invalid username or password.")
    return render_template("login.html")

@app.route("/logout")
def logout():
    session.clear()
    return redirect(url_for("login"))

@app.route("/peek", methods=["POST"])
@login_required
def peek():
    f = request.files.get("excel_file")
    if not f or not f.filename.endswith((".xlsx", ".xls")):
        return jsonify({"error": "Please upload a valid .xlsx file."}), 400
    filename = secure_filename(f.filename)
    filepath = os.path.join(UPLOAD_FOLDER, f"peek_{uuid.uuid4()}_{filename}")
    f.save(filepath)
    try:
        wb     = openpyxl.load_workbook(filepath, read_only=True, data_only=True)
        sheets = list(wb.sheetnames)
        wb.close()
    except Exception as e:
        return jsonify({"error": f"Could not read workbook: {e}"}), 400
    finally:
        try:
            os.remove(filepath)
        except Exception:
            pass
    return jsonify({"sheets": sheets})

@app.route("/upload", methods=["POST"])
@login_required
def upload():
    f         = request.files.get("excel_file")
    base_name = request.form.get("base_name", "").strip()
    try:
        sheet_names = json.loads(request.form.get("sheet_names", "[]"))
    except Exception:
        sheet_names = []

    if not f or not f.filename.endswith((".xlsx", ".xls")):
        return jsonify({"error": "Please upload a valid .xlsx file."}), 400
    if not sheet_names:
        return jsonify({"error": "Please select at least one sheet."}), 400
    if len(sheet_names) > 3:
        return jsonify({"error": "Maximum 3 sheets can be processed at once."}), 400

    filename = secure_filename(f.filename)
    filepath = os.path.join(UPLOAD_FOLDER, f"{uuid.uuid4()}_{filename}")
    f.save(filepath)

    job_id = str(uuid.uuid4())
    with jobs_lock:
        jobs[job_id] = {
            "status":      "running",
            "progress":    0,
            "total":       0,
            "error":       None,
            "zip_path":    None,
            "sheet_names": sheet_names,
            "sheet_status": {
                s: {"status": "queued", "progress": 0, "total": 0, "files": [], "error": None}
                for s in sheet_names
            },
        }

    threading.Thread(
        target=process_multi_sheet_job,
        args=(job_id, filepath, base_name, sheet_names),
        daemon=True
    ).start()

    return jsonify({"job_id": job_id})

@app.route("/status/<job_id>")
@login_required
def status(job_id):
    job = jobs.get(job_id)
    if not job:
        return jsonify({"error": "Job not found"}), 404
    return jsonify(job)

@app.route("/download/<job_id>")
@login_required
def download(job_id):
    job = jobs.get(job_id)
    if not job or job["status"] != "done":
        return "Not ready", 404
    return send_file(
        job["zip_path"],
        as_attachment=True,
        download_name="recordings.zip",
        mimetype="application/zip"
    )

if __name__ == "__main__":
    app.run(debug=True, port=5000)