#!/usr/bin/env python3
"""
Yandex Image Search API v2 client.

Adapted from microdrama/scripts/yandex_image_search.py for use in the V5 pipeline.
Requires YANDEX_FOLDER_ID and YANDEX_API_KEY environment variables.
"""

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

import requests

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 get_credentials() -> tuple[str, str]:
    """Get Yandex API credentials from environment."""
    folder_id = os.getenv("YANDEX_FOLDER_ID", "")
    api_key = os.getenv("YANDEX_API_KEY", "")
    return folder_id, api_key


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

    Returns list of dicts: {url, thumb_url, title, width, height}
    """
    size_map = {
        "ENORMOUS": "IMAGE_SIZE_ENORMOUS",
        "LARGE": "IMAGE_SIZE_LARGE",
        "MEDIUM": "IMAGE_SIZE_MEDIUM",
    }
    orient_map = {
        "HORIZONTAL": "IMAGE_ORIENTATION_HORIZONTAL",
        "VERTICAL": "IMAGE_ORIENTATION_VERTICAL",
        "SQUARE": "IMAGE_ORIENTATION_SQUARE",
    }

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

    if size.upper() in size_map:
        body["imageSpec"]["size"] = size_map[size.upper()]
    if orientation and orientation.upper() in orient_map:
        body["imageSpec"]["orientation"] = orient_map[orientation.upper()]

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

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

    if response.status_code != 200:
        print(f"  Yandex API error: {response.status_code} — {response.text[:200]}")
        return []

    result = response.json()

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

    xml_content = base64.b64decode(result["rawData"]).decode("utf-8")
    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)

        error = root.find(".//error")
        if error is not None:
            print(f"  Yandex API error: {error.text}")
            return []

        for group in root.findall(".//group"):
            doc = group.find("doc")
            if doc is None:
                continue

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

            thumb_elem = doc.find(".//thumbnail-url")
            title_elem = doc.find("title")
            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_elem.text,
                "thumb_url": thumb_elem.text if thumb_elem is not None else None,
                "title": "".join(title_elem.itertext()).strip() if title_elem is not None else "",
                "width": width,
                "height": height,
            })

    except ET.ParseError as e:
        print(f"  XML parse error: {e}")

    return results


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

        # Validate it's actually an image (check content-type or minimum size)
        if len(response.content) < 1000:
            return False

        output_path.parent.mkdir(parents=True, exist_ok=True)
        with open(output_path, "wb") as f:
            f.write(response.content)
        return True

    except Exception as e:
        print(f"  Download failed ({url[:80]}): {e}")
        return False


def search_and_download(
    query: str,
    output_dir: Path,
    prefix: str,
    count: int = 8,
    folder_id: str = None,
    api_key: str = None,
) -> List[Path]:
    """
    Search Yandex and download images. Returns list of downloaded file paths.
    """
    if not folder_id or not api_key:
        folder_id, api_key = get_credentials()

    if not folder_id or not api_key:
        print("  Warning: YANDEX_FOLDER_ID / YANDEX_API_KEY not set")
        return []

    results = search_images(query, folder_id, api_key, count=count)
    if not results:
        return []

    output_dir.mkdir(parents=True, exist_ok=True)
    downloaded = []

    for i, r in enumerate(results, 1):
        # Determine extension from URL
        ext = "jpg"
        url_lower = r["url"].lower()
        for e in ["png", "webp"]:
            if e in url_lower:
                ext = e
                break

        filepath = output_dir / f"{prefix}_{i:02d}.{ext}"
        if download_image(r["url"], filepath):
            downloaded.append(filepath)
        time.sleep(0.3)

    return downloaded
