app.py

import os
import re
import json
import shutil
import requests
import tempfile
import mimetypes
from urllib.parse import urlparse, parse_qs, quote
from bs4 import BeautifulSoup
import gradio as gr

try:
    from PIL import Image
    PIL_AVAILABLE = True
except ImportError:  # Pillow not installed — degrade gracefully
    PIL_AVAILABLE = False

HEADERS = {
    "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0 Safari/537.36"
}

# Optional: set a GITHUB_TOKEN env var to raise the GitHub API rate limit from
# 60/hr (unauthenticated, shared across everyone on a host like HF Spaces) to
# 5000/hr. A "public_repo" read-only token is enough.
GITHUB_TOKEN = os.environ.get("GITHUB_TOKEN")
GITHUB_API_HEADERS = {
    **HEADERS,
    "Accept": "application/vnd.github+json",
    "X-GitHub-Api-Version": "2022-11-28",
}
if GITHUB_TOKEN:
    GITHUB_API_HEADERS["Authorization"] = f"Bearer {GITHUB_TOKEN}"

KOMI_DOMAINS = ("komi.store", "komistore.app", "komistore.com", "github-store.org")

# Common conventional locations for an app icon inside a repo, used both to
# rank real tree results and as a blind fallback if the Trees API is
# unavailable (rate-limited / private repo).
ICON_FILENAME_RE = re.compile(
    r"(fastlane/.+/icon|ic_launcher(_round)?|playstore.?icon|app.?icon|logo|icon)[^/]*\.(png|webp|jpg|jpeg)$",
    re.I,
)
DENSITY_RANK = {"xxxhdpi": 6, "xxhdpi": 5, "xhdpi": 4, "hdpi": 3, "mdpi": 2, "ldpi": 1}
COMMON_ICON_GUESS_PATHS = [
    "fastlane/metadata/android/en-US/images/icon.png",
    "app/src/main/res/mipmap-xxxhdpi/ic_launcher.png",
    "app/src/main/res/mipmap-xxhdpi/ic_launcher.png",
    "src/main/res/mipmap-xxxhdpi/ic_launcher.png",
    ".github/logo.png",
    ".github/assets/logo.png",
    "assets/icon.png",
    "assets/logo.png",
    "docs/logo.png",
    "logo.png",
    "icon.png",
]

GITHUB_REPO_IN_TEXT_RE = re.compile(r"github\.com/([A-Za-z0-9_.-]+)/([A-Za-z0-9_.-]+)")


# ---------------------------------------------------------------------------
# Step 1: analyze the link and classify its source
# ---------------------------------------------------------------------------

def identify_source(input_str: str):
    input_str = input_str.strip()
    if not input_str:
        raise ValueError('Enter a Play Store, GitHub, F-Droid, or Komi Store URL or package id.')

    try:
        u = urlparse(input_str)
    except Exception:
        u = None

    if u and u.scheme in ('http', 'https'):
        netloc = u.netloc.lower()

        if 'play.google.com' in netloc:
            qs = parse_qs(u.query)
            if 'id' in qs:
                return 'play', qs['id'][0]
            m = re.search(r'id=([a-zA-Z0-9_\.]+)', input_str)
            if m:
                return 'play', m.group(1)
            raise ValueError('Invalid Google Play URL.')

        if 'f-droid.org' in netloc:
            m = re.search(r'/packages/([a-zA-Z0-9_.]+)', u.path)
            if m:
                return 'fdroid', m.group(1)
            qs = parse_qs(u.query)
            if 'fdid' in qs:
                return 'fdroid', qs['fdid'][0]
            raise ValueError('Invalid F-Droid URL: could not find a package id in the path.')

        if 'github.com' in netloc:
            path = u.path.strip('/')
            parts = [part for part in path.split('/') if part]
            if len(parts) >= 2:
                owner, repo = parts[0], re.sub(r'\.git$', '', parts[1])
                return 'github', {'owner': owner, 'repo': repo, 'url': f'https://github.com/{owner}/{repo}'}
            raise ValueError('GitHub URL must include owner and repository.')

        if any(d in netloc for d in KOMI_DOMAINS):
            return 'komi', input_str

    # bare "owner/repo" shorthand -> GitHub
    m = re.match(r'^([A-Za-z0-9_.-]+)/([A-Za-z0-9_.-]+)$', input_str)
    if m and '.' not in m.group(1):
        owner, repo = m.group(1), re.sub(r'\.git$', '', m.group(2))
        return 'github', {'owner': owner, 'repo': repo, 'url': f'https://github.com/{owner}/{repo}'}

    # bare package name -> Play Store
    if re.match(r'^[a-zA-Z0-9_\.]+\.[a-zA-Z0-9_\.]+$', input_str):
        return 'play', input_str

    raise ValueError('Unsupported source or invalid URL.')


