#!/usr/bin/env python3
"""
Yandex Image Search API v2 client for finding reference images.

Uses Yandex Cloud Search API v2 (REST).
Requires YANDEX_FOLDER_ID and YANDEX_API_KEY environment variables.

Example usage:
    python yandex_image_search.py "историк архивист портрет" --output references/search/ --download
"""

import os
import sys
import argparse
import requests
import xml.etree.ElementTree as ET
import base64
import json
from pathlib import Path
from typing import List, Dict, Optional
import time

# Load .env from parent directory
from dotenv import load_dotenv
env_path = Path(__file__).parent.parent.parent / ".env"
load_dotenv(env_path)


API_URL = "https://searchapi.api.cloud.yandex.net/v2/image/search"
USER_AGENT = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36"


def search_images(
    query: str,
    folder_id: str,
    api_key: str,
    count: int = 10,
    page: int = 0,
    size: str = "LARGE",
    orientation: Optional[str] = None,
    format: str = "JPEG",
) -> List[Dict]:
    """
    Search Yandex Images using API v2.

    Args:
        query: Search query in Russian
        folder_id: Yandex Cloud folder ID
        api_key: Yandex Cloud API key
        count: Number of results per page (max 100)
        page: Page number (0-indexed)
        size: Image size (ENORMOUS|LARGE|MEDIUM|SMALL|WALLPAPER)
        orientation: Orientation (HORIZONTAL|VERTICAL|SQUARE)
        format: Image format (JPEG|GIF|PNG)

    Returns:
        List of dicts with keys: url, thumb_url, title, width, height
    """
    # Map CLI values to API enum values
    size_map = {
        "ENORMOUS": "IMAGE_SIZE_ENORMOUS",
        "LARGE": "IMAGE_SIZE_LARGE",
        "MEDIUM": "IMAGE_SIZE_MEDIUM",
        "SMALL": "IMAGE_SIZE_SMALL",
        "WALLPAPER": "IMAGE_SIZE_WALLPAPER",
    }
    orient_map = {
        "HORIZONTAL": "IMAGE_ORIENTATION_HORIZONTAL",
        "VERTICAL": "IMAGE_ORIENTATION_VERTICAL",
        "SQUARE": "IMAGE_ORIENTATION_SQUARE",
    }
    format_map = {
        "JPEG": "IMAGE_FORMAT_JPEG",
        "GIF": "IMAGE_FORMAT_GIF",
        "PNG": "IMAGE_FORMAT_PNG",
    }

    # Build request body
    body = {
        "query": {
            "searchType": "SEARCH_TYPE_RU",
            "queryText": query,
            "familyMode": "FAMILY_MODE_STRICT",
            "page": str(page),
        },
        "imageSpec": {},
        "docsOnPage": str(min(count, 100)),
        "folderId": folder_id,
        "userAgent": USER_AGENT,
    }

    # Add optional image spec fields
    if size.upper() in size_map:
        body["imageSpec"]["size"] = size_map[size.upper()]
    if format.upper() in format_map:
        body["imageSpec"]["format"] = format_map[format.upper()]
    if orientation and orientation.upper() in orient_map:
        body["imageSpec"]["orientation"] = orient_map[orientation.upper()]

    # Make request
    headers = {
        "Authorization": f"Api-Key {api_key}",
        "Content-Type": "application/json",
    }

    print(f"Requesting: {API_URL}")
    print(f"Query: {query}")

    response = requests.post(API_URL, headers=headers, json=body, timeout=30)

    if response.status_code != 200:
        print(f"Error: {response.status_code}")
        print(f"Response: {response.text}")
        return []

    # Parse response
    result = response.json()

    if "rawData" not in result:
        print(f"No rawData in response: {result}")
        return []

    # Decode base64 XML
    xml_content = base64.b64decode(result["rawData"]).decode("utf-8")

    # Parse XML
    return parse_xml_response(xml_content)


