"""
image_fetcher.py — Scrapes real images from free public sources (no API keys).

Sources:
  - Wikipedia Commons API  (millions of free labeled images)
  - NASA Image API          (space, science — completely free)
  - Lorem Picsum / Unsplash Source (nature, people, tech)
  - Open Notify / public news RSS image tags

Images are saved to ./image_data/<category>/<filename>.jpg
A manifest file image_manifest.jsonl tracks everything downloaded.
"""

import os
import json
import time
import hashlib
import requests
import threading
from datetime import datetime, UTC
from pathlib import Path
from collections import defaultdict

# ── CONFIG ────────────────────────────────────────────────────────────────────
IMAGE_DIR      = Path('image_data')
MANIFEST_FILE  = 'image_manifest.jsonl'
MAX_PER_CAT    = 500      # max images stored per category
IMG_SIZE       = (224, 224)  # resize target for training

CATEGORIES = [
    'nature', 'technology', 'science', 'people',
    'animals', 'food', 'sports', 'architecture'
]

HEADERS = {'User-Agent': 'LivingNeuralNetwork/1.0 (educational project)'}

# ── WIKIPEDIA COMMONS ─────────────────────────────────────────────────────────
WIKI_SEARCH_TERMS = {
    'nature':       ['forest','ocean','mountain','river','sky','desert','waterfall'],
    'technology':   ['computer','robot','satellite','circuit','smartphone','laboratory'],
    'science':      ['microscope','telescope','dna','chemistry','physics experiment'],
    'people':       ['portrait','crowd','family','athlete','scientist','engineer'],
    'animals':      ['lion','eagle','dolphin','elephant','butterfly','wolf','deer'],
    'food':         ['fruit','vegetable','bread','cooking','meal','market food'],
    'sports':       ['football','basketball','swimming','running','cycling','tennis'],
    'architecture': ['bridge','skyscraper','cathedral','temple','stadium','museum'],
}

def fetch_wikipedia_images(category: str, term: str, max_imgs: int = 5) -> list:
    """Fetch image URLs from Wikipedia Commons for a search term."""
    results = []
    try:
        url = 'https://en.wikipedia.org/w/api.php'
        params = {
            'action':      'query',
            'generator':   'search',
            'gsrsearch':   f'file:{term}',
            'gsrnamespace': 6,
            'prop':        'imageinfo',
            'iiprop':      'url|size|mime',
            'iiurlwidth':  224,
            'gsrlimit':    max_imgs,
            'format':      'json',
        }
        r = requests.get(url, params=params, headers=HEADERS, timeout=10)
        r.raise_for_status()
        pages = r.json().get('query', {}).get('pages', {})
        for page in pages.values():
            info = page.get('imageinfo', [{}])[0]
            thumb_url = info.get('thumburl') or info.get('url', '')
            mime      = info.get('mime', '')
            if thumb_url and 'image' in mime:
                results.append({
                    'url':      thumb_url,
                    'category': category,
                    'source':   'wikipedia_commons',
                    'term':     term,
                })
    except Exception:
        pass
    return results


def fetch_nasa_images(max_imgs: int = 10) -> list:
    """Fetch images from NASA's free image API."""
    results = []
    nasa_terms = ['nebula', 'galaxy', 'astronaut', 'rocket', 'mars', 'earth from space']
    try:
        term = nasa_terms[int(time.time()) % len(nasa_terms)]
        r = requests.get(
            'https://images-api.nasa.gov/search',
            params={'q': term, 'media_type': 'image', 'page_size': max_imgs},
            headers=HEADERS, timeout=10
        )
        r.raise_for_status()
        items = r.json().get('collection', {}).get('items', [])
        for item in items:
            links = item.get('links', [])
            for link in links:
                if link.get('rel') == 'preview':
                    results.append({
                        'url':      link['href'],
                        'category': 'science',
                        'source':   'nasa',
                        'term':     term,
                    })
                    break
    except Exception:
        pass
    return results


def fetch_openverse_images(category: str, term: str, max_imgs: int = 5) -> list:
    """
    Openverse is a Creative Commons image search API — completely free, no key needed.
    https://api.openverse.org
    """
    results = []
    try:
        r = requests.get(
            'https://api.openverse.org/v1/images/',
            params={'q': term, 'page_size': max_imgs, 'license_type': 'commercial,modification'},
            headers=HEADERS, timeout=10
        )
        r.raise_for_status()
        for img in r.json().get('results', []):
            url = img.get('url', '')
            if url:
                results.append({
                    'url':      url,
                    'category': category,
                    'source':   'openverse',
                    'term':     term,
                    'title':    img.get('title', ''),
                })
    except Exception:
        pass
    return results


# ── DOWNLOADER ────────────────────────────────────────────────────────────────
def download_image(url: str, save_path: Path, timeout: int = 12) -> bool:
    """Download and save an image. Returns True on success."""
    try:
        r = requests.get(url, headers=HEADERS, timeout=timeout, stream=True)
        r.raise_for_status()
        content_type = r.headers.get('content-type', '')
        if 'image' not in content_type:
            return False
        save_path.parent.mkdir(parents=True, exist_ok=True)
        with open(save_path, 'wb') as f:
            for chunk in r.iter_content(8192):
                f.write(chunk)
        # Verify it's a valid image using basic header check
        with open(save_path, 'rb') as f:
            header = f.read(12)
        # JPEG, PNG, GIF, WEBP checks
        valid = (
            header[:2]  == b'\xff\xd8' or   # JPEG
            header[:8]  == b'\x89PNG\r\n\x1a\n' or  # PNG
            header[:6]  in (b'GIF87a', b'GIF89a') or  # GIF
            header[:4]  == b'RIFF'  # WEBP
        )
        if not valid:
            save_path.unlink(missing_ok=True)
            return False
        return True
    except Exception:
        try:
            save_path.unlink(missing_ok=True)
        except Exception:
            pass
        return False


