#!/usr/bin/env python3

"""Download episodes from Kaizoku's StrawVerse-compatible provider bridge."""

import argparse
import json
import os
import re
import shutil
import subprocess
import sys
from urllib.parse import urljoin, urlparse
import urllib.request
from pathlib import Path


PROJECT_ROOT = Path(__file__).resolve().parent
BRIDGE = PROJECT_ROOT / "providers" / "bridge.js"
PROVIDERS = ("anikoto", "anineko", "pahe")


def clean_component(value, default="Anime"):
    text = re.sub(r"[\\/:*?\"<>|]+", " ", str(value or "")).strip()
    text = re.sub(r"\s+", " ", text)
    return text.strip(". ") or default


def normalize_title(value):
    return re.sub(r"[^a-z0-9]+", " ", str(value or "").lower()).strip()


def parse_show_id(value, default_provider="anikoto"):
    text = str(value or "").strip()
    if ":" in text:
        provider, show_id = text.split(":", 1)
        provider = provider.strip().lower()
        if provider in PROVIDERS and show_id.strip():
            return provider, show_id.strip()
    return default_provider, text


def episode_key(value):
    text = str(value or "").strip()
    if not text:
        return ""
    try:
        number = float(text)
    except ValueError:
        return text
    if number.is_integer():
        return str(int(number))
    return str(number).rstrip("0").rstrip(".")


def bridge(command, provider, *args):
    cmd = ["node", str(BRIDGE), command, provider, *[str(arg) for arg in args]]
    proc = subprocess.run(cmd, cwd=PROJECT_ROOT, text=True, capture_output=True, check=False)
    if proc.returncode != 0:
        raise RuntimeError((proc.stderr or proc.stdout or f"{command} failed").strip())
    try:
        return json.loads(proc.stdout or "{}")
    except json.JSONDecodeError as exc:
        raise RuntimeError(f"Provider bridge returned invalid JSON: {exc}") from exc


def provider_order(primary):
    ordered = []
    primary = str(primary or "").strip().lower()
    if primary in PROVIDERS:
        ordered.append(primary)
    for provider in PROVIDERS:
        if provider not in ordered:
            ordered.append(provider)
    return ordered


def find_provider_show(provider, title, original_provider, original_show_id):
    if provider == original_provider:
        return original_show_id
    query = str(title or "").strip()
    if not query:
        return ""
    payload = bridge("search", provider, query, 1)
    results = payload.get("results") or []
    if not results:
        return ""
    target = normalize_title(query)
    scored = []
    for item in results:
        candidate = normalize_title(item.get("title"))
        if not candidate:
            continue
        if candidate == target:
            score = 100
        elif target and (candidate in target or target in candidate):
            score = 75 - abs(len(candidate) - len(target))
        else:
            shared = set(target.split()) & set(candidate.split())
            score = len(shared) * 5
        scored.append((score, item))
    if not scored:
        return ""
    scored.sort(key=lambda row: row[0], reverse=True)
    if scored[0][0] <= 0:
        return ""
    provider_id = str(scored[0][1].get("provider_id") or scored[0][1].get("id") or "").strip()
    if ":" in provider_id:
        provider_id = provider_id.split(":", 1)[1].strip()
    return provider_id


def episode_values(spec, episodes):
    values = []
    by_number = {episode_key(ep.get("number")): ep for ep in episodes}
    if not str(spec or "").strip():
        return episodes
    for part in re.split(r"[\s,]+", str(spec).strip()):
        if not part:
            continue
        if "-" in part:
            start, end = part.split("-", 1)
            try:
                left = int(float(start))
                right = int(float(end))
            except ValueError:
                continue
            step = 1 if right >= left else -1
            for number in range(left, right + step, step):
                item = by_number.get(str(number))
                if item and item not in values:
                    values.append(item)
            continue
        item = by_number.get(episode_key(part)) if re.match(r"^\d+(?:\.\d+)?$", part) else None
        if item and item not in values:
            values.append(item)
    return values


