from flask import Flask, request, jsonify
from flask_cors import CORS
import requests
import time
import json
from typing import Optional, Dict, Any
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = Flask(__name__)
CORS(app)

# Configuration
GRADIO_API_URL = "https://mohamedislegend4-z-image-turbo-api.hf.space"
MAX_POLL_ATTEMPTS = 120  # 2 minutes with 1 second intervals
POLL_INTERVAL = 1  # seconds

class GradioAPIClient:
    """Client for interacting with Gradio APIs"""
    
    def __init__(self, base_url: str):
        self.base_url = base_url.rstrip('/')
        self.session = requests.Session()
        self.session.timeout = 30
    
    def call_endpoint(self, endpoint: str, params: Dict[str, Any]) -> Dict[str, Any]:
        """
        Call a Gradio endpoint and get the event_id
        
        Args:
            endpoint: The endpoint name (e.g., "generate")
            params: Dictionary of parameters for the endpoint
        
        Returns:
            Dictionary with event_id and data
        """
        url = f"{self.base_url}/gradio_api/call/v2/{endpoint}"
        
        logger.info(f"Calling endpoint: {endpoint}")
        logger.info(f"Parameters: {params}")
        
        try:
            response = self.session.post(
                url,
                json=params,
                headers={"Content-Type": "application/json"}
            )
            response.raise_for_status()
            
            data = response.json()
            logger.info(f"Response: {data}")
            
            return data
        except requests.exceptions.RequestException as e:
            logger.error(f"Error calling endpoint: {e}")
            raise
    
    def poll_result(self, endpoint: str, event_id: str) -> Optional[Dict[str, Any]]:
        """
        Poll the result of an async call
        
        Args:
            endpoint: The endpoint name
            event_id: The event ID from the initial call
        
        Returns:
            The result data when ready, or None if timeout
        """
        url = f"{self.base_url}/gradio_api/call/{endpoint}/{event_id}"
        
        logger.info(f"Polling result for {endpoint}/{event_id}")
        
        for attempt in range(MAX_POLL_ATTEMPTS):
            try:
                response = self.session.get(url, stream=True)
                
                # Read all content
                content = b""
                for chunk in response.iter_content(chunk_size=1024):
                    if chunk:
                        content += chunk
                
                # Try to parse JSON
                try:
                    data = json.loads(content.decode('utf-8'))
                    
                    # Check if we have a result
                    if "result" in data and data["result"] is not None:
                        logger.info(f"Got result on attempt {attempt + 1}")
                        return data["result"]
                    
                    logger.info(f"Attempt {attempt + 1}: No result yet, waiting...")
                    time.sleep(POLL_INTERVAL)
                    
                except json.JSONDecodeError as e:
                    logger.warning(f"Failed to parse JSON: {e}")
                    logger.debug(f"Content: {content[:500]}")
                    time.sleep(POLL_INTERVAL)
                    
            except requests.exceptions.RequestException as e:
                logger.warning(f"Poll request failed: {e}")
                time.sleep(POLL_INTERVAL)
        
        logger.error(f"Polling timeout after {MAX_POLL_ATTEMPTS} attempts")
        return None
    
    def generate_image(self, prompt: str, steps: int = 20, 
                      height: int = 512, width: int = 512) -> Optional[Dict[str, Any]]:
        """
        Generate an image using the Z-Image-Turbo API
        
        Args:
            prompt: Text description of the image
            steps: Number of inference steps (1-50)
            height: Image height in pixels
            width: Image width in pixels
        
        Returns:
            Dictionary with image data including URL, or None on error
        """
        params = {
            "prompt": prompt,
            "steps": steps,
            "height": height,
            "width": width
        }
        
        # Step 1: Call the endpoint
        try:
            response = self.call_endpoint("generate", params)
            event_id = response.get("event_id")
            
            if not event_id:
                logger.error("No event_id in response")
                return None
            
            # Step 2: Poll for result
            result = self.poll_result("generate", event_id)
            
            if result is None:
                logger.error("Failed to get result")
                return None
            
            # Result should be a list with image data
            if isinstance(result, list) and len(result) > 0:
                image_data = result[0]
                logger.info(f"Image data: {image_data}")
                return image_data
            
            logger.error("Invalid result format")
            return None
            
        except Exception as e:
            logger.error(f"Error generating image: {e}")
            return None