def describe_source(source: str, data) -> str:
    if source == 'play':
        return f'Detected: Google Play Store — package "{data}"'
    if source == 'fdroid':
        return f'Detected: F-Droid — package "{data}"'
    if source == 'github':
        return f'Detected: GitHub repository — {data["owner"]}/{data["repo"]}'
    if source == 'komi':
        return 'Detected: Komi Store (client-rendered — will resolve via underlying GitHub/Codeberg repo)'
    return 'Detected: generic page'


# ---------------------------------------------------------------------------
# Shared helpers
# ---------------------------------------------------------------------------

def fetch_page(url: str) -> str:
    r = requests.get(url, headers=HEADERS, timeout=20)
    r.raise_for_status()
    return r.text


def fetch_json(url: str, headers=None):
    r = requests.get(url, headers=headers or HEADERS, timeout=20)
    r.raise_for_status()
    return r.json()


def find_icon_url_from_html(html: str, source: str = 'generic', repo_data=None):
    soup = BeautifulSoup(html, 'lxml')

    meta = soup.find('meta', property='og:image')
    if meta and meta.get('content'):
        return meta['content']

    meta = soup.find('meta', attrs={'name': 'twitter:image'})
    if meta and meta.get('content'):
        return meta['content']

    link = soup.find('link', rel='image_src')
    if link and link.get('href'):
        return link['href']

    script = soup.find('script', type='application/ld+json')
    if script and script.string:
        try:
            data = json.loads(script.string)
            if isinstance(data, dict) and 'image' in data:
                return data['image']
        except Exception:
            pass

    candidates = []
    for img in soup.find_all('img'):
        src = img.get('src') or img.get('data-src')
        alt = (img.get('alt') or '').lower()
        cls = ' '.join(img.get('class') or []).lower()
        if not src:
            continue
        if 'icon' in alt or 'logo' in alt or 'avatar' in alt or 'app' in alt:
            candidates.append(src)
        elif 'avatar' in cls or 'logo' in cls or 'icon' in cls:
            candidates.append(src)

    if candidates:
        return candidates[0]

    if source == 'github' and repo_data:
        owner = repo_data.get('owner')
        if owner:
            return f'https://github.com/{owner}.png?size=1024'

    raise RuntimeError('Could not find icon URL on the page.')


def upscale_google_image_url(url: str, target_size: int = 1024) -> str:
    new = re.sub(r'=s\d+(-[a-zA-Z0-9_-]+)?', f'=s{target_size}', url)
    new = re.sub(r'=w\d+-h\d+(-[a-zA-Z0-9_-]+)?', f'=w{target_size}-h{target_size}-rw', new)
    return new


def download_image_to_temp(url: str) -> str:
    r = requests.get(url, headers=HEADERS, timeout=30, stream=True)
    r.raise_for_status()
    content_type = r.headers.get('Content-Type', '')
    if 'image' not in content_type and not re.search(r'\.(png|jpe?g|webp|gif)$', urlparse(url).path, re.I):
        raise RuntimeError(f'Not an image response ({content_type or "unknown type"}).')
    ext = mimetypes.guess_extension(content_type.split(';')[0]) or os.path.splitext(urlparse(url).path)[1] or '.png'
    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=ext, prefix='icon_fetch_')
    size = 0
    with open(tmp.name, 'wb') as f:
        for chunk in r.iter_content(4096):
            if chunk:
                size += len(chunk)
                f.write(chunk)
    if size < 256:  # guard against 1x1 tracking pixels / empty placeholder responses
        raise RuntimeError('Downloaded file is too small to be a real icon.')
    return tmp.name


def image_dimensions(path: str):
    if not PIL_AVAILABLE:
        return (0, 0)
    try:
        with Image.open(path) as im:
            return im.size
    except Exception:
        return (0, 0)