def requested_episode_numbers(spec, episodes):
    text = str(spec or "").strip()
    if not text:
        return [episode.get("number") for episode in episodes or [] if episode.get("number") not in (None, "")]
    numbers = []
    for part in re.split(r"[\s,]+", text):
        if not part:
            continue
        if "-" in part:
            start, end = part.split("-", 1)
            try:
                left = int(float(start))
                right = int(float(end))
            except ValueError:
                continue
            step = 1 if right >= left else -1
            numbers.extend(range(left, right + step, step))
            continue
        if re.match(r"^\d+(?:\.\d+)?$", part):
            numbers.append(float(part) if "." in part else int(part))
    return numbers


def wanted_episodes(spec, episodes):
    by_number = {episode_key(ep.get("number")): ep for ep in episodes or []}
    wanted = []
    for number in requested_episode_numbers(spec, episodes):
        key = episode_key(number)
        wanted.append(by_number.get(key) or {"number": number, "id": ""})
    return wanted


def episode_by_number(number, episodes):
    wanted = episode_key(number)
    for episode in episodes or []:
        current = episode_key(episode.get("number"))
        if current == wanted:
            return episode
    return None


def ffmpeg_headers(headers):
    pairs = []
    for key, value in (headers or {}).items():
        if value:
            pairs.append(f"{key}: {value}")
    return "\r\n".join(pairs) + ("\r\n" if pairs else "")


def request_headers(headers=None, referer=None):
    merged = {
        "User-Agent": "Mozilla/5.0",
        "Accept": "*/*",
        "Accept-Language": "en-US,en;q=0.9",
        "Cache-Control": "no-cache",
        "Pragma": "no-cache",
    }
    merged.update(headers or {})
    if referer and not (merged.get("Referer") or merged.get("referer")):
        merged["Referer"] = referer
    final_referer = merged.get("Referer") or merged.get("referer")
    if final_referer and not (merged.get("Origin") or merged.get("origin")):
        try:
            parsed = urlparse(final_referer)
            if parsed.scheme and parsed.netloc:
                merged["Origin"] = f"{parsed.scheme}://{parsed.netloc}"
        except Exception:
            pass
    return merged


def fetch_bytes(url, headers=None, timeout=30):
    request = urllib.request.Request(url, headers=request_headers(headers))
    with urllib.request.urlopen(request, timeout=timeout) as response:
        return response.read()


def fetch_text(url, headers=None, timeout=30):
    return fetch_bytes(url, headers=headers, timeout=timeout).decode("utf-8", errors="replace")


def best_hls_variant_url(master_url, playlist_text):
    variants = []
    pending_bandwidth = 0
    for raw_line in str(playlist_text or "").splitlines():
        line = raw_line.strip()
        if not line:
            continue
        if line.startswith("#EXT-X-STREAM-INF"):
            match = re.search(r"BANDWIDTH=(\d+)", line)
            pending_bandwidth = int(match.group(1)) if match else 0
            continue
        if pending_bandwidth and not line.startswith("#"):
            variants.append((pending_bandwidth, urljoin(master_url, line)))
            pending_bandwidth = 0
    if not variants:
        return master_url
    variants.sort(key=lambda item: item[0], reverse=True)
    return variants[0][1]


def resolve_hls_input_url(stream):
    url = stream.get("url")
    if not stream.get("isM3U8") and ".m3u8" not in str(url):
        return url
    data = fetch_text(url, headers=stream.get("headers") or {})
    if "#EXT-X-STREAM-INF" not in data:
        return url
    return best_hls_variant_url(url, data)


def hls_media_playlist(stream, input_url):
    playlist = fetch_text(input_url, headers=stream.get("headers") or {})
    if "#EXT-X-STREAM-INF" in playlist:
        input_url = best_hls_variant_url(input_url, playlist)
        playlist = fetch_text(input_url, headers=stream.get("headers") or {})
    return input_url, playlist


