# app.py - PyIQA Image Quality Assessment API for Horsh
import gradio as gr
import pyiqa
import torch
from PIL import Image
import requests
from io import BytesIO
import os

# Загрузить модель при старте
print("🚀 Loading ARNIQA model...")
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
metric = pyiqa.create_metric('arniqa', device=device)
print(f"✅ Model loaded on {device}")

def assess_image_quality(image):
    """
    Оценить качество изображения
    Args:
        image: PIL Image
    Returns:
        dict: {'quality': 0.75, 'score': 75, 'status': 'approved'}
    """
    try:
        # Сохранить временно
        temp_path = '/tmp/temp_image.jpg'
        image.save(temp_path)

        # Оценить качество (50-200ms)
        with torch.no_grad():
            score = metric(temp_path).item()

        quality = score / 100.0  # Нормализовать 0-1
        status = 'approved' if quality >= 0.3 else 'rejected'

        return {
            'quality': round(quality, 3),
            'score': round(score, 2),
            'status': status,
            'threshold': 0.3,
            'model': 'ARNIQA (WACV 2024)',
            'device': str(device)
        }

    except Exception as e:
        return {
            'error': str(e),
            'quality': 0.0,
            'score': 0.0,
            'status': 'error'
        }

def assess_from_url(url):
    """Оценить по URL"""
    try:
        if not url or not url.startswith('http'):
            return {'error': 'Invalid URL. Must start with http:// or https://'}

        response = requests.get(url, timeout=15)
        response.raise_for_status()

        img = Image.open(BytesIO(response.content))

        # Конвертировать в RGB если нужно
        if img.mode != 'RGB':
            img = img.convert('RGB')

        return assess_image_quality(img)
    except Exception as e:
        return {
            'error': f'Failed to load image: {str(e)}',
            'quality': 0.0,
            'score': 0.0,
            'status': 'error'
        }

# Gradio Interface
with gr.Blocks(
    title="PyIQA Image Quality Assessment API",
    theme=gr.themes.Soft()
) as demo:

    gr.Markdown("""
    # 📸 Image Quality Assessment API

    **Powered by ARNIQA** (WACV 2024) - State-of-the-art no-reference quality assessment

    This API is used by [Horsh](https://hor.sh) photo-sharing app for automatic quality control.
    """)

    with gr.Tab("🖼️ Upload Image"):
        with gr.Row():
            with gr.Column():
                image_input = gr.Image(type="pil", label="Upload Image")
                upload_btn = gr.Button("Assess Quality", variant="primary", size="lg")

                gr.Markdown("""
                **Threshold:** Quality score ≥ 0.3 = Approved
                **Model:** ARNIQA (Learning Distortion Manifold)
                **Speed:** ~50-200ms per image
                """)

            with gr.Column():
                upload_output = gr.JSON(label="Quality Assessment Result")

        upload_btn.click(
            fn=assess_image_quality,
            inputs=image_input,
            outputs=upload_output
        )

    with gr.Tab("🔗 Image URL"):
        with gr.Row():
            with gr.Column():
                url_input = gr.Textbox(
                    label="Image URL",
                    placeholder="https://example.com/photo.jpg",
                    lines=1
                )
                url_btn = gr.Button("Assess Quality", variant="primary", size="lg")

                gr.Markdown("""
                **Example URLs:**
                - https://picsum.photos/800/600
                - https://images.unsplash.com/photo-1506905925346-21bda4d32df4
                """)

            with gr.Column():
                url_output = gr.JSON(label="Quality Assessment Result")

        url_btn.click(
            fn=assess_from_url,
            inputs=url_input,
            outputs=url_output
        )

    with gr.Tab("📚 API Documentation"):
        gr.Markdown("""
        ## REST API Usage

        This Space exposes a REST API for programmatic access.

        ### Python (gradio_client)

        ```python
        from gradio_client import Client

        client = Client("YOUR_USERNAME/pyiqa-api")

        # Assess from URL
        result = client.predict(
            "https://example.com/photo.jpg",
            api_name="/assess_from_url"
        )
        print(result['quality'])  # 0.756
        ```

        ### Python (requests)

        ```python
        import requests
        import json

        url = "https://YOUR_USERNAME-pyiqa-api.hf.space/api/predict"

        response = requests.post(url, json={
            "data": ["https://example.com/photo.jpg"]
        })

        result = response.json()
        quality = result['data'][0]['quality']
        print(f"Quality: {quality}")
        ```

        ### Flutter/Dart

        ```dart
        import 'package:http/http.dart' as http;
        import 'dart:convert';

        Future<double> assessQuality(String imageUrl) async {
          final response = await http.post(
            Uri.parse('https://YOUR_USERNAME-pyiqa-api.hf.space/api/predict'),
            headers: {'Content-Type': 'application/json'},
            body: jsonEncode({'data': [imageUrl]}),
          );

          if (response.statusCode == 200) {
            final result = jsonDecode(response.body);
            return result['data'][0]['quality'];
          }
          throw Exception('Failed to assess quality');
        }
        ```

        ### cURL

        ```bash
        curl -X POST https://YOUR_USERNAME-pyiqa-api.hf.space/api/predict \\
          -H "Content-Type: application/json" \\
          -d '{"data": ["https://example.com/photo.jpg"]}'
        ```

        ### JavaScript/TypeScript

        ```javascript
        const response = await fetch('https://YOUR_USERNAME-pyiqa-api.hf.space/api/predict', {
          method: 'POST',
          headers: { 'Content-Type': 'application/json' },
          body: JSON.stringify({ data: ['https://example.com/photo.jpg'] })
        });

        const result = await response.json();
        const quality = result.data[0].quality;
        console.log('Quality:', quality);
        ```

        ## Response Format

        ```json
        {
          "quality": 0.756,
          "score": 75.6,
          "status": "approved",
          "threshold": 0.3,
          "model": "ARNIQA (WACV 2024)",
          "device": "cpu"
        }
        ```

        ## Rate Limiting

        - Free tier: No hard limits, but please be reasonable
        - If you need high volume (>10k requests/day), contact us

        ## Model Information

        - **Paper:** [ARNIQA: Learning Distortion Manifold for Image Quality Assessment](https://arxiv.org/abs/2310.14918)
        - **Conference:** WACV 2024 (Oral)
        - **Code:** [miccunifi/ARNIQA](https://github.com/miccunifi/ARNIQA)
        - **PyIQA:** [chaofengc/IQA-PyTorch](https://github.com/chaofengc/IQA-PyTorch)

        ## About Horsh

        This API powers quality control for [Horsh](https://hor.sh) - a photo-sharing app with AI-powered moderation.
        """)

    gr.Markdown("""
    ---
    **Note:** This is a public API. Please use responsibly. For production use, consider running your own instance.

    **License:** Apache-2.0 | **Model:** ARNIQA (WACV 2024) | **Built with:** PyIQA + Gradio
    """)

# Запустить
if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False
    )