"""Karaoke mode — vocal isolation via unsandbox zerotrust containers.

Downloads media from S3, sends it + voxsplit.c into a zerotrust container
(no network, destroyed after use), gets back both instrumentals and vocals,
uploads them to S3.

Uses the unsandbox Python SDK (~/git/un-inception/un.py) for API auth.
"""

import base64
import logging
import os
import sys
import tempfile
import threading
from concurrent.futures import ThreadPoolExecutor, as_completed

log = logging.getLogger(__name__)

# Use the unsandbox Python SDK for API requests
sys.path.insert(0, os.path.expanduser("~/git/un-inception"))
from un import _make_request, _resolve_credentials  # noqa: E402

VOXSPLIT_SOURCE = os.environ.get(
    "VOXSPLIT_SOURCE_PATH",
    os.path.expanduser("~/git/uncloseai-cli/audiotoys/voxsplit.c"),
)


def process_karaoke(s3_client, bucket, s3_key, s3_path,
                    is_video, extension, public_key=None, secret_key=None):
    """Download from S3 -> process in unsandbox zerotrust -> upload both tracks.

    Produces two files under s3_path:
        {s3_path}/instrumentals
        {s3_path}/vocals

    Returns dict {"instrumentals": size, "vocals": size} on success, None on failure.
    """
    public_key, secret_key = _resolve_credentials(public_key, secret_key)

    with tempfile.TemporaryDirectory(prefix="karaoke_") as tmpdir:
        # 1. Download original from S3
        local_path = os.path.join(tmpdir, "media")
        s3_client.download_file(bucket, s3_key, local_path)

        # 2. Read and base64 encode
        with open(local_path, "rb") as f:
            media_b64 = base64.b64encode(f.read()).decode()
        with open(VOXSPLIT_SOURCE, "r") as f:
            voxsplit_src = f.read()
        voxsplit_b64 = base64.b64encode(voxsplit_src.encode()).decode()

        # 3. Build execution script — produce both tracks
        if is_video:
            inst_name = f"instrumentals.{extension}"
            vocals_name = f"vocals.{extension}"
            script = (
                "gcc -O2 -o /tmp/voxsplit /tmp/input/voxsplit.c -lm && "
                "/tmp/voxsplit /tmp/input/media -o /tmp/split && "
                "mkdir -p /tmp/artifacts && "
                f"ffmpeg -y -i /tmp/input/media -i /tmp/split-instrumental.wav "
                f"-c:v copy -map 0:v -map 1:a -shortest /tmp/artifacts/{inst_name} && "
                f"ffmpeg -y -i /tmp/input/media -i /tmp/split-vocal.wav "
                f"-c:v copy -map 0:v -map 1:a -shortest /tmp/artifacts/{vocals_name}"
            )
        else:
            inst_name = "instrumentals.wav"
            vocals_name = "vocals.wav"
            script = (
                "gcc -O2 -o /tmp/voxsplit /tmp/input/voxsplit.c -lm && "
                "/tmp/voxsplit /tmp/input/media -o /tmp/split && "
                "mkdir -p /tmp/artifacts && "
                "cp /tmp/split-instrumental.wav /tmp/artifacts/instrumentals.wav && "
                "cp /tmp/split-vocal.wav /tmp/artifacts/vocals.wav"
            )

        # 4. Call unsandbox execute API via SDK
        payload = {
            "language": "bash",
            "code": script,
            "network_mode": "zerotrust",
            "ttl": 300,
            "vcpu": 2,
            "input_files": [
                {"filename": "media", "content": media_b64},
                {"filename": "voxsplit.c", "content": voxsplit_b64},
            ],
        }

        try:
            result = _make_request("POST", "/execute", public_key, secret_key, payload)
        except Exception:
            return None

        if not result.get("success") or result.get("exit_code") != 0:
            return None

        # 5. Extract artifacts and upload to S3
        artifacts = result.get("artifacts", [])
        if len(artifacts) < 2:
            return None

        # Map artifact filenames to track names
        artifact_map = {}
        for a in artifacts:
            fname = a.get("filename", "")
            if fname.startswith("instrumentals"):
                artifact_map["instrumentals"] = a
            elif fname.startswith("vocals"):
                artifact_map["vocals"] = a

        if "instrumentals" not in artifact_map or "vocals" not in artifact_map:
            return None

        sizes = {}
        content_type = "video/mp4" if is_video else "audio/wav"

        for track_name in ("instrumentals", "vocals"):
            artifact = artifact_map[track_name]
            data = base64.b64decode(artifact["content_base64"])
            local_file = os.path.join(tmpdir, track_name)
            with open(local_file, "wb") as f:
                f.write(data)

            s3_client.upload_file(
                local_file, bucket, f"{s3_path}/{track_name}",
                ExtraArgs={
                    "CacheControl": "private, max-age=172800",
                    "ContentType": content_type,
                },
            )
            sizes[track_name] = len(data)

        return sizes


def tracks_exist(s3_client, bucket, s3_path):
    """Check if instrumentals and vocals already exist in S3."""
    for track in ("instrumentals", "vocals"):
        try:
            s3_client.head_object(Bucket=bucket, Key=f"{s3_path}/{track}")
        except Exception:
            return False
    return True