def ffmpeg_base_command(stream, input_url):
    cmd = [
        "ffmpeg",
        "-hide_banner",
        "-loglevel",
        "info",
        "-y",
        "-protocol_whitelist",
        "file,http,https,tcp,tls,crypto",
        "-allowed_extensions",
        "ALL",
        "-allowed_segment_extensions",
        "ALL",
    ]
    headers = ffmpeg_headers(stream.get("headers") or {})
    if headers:
        cmd.extend(["-headers", headers])
    referer = (stream.get("headers") or {}).get("Referer")
    if referer:
        cmd.extend(["-referer", referer])
    cmd.extend(["-i", input_url])
    return cmd


def strip_png_header(data):
    png_header = b"\x89PNG\r\n\x1a\n"
    if not data.startswith(png_header):
        return data
    iend_offset = data.find(b"IEND")
    if iend_offset != -1 and iend_offset < 1024:
        return data[iend_offset + 8:]
    return data


def parse_hls_segments(playlist_url, playlist_text):
    segments = []
    key_url = None
    iv = None
    media_sequence = 1
    segment_index = 0
    for raw_line in str(playlist_text or "").splitlines():
        line = raw_line.strip()
        if not line:
            continue
        if line.startswith("#EXT-X-MEDIA-SEQUENCE:"):
            try:
                media_sequence = int(line.split(":", 1)[1])
            except ValueError:
                media_sequence = 1
            continue
        if line.startswith("#EXT-X-KEY:"):
            attrs = {
                match.group(1): match.group(2) if match.group(2) is not None else match.group(3)
                for match in re.finditer(r'([A-Z0-9_-]+)=(?:"([^"]*)"|([^,]*))', line.split(":", 1)[1])
            }
            if str(attrs.get("METHOD") or "").upper() == "AES-128":
                key_url = urljoin(playlist_url, attrs.get("URI") or "")
                iv = attrs.get("IV")
            else:
                key_url = None
                iv = None
            continue
        if line.startswith("#"):
            continue
        segment = {"url": urljoin(playlist_url, line)}
        if key_url:
            segment["key_url"] = key_url
            segment["iv"] = iv or str(media_sequence + segment_index)
        segments.append(segment)
        segment_index += 1
    return segments


def hls_segments_need_native_download(segments):
    for segment in segments or []:
        path = urlparse(segment.get("url") or "").path
        name = path.rsplit("/", 1)[-1]
        if "." not in name:
            return True
        if "/ad-site-i18n/" in path:
            return True
    return False


def decrypt_aes128_segment(data, key, iv_value):
    openssl = shutil.which("openssl")
    if not openssl:
        raise RuntimeError("Encrypted HLS segment requires openssl, but openssl is not available.")
    iv = bytearray(16)
    text = str(iv_value or "").strip()
    if text.lower().startswith("0x"):
        raw = bytes.fromhex(text[2:])
        iv[:len(raw[:16])] = raw[:16]
    else:
        iv[-4:] = int(text or "0").to_bytes(4, "big")
    proc = subprocess.run(
        [openssl, "enc", "-d", "-aes-128-cbc", "-K", key.hex(), "-iv", bytes(iv).hex()],
        input=data,
        capture_output=True,
        check=False,
    )
    if proc.returncode != 0:
        raise RuntimeError((proc.stderr or b"openssl AES-128 decrypt failed").decode("utf-8", errors="replace").strip())
    return proc.stdout