# Initialize client
client = GradioAPIClient(GRADIO_API_URL)


@app.route('/health', methods=['GET'])
def health():
    """Health check endpoint"""
    return jsonify({"status": "ok"})


@app.route('/api/generate', methods=['POST'])
def generate():
    """
    Generate an image via API
    
    Request body:
    {
        "prompt": "A beautiful sunset over mountains",
        "steps": 20,
        "height": 512,
        "width": 512
    }
    
    Response:
    {
        "success": true,
        "image_url": "https://...",
        "image_path": "...",
        "prompt": "...",
        "size": 12345,
        "mime_type": "image/png"
    }
    """
    try:
        data = request.get_json()
        
        # Validate input
        if not data or "prompt" not in data:
            return jsonify({
                "success": False,
                "error": "Missing required parameter: 'prompt'"
            }), 400
        
        prompt = data.get("prompt")
        steps = min(max(int(data.get("steps", 20)), 1), 50)  # Clamp 1-50
        height = int(data.get("height", 512))
        width = int(data.get("width", 512))
        
        logger.info(f"Generating image for prompt: {prompt}")
        
        # Call the API
        image_data = client.generate_image(prompt, steps, height, width)
        
        if not image_data:
            return jsonify({
                "success": False,
                "error": "Failed to generate image"
            }), 500
        
        # Extract relevant fields
        result = {
            "success": True,
            "prompt": prompt,
            "steps": steps,
            "height": height,
            "width": width,
        }
        
        # Add image URL if available
        if isinstance(image_data, dict):
            if "url" in image_data:
                result["image_url"] = image_data["url"]
            if "path" in image_data:
                result["image_path"] = image_data["path"]
            if "size" in image_data:
                result["size"] = image_data["size"]
            if "mime_type" in image_data:
                result["mime_type"] = image_data["mime_type"]
            if "orig_name" in image_data:
                result["filename"] = image_data["orig_name"]
        
        return jsonify(result), 200
        
    except Exception as e:
        logger.error(f"Error in /api/generate: {e}", exc_info=True)
        return jsonify({
            "success": False,
            "error": str(e)
        }), 500


@app.route('/api/generate-sync', methods=['POST'])
def generate_sync():
    """
    Synchronous image generation endpoint
    Same as /api/generate but with better error handling
    """
    return generate()


@app.route('/api/info', methods=['GET'])
def get_info():
    """Get API info and available parameters"""
    return jsonify({
        "api_name": "Z-Image-Turbo API Wrapper",
        "version": "1.0.0",
        "endpoints": [
            {
                "path": "/api/generate",
                "method": "POST",
                "description": "Generate an image from a text prompt",
                "parameters": {
                    "prompt": {
                        "type": "string",
                        "required": True,
                        "description": "Text description of the image"
                    },
                    "steps": {
                        "type": "integer",
                        "required": False,
                        "default": 20,
                        "min": 1,
                        "max": 50,
                        "description": "Number of inference steps"
                    },
                    "height": {
                        "type": "integer",
                        "required": False,
                        "default": 512,
                        "description": "Image height in pixels"
                    },
                    "width": {
                        "type": "integer",
                        "required": False,
                        "default": 512,
                        "description": "Image width in pixels"
                    }
                }
            }
        ]
    })


@app.errorhandler(404)
def not_found(e):
    return jsonify({"error": "Not found"}), 404


@app.errorhandler(500)
def server_error(e):
    return jsonify({"error": "Internal server error"}), 500


if __name__ == '__main__':
    print("=" * 60)
    print("Z-Image-Turbo API Wrapper")
    print("=" * 60)
    print(f"Gradio API URL: {GRADIO_API_URL}")
    print("Starting Flask server...")
    print("=" * 60)
    app.run(host='0.0.0.0', port=5000, debug=True)