# Background backfill guard — one per shop, same pattern as reforge_discovery_ring_async
_backfill_lock = threading.Lock()
_backfill_running = set()


def backfill_karaoke_async(shop_id, session_factory, app_settings):
    """Backfill vocal isolation tracks for a shop's catalog using a thread pool.

    Pool size matches the unsandbox account's concurrency limit (from validate_keys).
    """
    shop_id_str = str(shop_id)

    with _backfill_lock:
        if shop_id_str in _backfill_running:
            return  # already running for this shop
        _backfill_running.add(shop_id_str)

    def _backfill():
        try:
            import boto3
            from sqlalchemy.orm import Session as SASession
            from ..models.product import Product, get_media_type
            from ..models.shop import Shop
            from un import validate_keys

            session = SASession(bind=session_factory().get_bind())

            def _make_s3():
                return boto3.session.Session().client(
                    "s3",
                    region_name=app_settings["bucket.secure_uploads.region"],
                    endpoint_url=app_settings["bucket.secure_uploads.post_endpoint"],
                    aws_access_key_id=app_settings["bucket.secure_uploads.access_key"],
                    aws_secret_access_key=app_settings["bucket.secure_uploads.secret_key"],
                )

            s3 = _make_s3()
            bucket = app_settings["bucket.secure_uploads"]

            try:
                shop = session.get(Shop, shop_id)
                if not shop or not shop.unsandbox_public_key or not shop.unsandbox_secret_key:
                    return

                pk, sk = shop.unsandbox_public_key, shop.unsandbox_secret_key

                # Query account concurrency from unsandbox API
                try:
                    key_info = validate_keys(pk, sk)
                    concurrency = int(key_info.get("concurrency", 1))
                except Exception:
                    concurrency = 1
                concurrency = max(1, concurrency)
                log.info("Backfill karaoke for shop %s — %d concurrent workers", shop.name, concurrency)

                # Collect work items
                work = []
                products = session.query(Product).filter(Product.shop_id == shop_id).all()
                for product in products:
                    for file_key in ("product", "preview"):
                        ext = product.extensions.get(file_key)
                        if not ext:
                            continue
                        media_type = get_media_type(ext)
                        if media_type not in ("video", "audio"):
                            continue
                        if tracks_exist(s3, bucket, product.s3_path):
                            continue
                        work.append((product.id, product.s3_path, product.title, file_key, ext, media_type))

                if not work:
                    log.info("No products need backfill for shop %s", shop.name)
                    return

                MAX_RETRIES = 3
                BACKOFF_BASE = 5  # seconds

                # Process a single product with retries (runs in pool thread)
                def _process_one(item):
                    product_id, s3_path, title, file_key, ext, media_type = item
                    thread_s3 = _make_s3()
                    is_video = (media_type == "video")

                    for attempt in range(MAX_RETRIES):
                        log.info("Backfill karaoke: %s (%s) attempt %d/%d",
                                 title, file_key, attempt + 1, MAX_RETRIES)
                        try:
                            sizes = process_karaoke(
                                thread_s3, bucket,
                                f"{s3_path}/{file_key}",
                                s3_path, is_video, ext,
                                public_key=pk, secret_key=sk,
                            )
                            if sizes:
                                return product_id, title, file_key, ext, media_type, sizes
                        except Exception:
                            log.exception("Backfill attempt %d failed for %s", attempt + 1, title)

                        if attempt < MAX_RETRIES - 1:
                            import time as _time
                            wait = BACKOFF_BASE * (2 ** attempt)
                            log.info("  Retrying %s in %ds...", title, wait)
                            _time.sleep(wait)

                    return product_id, title, file_key, ext, media_type, None

                # Fan out with thread pool
                with ThreadPoolExecutor(max_workers=concurrency) as pool:
                    futures = {pool.submit(_process_one, item): item for item in work}
                    for future in as_completed(futures):
                        try:
                            product_id, title, file_key, ext, media_type, sizes = future.result()
                        except Exception:
                            item = futures[future]
                            log.exception("Backfill worker crashed for %s", item[2])
                            continue

                        if sizes:
                            # DB write only on success — no partial state
                            product = session.get(Product, product_id)
                            if not product:
                                continue
                            track_ext = ext if media_type == "video" else "wav"
                            for track_name in ("instrumentals", "vocals"):
                                product.set_file_metadata(track_name, track_ext, f"{track_name}.{track_ext}")
                            tmp = product.file_bytes
                            tmp.update(sizes)
                            product.file_bytes = tmp
                            session.add(product)
                            session.commit()
                            log.info("  OK %s (inst=%dB, vox=%dB)", title, sizes["instrumentals"], sizes["vocals"])
                        else:
                            log.warning("  FAILED %s after %d attempts", title, MAX_RETRIES)

            except Exception:
                session.rollback()
                log.exception("Background karaoke backfill failed for shop %s", shop_id_str)
            finally:
                session.close()
        finally:
            with _backfill_lock:
                _backfill_running.discard(shop_id_str)

    t = threading.Thread(target=_backfill, daemon=True)
    t.start()