def download_hls_segments(stream, input_url, target, partial):
    input_url, playlist = hls_media_playlist(stream, input_url)
    segments = parse_hls_segments(input_url, playlist)
    if not segments:
        return 1
    ts_file = target.with_suffix(target.suffix + ".ts.part")
    segment_dir = target.parent / f".segments_{target.stem}"
    if segment_dir.exists():
        shutil.rmtree(segment_dir)
    segment_dir.mkdir(parents=True, exist_ok=True)
    key_cache = {}
    try:
        print(f"Downloading {len(segments)} HLS playlist segments directly...")
        with ts_file.open("wb") as joined:
            for index, segment in enumerate(segments, start=1):
                data = fetch_bytes(segment["url"], headers=stream.get("headers") or {}, timeout=60)
                data = strip_png_header(data)
                key_url = segment.get("key_url")
                if key_url:
                    if key_url not in key_cache:
                        key_cache[key_url] = fetch_bytes(key_url, headers=stream.get("headers") or {}, timeout=30)
                    data = decrypt_aes128_segment(data, key_cache[key_url], segment.get("iv"))
                if not data:
                    raise RuntimeError(f"Segment {index} was empty.")
                joined.write(data)
                if index == 1 or index == len(segments) or index % 25 == 0:
                    print(f"Downloaded segment {index}/{len(segments)}")
        cmd = [
            "ffmpeg",
            "-hide_banner",
            "-loglevel",
            "info",
            "-y",
            "-f",
            "mpegts",
            "-i",
            str(ts_file),
            "-c",
            "copy",
            "-bsf:a",
            "aac_adtstoasc",
            "-movflags",
            "+faststart",
            "-f",
            "mp4",
            str(partial),
        ]
        code = run_ffmpeg(cmd, timeout=300)
        if code == 0 and partial.exists() and partial.stat().st_size > 0:
            partial.replace(target)
            return 0
        return code
    finally:
        try:
            ts_file.unlink()
        except FileNotFoundError:
            pass
        shutil.rmtree(segment_dir, ignore_errors=True)


def run_ffmpeg(cmd, timeout=180):
    try:
        return subprocess.run(cmd, check=False, timeout=timeout).returncode
    except subprocess.TimeoutExpired:
        print(f"ffmpeg timed out after {timeout}s; trying fallback.", file=sys.stderr)
        return 124


def download_subtitles(subtitles, output_base):
    saved = []
    english = [
        sub for sub in subtitles or []
        if sub.get("url") and re.search(r"eng|english", str(sub.get("lang") or ""), re.I)
    ] or [sub for sub in subtitles or [] if sub.get("url")]
    for index, sub in enumerate(english[:1], start=1):
        url = sub["url"]
        suffix = ".srt" if ".srt" in url.lower() else ".vtt"
        target = output_base.with_suffix(f".eng{suffix}" if index == 1 else f".eng-{index}{suffix}")
        request = urllib.request.Request(url, headers=request_headers(referer=sub.get("referer") or ""))
        with urllib.request.urlopen(request, timeout=30) as response:
            target.write_bytes(response.read())
        saved.append(target)
    return saved


def download_episode(stream, target):
    input_url = resolve_hls_input_url(stream)
    partial = target.with_suffix(target.suffix + ".part")
    for path in (target, partial):
        try:
            path.unlink()
        except FileNotFoundError:
            pass
    if stream.get("isM3U8") or ".m3u8" in str(input_url):
        try:
            media_url, playlist = hls_media_playlist(stream, input_url)
            segments = parse_hls_segments(media_url, playlist)
            if hls_segments_need_native_download(segments):
                print("HLS playlist uses extensionless provider segments; skipping direct ffmpeg.")
                code = download_hls_segments(stream, media_url, target, partial)
                if code == 0:
                    return 0
                return code
        except Exception as exc:
            print(f"Native HLS preflight failed: {exc}", file=sys.stderr)
    attempts = [
        [
            *ffmpeg_base_command(stream, input_url),
            "-map",
            "0:v:0?",
            "-map",
            "0:a:0?",
            "-c",
            "copy",
            "-bsf:a",
            "aac_adtstoasc",
            "-movflags",
            "+faststart",
            "-f",
            "mp4",
            str(partial),
        ],
        [
            *ffmpeg_base_command(stream, input_url),
            "-map",
            "0",
            "-c",
            "copy",
            "-bsf:a",
            "aac_adtstoasc",
            "-movflags",
            "+faststart",
            "-f",
            "mp4",
            str(partial),
        ],
    ]
    code = 1
    for cmd in attempts:
        try:
            partial.unlink()
        except FileNotFoundError:
            pass
        code = run_ffmpeg(cmd)
        if code == 0 and partial.exists() and partial.stat().st_size > 0:
            partial.replace(target)
            return 0
    if stream.get("isM3U8") or ".m3u8" in str(input_url):
        try:
            code = download_hls_segments(stream, input_url, target, partial)
            if code == 0:
                return 0
        except Exception as exc:
            print(f"Segment downloader failed: {exc}", file=sys.stderr)
            code = 1
    try:
        partial.unlink()
    except FileNotFoundError:
        pass
    return code