def parse_xml_response(xml_content: str) -> List[Dict]:
    """Parse Yandex XML response into result dicts."""
    results = []

    try:
        root = ET.fromstring(xml_content)

        # Check for errors
        error = root.find(".//error")
        if error is not None:
            print(f"API Error: {error.text}", file=sys.stderr)
            return []

        # Find all document groups
        for group in root.findall(".//group"):
            doc = group.find("doc")
            if doc is None:
                continue

            # Extract image URL
            url_elem = doc.find("url")
            if url_elem is None or not url_elem.text:
                continue
            url = url_elem.text

            # Extract thumbnail
            thumb_elem = doc.find(".//thumbnail-url")
            thumb_url = thumb_elem.text if thumb_elem is not None else None

            # Extract title
            title_elem = doc.find("title")
            title = "".join(title_elem.itertext()).strip() if title_elem is not None else ""

            # Extract dimensions
            props = doc.find(".//image-properties")
            width = height = 0
            if props is not None:
                w = props.find("width")
                h = props.find("height")
                width = int(w.text) if w is not None and w.text else 0
                height = int(h.text) if h is not None and h.text else 0

            results.append({
                "url": url,
                "thumb_url": thumb_url,
                "title": title,
                "width": width,
                "height": height,
            })

    except ET.ParseError as e:
        print(f"XML parse error: {e}", file=sys.stderr)
        print(f"XML preview: {xml_content[:500]}", file=sys.stderr)

    return results


def download_image(url: str, output_path: Path, timeout: int = 30) -> bool:
    """Download image from URL."""
    try:
        headers = {"User-Agent": USER_AGENT}
        response = requests.get(url, headers=headers, timeout=timeout, allow_redirects=True)
        response.raise_for_status()

        with open(output_path, "wb") as f:
            f.write(response.content)
        return True

    except Exception as e:
        print(f"Failed to download {url}: {e}", file=sys.stderr)
        return False


def main():
    parser = argparse.ArgumentParser(description="Search Yandex Images for references")
    parser.add_argument("query", help="Search query (in Russian)")
    parser.add_argument("--output", "-o", help="Output directory", default=".")
    parser.add_argument("--count", "-n", type=int, default=5, help="Number of images")
    parser.add_argument("--size", choices=["enormous", "large", "medium", "small", "wallpaper"], default="large")
    parser.add_argument("--orient", choices=["horizontal", "vertical", "square"])
    parser.add_argument("--format", choices=["jpeg", "gif", "png"], default="jpeg")
    parser.add_argument("--download", "-d", action="store_true", help="Download images")
    parser.add_argument("--prefix", "-p", default="image", help="Filename prefix")

    args = parser.parse_args()

    # Get credentials from environment
    folder_id = os.getenv("YANDEX_FOLDER_ID")
    api_key = os.getenv("YANDEX_API_KEY")

    if not folder_id or not api_key:
        print("Error: Set YANDEX_FOLDER_ID and YANDEX_API_KEY environment variables")
        sys.exit(1)

    print(f"Searching: {args.query}\n")

    # Search
    results = search_images(
        query=args.query,
        folder_id=folder_id,
        api_key=api_key,
        count=args.count,
        size=args.size.upper(),
        orientation=args.orient.upper() if args.orient else None,
        format=args.format.upper(),
    )

    print(f"\nFound {len(results)} results\n")

    # Output directory
    output_dir = Path(args.output)
    output_dir.mkdir(parents=True, exist_ok=True)

    # Save query
    query_file = output_dir / f"query_{args.prefix}.txt"
    with open(query_file, "w") as f:
        f.write(f"Query: {args.query}\n")
        f.write(f"Size: {args.size}\n")
        f.write(f"Orient: {args.orient or 'any'}\n")
        f.write(f"Results: {len(results)}\n\n")
        for i, r in enumerate(results, 1):
            f.write(f"{i}. {r['url']}\n")
            f.write(f"   Title: {r['title']}\n")
            f.write(f"   Size: {r['width']}x{r['height']}\n\n")

    print(f"Query saved to: {query_file}")

    # Print results
    for i, r in enumerate(results, 1):
        title_preview = r['title'][:60] + "..." if len(r['title']) > 60 else r['title']
        print(f"{i}. {title_preview}")
        print(f"   URL: {r['url']}")
        print(f"   Size: {r['width']}x{r['height']}")
        print()

    # Download if requested
    if args.download and results:
        print("Downloading images...")
        for i, r in enumerate(results, 1):
            # Determine extension
            ext = "jpg"
            url_lower = r['url'].lower()
            for e in ['png', 'gif', 'webp']:
                if e in url_lower:
                    ext = e
                    break

            filename = f"{args.prefix}_{i:02d}.{ext}"
            filepath = output_dir / filename

            if download_image(r['url'], filepath):
                print(f"  Downloaded: {filename}")

            time.sleep(0.5)  # Be nice to servers


if __name__ == "__main__":
    main()