def save_as_logo_png(path: str) -> str:
    out_dir = tempfile.mkdtemp(prefix='icon_fetch_')
    out_path = os.path.join(out_dir, 'logo.png')
    if PIL_AVAILABLE:
        try:
            with Image.open(path) as im:
                if im.mode not in ('RGB', 'RGBA', 'LA', 'L'):
                    im = im.convert('RGBA')
                im.save(out_path, format='PNG')
            return out_path
        except Exception:
            pass
    shutil.copyfile(path, out_path)  # fallback: raw copy, keep original bytes
    return out_path


# ---------------------------------------------------------------------------
# Step 2: per-source, quality-ranked candidate resolution
# ---------------------------------------------------------------------------

def candidates_play(package_id: str):
    url = f'https://play.google.com/store/apps/details?id={package_id}&hl=en&gl=US'
    html = fetch_page(url)
    icon_url = find_icon_url_from_html(html, source='play')
    return [upscale_google_image_url(icon_url, target_size=1024)]


def candidates_fdroid(package_id: str):
    urls = []
    try:
        data = fetch_json(f'https://f-droid.org/api/v1/packages/{package_id}')
        vcode = data.get('suggestedVersionCode')
        if not vcode and data.get('packages'):
            vcode = data['packages'][0].get('versionCode')
        if vcode:
            for density in ('640', '480', '320', '240', '120'):
                urls.append(f'https://f-droid.org/repo/icons-{density}/{package_id}.{vcode}.png')
            urls.append(f'https://f-droid.org/repo/icons/{package_id}.{vcode}.png')
    except Exception:
        pass

    for page in (f'https://f-droid.org/en/packages/{package_id}/', f'https://f-droid.org/packages/{package_id}/'):
        try:
            html = fetch_page(page)
            urls.append(find_icon_url_from_html(html, source='fdroid'))
        except Exception:
            continue

    if not urls:
        raise RuntimeError(f'Could not resolve an icon for F-Droid package "{package_id}".')
    return urls


def score_github_path(path: str) -> int:
    lower = path.lower()
    score = 0
    if 'fastlane' in lower and lower.endswith('icon.png'):
        score += 1000
    if 'playstore' in lower:
        score += 200
    for density, rank in DENSITY_RANK.items():
        if density in lower:
            score += rank * 10
    m = re.search(r'(\d{2,4})', lower)
    if m:
        score += min(int(m.group(1)), 2000) // 10
    if lower.rsplit('/', 1)[-1] in ('icon.png', 'logo.png', 'app_icon.png'):
        score += 50
    return score


def candidates_github(owner: str, repo: str, page_url: str = None):
    urls = []
    branch = None
    avatar = None

    try:
        meta = fetch_json(f'https://api.github.com/repos/{owner}/{repo}', headers=GITHUB_API_HEADERS)
        branch = meta.get('default_branch')
        avatar = (meta.get('owner') or {}).get('avatar_url')
    except Exception:
        pass

    branches_to_try = [branch] if branch else ['main', 'master']

    tree_paths = []
    for b in branches_to_try:
        if not b:
            continue
        try:
            tree = fetch_json(
                f'https://api.github.com/repos/{owner}/{repo}/git/trees/{b}?recursive=1',
                headers=GITHUB_API_HEADERS,
            ).get('tree', [])
            tree_paths = [item['path'] for item in tree if item.get('type') == 'blob' and ICON_FILENAME_RE.search(item.get('path', ''))]
            branch = b
            break
        except Exception:
            continue

    if tree_paths:
        tree_paths.sort(key=score_github_path, reverse=True)
        for path in tree_paths[:5]:
            encoded = '/'.join(quote(part) for part in path.split('/'))
            urls.append(f'https://raw.githubusercontent.com/{owner}/{repo}/{branch}/{encoded}')
    else:
        # Trees API unavailable (rate-limited / private repo) — blind-guess the
        # conventional icon locations against each candidate branch.
        for b in branches_to_try:
            if not b:
                continue
            for guess in COMMON_ICON_GUESS_PATHS:
                urls.append(f'https://raw.githubusercontent.com/{owner}/{repo}/{b}/{guess}')

    if avatar:
        sep = '&' if '?' in avatar else '?'
        urls.append(f'{avatar}{sep}s=1024')
    else:
        urls.append(f'https://github.com/{owner}.png?size=1024')

    if page_url:
        try:
            html = fetch_page(page_url)
            og = find_icon_url_from_html(html, source='github', repo_data={'owner': owner})
            if og and og not in urls:
                urls.append(og)
        except Exception:
            pass

    return urls