def provider_episode_candidates(primary_provider, primary_show_id, title, wanted_episode):
    number = wanted_episode.get("number")
    for provider in provider_order(primary_provider):
        try:
            provider_show_id = find_provider_show(provider, title, primary_provider, primary_show_id)
            if not provider_show_id:
                print(f"Fallback skipped {provider}: no matching title.")
                continue
            episodes = bridge("episodes", provider, provider_show_id).get("episodes") or []
            episode = episode_by_number(number, episodes)
            if not episode:
                print(f"Fallback skipped {provider}: episode {number} is not listed.")
                continue
            yield provider, provider_show_id, episode
        except Exception as exc:
            print(f"Fallback skipped {provider}: {exc}")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--provider", default="anikoto")
    parser.add_argument("--show-id", required=True)
    parser.add_argument("--title", required=True)
    parser.add_argument("--episodes", required=True)
    parser.add_argument("--mode", choices=("sub", "dub"), default="sub")
    parser.add_argument("--quality", default="best")
    parser.add_argument("--output-dir", required=True)
    args = parser.parse_args()

    provider, provider_show_id = parse_show_id(args.show_id, args.provider)
    info = bridge("info", provider, provider_show_id)
    episodes = bridge("episodes", provider, provider_show_id).get("episodes") or []
    wanted = wanted_episodes(args.episodes, episodes)
    if not wanted:
        raise SystemExit(f"No matching episodes for spec {args.episodes!r}.")

    series_title = clean_component(args.title or info.get("title"))
    output_dir = Path(args.output_dir).expanduser()
    output_dir.mkdir(parents=True, exist_ok=True)

    print(f"Provider: {provider}")
    print(f"Title: {series_title}")
    print(f"Episodes: {', '.join(str(ep.get('number')) for ep in wanted)}")
    exit_code = 0
    for ep in wanted:
        number = ep.get("number")
        ep_id = ep.get("id")
        padded = f"{int(float(number)):02d}" if str(number).replace(".", "", 1).isdigit() else str(number)
        basename = f"{series_title} - S01E{padded}"
        target = output_dir / f"{basename}.mp4"
        last_error = None
        downloaded = False
        for active_provider, _active_show_id, active_ep in provider_episode_candidates(provider, provider_show_id, args.title or info.get("title"), ep):
            try:
                active_ep_id = active_ep.get("id")
                print(f"Resolving episode {number} on {active_provider} ({args.mode}, {args.quality})...")
                stream = bridge("resolve", active_provider, active_ep_id, args.mode, args.quality)
                print(f"Downloading episode {number} from {active_provider} {stream.get('quality') or 'auto'}...")
                code = download_episode(stream, target)
                if code != 0:
                    last_error = RuntimeError(f"ffmpeg failed on {active_provider} with exit code {code}")
                    print(str(last_error), file=sys.stderr)
                    continue
                try:
                    for subtitle_path in download_subtitles(stream.get("subtitles") or [], target):
                        print(f"Saved subtitle: {subtitle_path.name}")
                except Exception as exc:
                    print(f"Subtitle download skipped: {exc}")
                print(f"Saved: {target.name}")
                downloaded = True
                break
            except Exception as exc:
                last_error = exc
                print(f"Provider {active_provider} failed for episode {number}: {exc}", file=sys.stderr)
        if not downloaded:
            print(f"Episode {number} failed on all providers: {last_error}", file=sys.stderr)
            exit_code = 1
            break
    raise SystemExit(exit_code)


if __name__ == "__main__":
    main()