'''FastAPI endpoints for Neural Style Transfer - Standalone version'''

import numpy as np
import tensorflow as tf
import tensorflow_hub as hub
from PIL import Image
import os
import logging
from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.responses import JSONResponse, HTMLResponse
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
import uvicorn
import io
import base64

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

np.set_printoptions(suppress=True)

# Load model
try:
    logger.info("Loading TensorFlow Hub model...")
    model = hub.load('https://tfhub.dev/google/magenta/arbitrary-image-stylization-v1-256/2')
    logger.info("Model loaded successfully!")
except Exception as e:
    logger.error(f"Error loading model: {str(e)}")
    raise

def tensor_to_image(tensor):
    try:
        tensor *= 255
        tensor = np.array(tensor, dtype=np.uint8)
        if tensor.ndim > 3:
            tensor = tensor[0]
        return Image.fromarray(tensor)
    except Exception as e:
        logger.error(f"Error in tensor_to_image: {str(e)}")
        raise

def transform_my_model(content_image, style_image):
    try:
        if content_image is None or style_image is None:
            raise ValueError("Both content and style images are required")
            
        logger.info("Processing images...")
        content_image = content_image.astype(np.float32)[np.newaxis, ...] / 255.0
        style_image = style_image.astype(np.float32)[np.newaxis, ...] / 255.0
        stylized_image = model(tf.constant(content_image), tf.constant(style_image))[0]
        logger.info("Style transfer completed successfully!")
        return tensor_to_image(stylized_image)
    except Exception as e:
        logger.error(f"Error in transform_my_model: {str(e)}")
        raise

# Create FastAPI app
app = FastAPI(title="Neural Style Transfer API")

# Add CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Serve static files
if os.path.exists("static"):
    app.mount("/static", StaticFiles(directory="static"), name="static")

# API endpoint for style transfer
@app.post("/api/style-transfer")
async def style_transfer_api(
    content_image: UploadFile = File(...),
    style_image: UploadFile = File(...)
):
    try:
        # Read images
        content_bytes = await content_image.read()
        style_bytes = await style_image.read()
        
        # Convert to PIL Images
        content_img = Image.open(io.BytesIO(content_bytes)).convert("RGB")
        style_img = Image.open(io.BytesIO(style_bytes)).convert("RGB")
        
        # Convert to numpy arrays
        content_array = np.array(content_img)
        style_array = np.array(style_img)
        
        # Process style transfer
        result_image = transform_my_model(content_array, style_array)
        
        # Convert result to base64
        buffer = io.BytesIO()
        result_image.save(buffer, format="PNG")
        img_str = base64.b64encode(buffer.getvalue()).decode()
        
        return JSONResponse({
            "success": True,
            "image": f"data:image/png;base64,{img_str}"
        })
    except Exception as e:
        logger.error(f"API Error: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

# Endpoint to get available images
@app.get("/api/images")
async def get_images():
    try:
        content_images = []
        style_images = []
        
        if os.path.exists("Content_Images"):
            content_images = sorted([f for f in os.listdir("Content_Images") 
                            if f.lower().endswith(('.png', '.jpg', '.jpeg'))])
        
        if os.path.exists("Style_Images"):
            style_images = sorted([f for f in os.listdir("Style_Images") 
                          if f.lower().endswith(('.png', '.jpg', '.jpeg'))])
        
        return JSONResponse({
            "content_images": content_images,
            "style_images": style_images
        })
    except Exception as e:
        logger.error(f"Error getting images: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

# Serve test.html
@app.get("/", response_class=HTMLResponse)
async def read_root():
    """Serve the test.html interface"""
    if os.path.exists("test.html"):
        with open("test.html", "r", encoding="utf-8") as f:
            return HTMLResponse(content=f.read())
    return HTMLResponse(content="""
    <html>
        <head>
            <title>Neural Style Transfer API</title>
            <style>
                body { font-family: Arial, sans-serif; padding: 40px; background: #f5f5f5; }
                .container { max-width: 800px; margin: 0 auto; background: white; padding: 30px; border-radius: 10px; }
                h1 { color: #333; }
                ul { line-height: 1.8; }
                code { background: #f0f0f0; padding: 2px 6px; border-radius: 3px; }
            </style>
        </head>
        <body>
            <div class="container">
                <h1>Neural Style Transfer API</h1>
                <p>API Endpoints:</p>
                <ul>
                    <li><code>POST /api/style-transfer</code> - Process style transfer</li>
                    <li><code>GET /api/images</code> - Get available images</li>
                </ul>
                <p style="margin-top: 20px; color: #666;">
                    Note: test.html file not found. Please ensure test.html is in the same directory.
                </p>
            </div>
        </body>
    </html>
    """)

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)