def candidates_komi(page_url: str):
    html = fetch_page(page_url)

    generic_urls = []
    try:
        generic_urls.append(find_icon_url_from_html(html, source='komi'))
    except Exception:
        pass

    # Komi Store is a discovery layer over GitHub/Codeberg/Forgejo releases and
    # its app pages are rendered client-side, so the raw HTML rarely has a
    # usable icon. Look for the underlying GitHub repo reference and delegate
    # to the (much more reliable) GitHub pipeline.
    github_urls = []
    match = GITHUB_REPO_IN_TEXT_RE.search(html)
    if match:
        owner, repo = match.group(1), re.sub(r'\.git$', '', match.group(2))
        try:
            github_urls = candidates_github(owner, repo)
        except Exception:
            github_urls = []

    urls = github_urls + generic_urls
    if not urls:
        raise RuntimeError(
            'This Komi Store page is rendered client-side, so no icon is exposed to a '
            'plain HTTP fetch and no GitHub/Codeberg repo link was found on the page. '
            'Paste the underlying GitHub repo link instead for a high-quality logo.'
        )
    return urls


# ---------------------------------------------------------------------------
# Step 3: download candidates and keep the highest-quality result
# ---------------------------------------------------------------------------

def fetch_icon_from_source(input_str: str):
    source, data = identify_source(input_str)
    detected = describe_source(source, data)

    if source == 'play':
        urls = candidates_play(data)
    elif source == 'fdroid':
        urls = candidates_fdroid(data)
    elif source == 'github':
        urls = candidates_github(data['owner'], data['repo'], page_url=data['url'])
    elif source == 'komi':
        urls = candidates_komi(data)
    else:
        page_html = fetch_page(input_str)
        urls = [find_icon_url_from_html(page_html)]

    downloaded = []
    for url_try in urls:
        try:
            path = download_image_to_temp(url_try)
        except Exception:
            continue
        downloaded.append((path, image_dimensions(path), url_try))
        # Once we have a few real candidates, compare and stop — no need to
        # hammer every fallback once quality is established.
        if len(downloaded) >= 3:
            break

    if not downloaded:
        return None, None, f'{detected}\nFailed to download a logo from any candidate source.'

    downloaded.sort(key=lambda d: (d[1][0] * d[1][1]) if d[1] else 0, reverse=True)
    best_path, best_dims, best_url = downloaded[0]
    final_path = save_as_logo_png(best_path)

    dims_str = f'{best_dims[0]}x{best_dims[1]}' if best_dims != (0, 0) else 'unknown size'
    status = f'{detected}\nOK — {dims_str} from {urlparse(best_url).netloc}'
    return final_path, final_path, status


def gr_fetch(input_text: str):
    try:
        path_img, path_file, status = fetch_icon_from_source(input_text)
    except Exception as exc:
        return None, None, f'Error: {exc}'
    if path_img is None:
        return None, None, status
    return path_img, path_file, status


if __name__ == '__main__':
    with gr.Blocks(title='Icon Fetcher') as demo:
        gr.Markdown(
            '''# Icon Fetcher\n'''
            '''Paste a Google Play Store, GitHub, F-Droid, or Komi Store app URL (or a Play Store package id / bare `owner/repo`). '''
            '''Click Fetch — the source is analyzed first, then the highest-quality `logo.png` available is downloaded.'''
        )
        with gr.Row():
            inp = gr.Textbox(label='Play Store / GitHub / F-Droid / Komi Store URL or id', placeholder='https://play.google.com/store/apps/details?id=com.example.app')
            btn = gr.Button('Fetch icon')
        with gr.Row():
            img = gr.Image(type='filepath', label='Icon (preview)')
            dl = gr.File(label='Download logo.png')
        status = gr.Textbox(label='Status', interactive=False, lines=2)

        btn.click(fn=gr_fetch, inputs=[inp], outputs=[img, dl, status])

    demo.launch(server_name='0.0.0.0', share=False)