def write_manifest(record: dict):
    """Append one image record to the manifest file."""
    try:
        with open(MANIFEST_FILE, 'a', encoding='utf-8') as f:
            f.write(json.dumps(record, ensure_ascii=False) + '\n')
    except Exception:
        pass


def get_manifest_stats() -> dict:
    """Count images per category from manifest."""
    counts = defaultdict(int)
    total  = 0
    if not os.path.exists(MANIFEST_FILE):
        return {'total': 0, 'by_category': {}}
    try:
        with open(MANIFEST_FILE, 'r') as f:
            for line in f:
                line = line.strip()
                if line:
                    try:
                        rec = json.loads(line)
                        counts[rec.get('category', 'unknown')] += 1
                        total += 1
                    except Exception:
                        pass
    except Exception:
        pass
    return {'total': total, 'by_category': dict(counts)}


def get_category_image_count(category: str) -> int:
    cat_dir = IMAGE_DIR / category
    if not cat_dir.exists():
        return 0
    return len(list(cat_dir.glob('*.jpg')) + list(cat_dir.glob('*.png')) +
               list(cat_dir.glob('*.webp')))


# ── MAIN FETCHER CLASS ────────────────────────────────────────────────────────
class ImageFetcher:
    def __init__(self):
        self.total_downloaded = 0
        self.total_failed     = 0
        self.log              = []
        self.category_counts  = defaultdict(int)
        IMAGE_DIR.mkdir(exist_ok=True)
        for cat in CATEGORIES:
            (IMAGE_DIR / cat).mkdir(exist_ok=True)
        # Load existing counts from manifest
        stats = get_manifest_stats()
        self.total_downloaded = stats['total']
        for cat, count in stats['by_category'].items():
            self.category_counts[cat] = count

    def _log(self, msg: str):
        ts = datetime.now(UTC).strftime('%H:%M:%S')
        entry = f"[{ts}] {msg}"
        self.log.append(entry)
        if len(self.log) > 200:
            self.log = self.log[-200:]
        return entry

    def fetch_round(self) -> dict:
        """
        Fetch one round of images across all categories.
        Returns summary of what was downloaded.
        """
        downloaded = 0
        failed     = 0
        candidates = []

        # Wikipedia Commons — most reliable
        import random
        for cat, terms in WIKI_SEARCH_TERMS.items():
            if get_category_image_count(cat) >= MAX_PER_CAT:
                continue
            term = random.choice(terms)
            imgs = fetch_wikipedia_images(cat, term, max_imgs=4)
            candidates.extend(imgs)

        # NASA (science category)
        if get_category_image_count('science') < MAX_PER_CAT:
            candidates.extend(fetch_nasa_images(6))

        # Openverse — random category
        cat  = random.choice(CATEGORIES)
        term = random.choice(WIKI_SEARCH_TERMS.get(cat, [cat]))
        candidates.extend(fetch_openverse_images(cat, term, 5))

        # Download each candidate
        for img_info in candidates:
            cat = img_info['category']
            url = img_info['url']

            # Make a unique filename from URL hash
            url_hash  = hashlib.md5(url.encode()).hexdigest()[:12]
            ext       = url.split('?')[0].split('.')[-1].lower()
            if ext not in ('jpg', 'jpeg', 'png', 'gif', 'webp'):
                ext = 'jpg'
            filename  = f"{url_hash}.{ext}"
            save_path = IMAGE_DIR / cat / filename

            if save_path.exists():
                continue  # Already have it

            if download_image(url, save_path):
                record = {
                    **img_info,
                    'filename':  str(save_path),
                    'timestamp': datetime.now(UTC).isoformat(),
                }
                write_manifest(record)
                self.category_counts[cat] += 1
                self.total_downloaded += 1
                downloaded += 1
            else:
                failed += 1
                self.total_failed += 1

            time.sleep(0.15)  # polite rate limiting

        self._log(
            f"Round complete: +{downloaded} images saved, "
            f"{failed} failed | Total: {self.total_downloaded}"
        )
        return {'downloaded': downloaded, 'failed': failed,
                'total': self.total_downloaded}

    def get_stats(self) -> dict:
        by_cat = {
            cat: get_category_image_count(cat)
            for cat in CATEGORIES
        }
        disk_mb = 0
        try:
            for p in IMAGE_DIR.rglob('*'):
                if p.is_file():
                    disk_mb += p.stat().st_size
            disk_mb = round(disk_mb / 1_000_000, 1)
        except Exception:
            pass
        return {
            'total':      self.total_downloaded,
            'by_category': by_cat,
            'disk_mb':    disk_mb,
            'recent_log': self.log[-15:],
        }

    def get_sample_paths(self, category: str, n: int = 9) -> list:
        """Return paths of n sample images from a category."""
        cat_dir = IMAGE_DIR / category
        if not cat_dir.exists():
            return []
        files = list(cat_dir.iterdir())
        import random
        sample = random.sample(files, min(n, len(files)))
        return [str(p) for p in sample]