import os
import cv2
import torch
from flask import Flask, request, jsonify, send_file
from basicsr.archs.srvgg_arch import SRVGGNetCompact
from gfpgan.utils import GFPGANer
from realesrgan.utils import RealESRGANer
import uuid
import tempfile
# ウェイトファイルをダウンロード（存在しない場合）
if not os.path.exists('realesr-general-x4v3.pth'):
    os.system("wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth -P .")
if not os.path.exists('GFPGANv1.2.pth'):
    os.system("wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.2.pth -P .")
if not os.path.exists('GFPGANv1.3.pth'):
    os.system("wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth -P .")
if not os.path.exists('GFPGANv1.4.pth'):
    os.system("wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth -P .")
if not os.path.exists('RestoreFormer.pth'):
    os.system("wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.4/RestoreFormer.pth -P .")
if not os.path.exists('CodeFormer.pth'):
    os.system("wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.4/CodeFormer.pth -P .")

app = Flask(__name__)

# モデルの初期化
model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu')
model_path = 'realesr-general-x4v3.pth'
half = True if torch.cuda.is_available() else False
upsampler = RealESRGANer(scale=4, model_path=model_path, model=model, tile=0, tile_pad=10, pre_pad=0, half=half)

os.makedirs('output', exist_ok=True)

@app.route('/api/restore', methods=['POST'])
def restore_image():
    try:
        # リクエストからパラメータを取得
        if 'file' not in request.files:
            return jsonify({'error': 'No file uploaded'}), 400
        
        file = request.files['file']
        version = request.form.get('version', 'v1.4')
        scale = float(request.form.get('scale', 2))
        # weight = float(request.form.get('weight', 50)) / 100  # CodeFormer用のweightパラメータが必要な場合
        
        # 一時ファイルに保存
        temp_dir = tempfile.mkdtemp()
        input_path = os.path.join(temp_dir, file.filename)
        file.save(input_path)
        
        # 画像処理
        extension = os.path.splitext(os.path.basename(str(input_path)))[1]
        img = cv2.imread(input_path, cv2.IMREAD_UNCHANGED)
        
        if len(img.shape) == 3 and img.shape[2] == 4:
            img_mode = 'RGBA'
        elif len(img.shape) == 2:
            img_mode = None
            img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
        else:
            img_mode = None

        h, w = img.shape[0:2]
        if h < 300:
            img = cv2.resize(img, (w * 2, h * 2), interpolation=cv2.INTER_LANCZOS4)

        # バージョンに応じてモデルを選択
        if version == 'v1.2':
            face_enhancer = GFPGANer(
                model_path='GFPGANv1.2.pth', upscale=2, arch='clean', channel_multiplier=2, bg_upsampler=upsampler)
        elif version == 'v1.3':
            face_enhancer = GFPGANer(
                model_path='GFPGANv1.3.pth', upscale=2, arch='clean', channel_multiplier=2, bg_upsampler=upsampler)
        elif version == 'v1.4':
            face_enhancer = GFPGANer(
                model_path='GFPGANv1.4.pth', upscale=2, arch='clean', channel_multiplier=2, bg_upsampler=upsampler)
        elif version == 'RestoreFormer':
            face_enhancer = GFPGANer(
                model_path='RestoreFormer.pth', upscale=2, arch='RestoreFormer', channel_multiplier=2, bg_upsampler=upsampler)
        elif version == 'CodeFormer':
            face_enhancer = GFPGANer(
                model_path='CodeFormer.pth', upscale=2, arch='CodeFormer', channel_multiplier=2, bg_upsampler=upsampler)
        elif version == 'RealESR-General-x4v3':
            face_enhancer = GFPGANer(
                model_path='realesr-general-x4v3.pth', upscale=2, arch='realesr-general', channel_multiplier=2, bg_upsampler=upsampler)

        # 画像を拡張
        _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True)
        
        # スケール調整
        if scale != 2:
            interpolation = cv2.INTER_AREA if scale < 2 else cv2.INTER_LANCZOS4
            h, w = img.shape[0:2]
            output = cv2.resize(output, (int(w * scale / 2), int(h * scale / 2)), interpolation=interpolation)
        
        # 出力ファイルを保存
        output_filename = f'output_{uuid.uuid4().hex}'
        if img_mode == 'RGBA':
            output_path = os.path.join('output', f'{output_filename}.png')
            cv2.imwrite(output_path, output)
            mimetype = 'image/png'
        else:
            output_path = os.path.join('output', f'{output_filename}.jpg')
            cv2.imwrite(output_path, output)
            mimetype = 'image/jpeg'
        
        # 結果を返す
        return send_file(output_path, mimetype=mimetype, as_attachment=True, download_name=os.path.basename(output_path))
    
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/')
def index():
    return """
<!DOCTYPE html>
<html>
<head>
    <title>Image Upscaling & Restoration API</title>
    <style>
        body { font-family: Arial, sans-serif; max-width: 800px; margin: 0 auto; padding: 20px; }
        .container { border: 1px solid #ddd; padding: 20px; border-radius: 5px; }
        .form-group { margin-bottom: 15px; }
        label { display: block; margin-bottom: 5px; }
        input, select { width: 100%; padding: 8px; box-sizing: border-box; }
        button { background-color: #4CAF50; color: white; padding: 10px 15px; border: none; border-radius: 4px; cursor: pointer; }
        button:hover { background-color: #45a049; }
        #result { margin-top: 20px; }
        #preview { max-width: 100%; margin-top: 10px; }
        #apiUsage { background-color: #f5f5f5; padding: 15px; border-radius: 5px; margin-top: 20px; font-family: monospace; white-space: pre-wrap; }
        #apiUsage h3 { margin-top: 0; }
        #formDataPreview { max-height: 200px; overflow-y: auto; margin-bottom: 10px; }
        .code-block { background-color: #f8f8f8; padding: 10px; border-radius: 4px; border-left: 3px solid #4CAF50; }
        .comment { color: #666; font-style: italic; }
        
        .loader {
            width: 48px;
            height: 48px;
            border: 5px solid #4CAF50;
            border-bottom-color: transparent;
            border-radius: 50%;
            display: inline-block;
            box-sizing: border-box;
            animation: rotation 1s linear infinite;
            margin: 20px auto;
            display: none; /* 初期状態では非表示 */
        }

        @keyframes rotation {
            0% {
                transform: rotate(0deg);
            }
            100% {
                transform: rotate(360deg);
            }
        }
    </style>
</head>
<body>
    <h1>Image Upscaling & Restoration API</h1>
    <div class="container">
        <form id="uploadForm" enctype="multipart/form-data">
            <div class="form-group">
                <label for="file">Upload Image:</label>
                <input type="file" id="file" name="file" required>
            </div>
            <div class="form-group">
                <label for="version">Version:</label>
                <select id="version" name="version">
                    <option value="v1.2">GFPGANv1.2</option>
                    <option value="v1.3">GFPGANv1.3</option>
                    <option value="v1.4" selected>GFPGANv1.4</option>
                    <option value="RestoreFormer">RestoreFormer</option>
                    <option value="CodeFormer">CodeFormer</option>
                    <option value="RealESR-General-x4v3">RealESR-General-x4v3</option>
                </select>
            </div>
            <div class="form-group">
                <label for="scale">Rescaling factor:</label>
                <input type="number" id="scale" name="scale" value="2" step="0.1" min="1" max="4" required>
            </div>
            <button type="submit" id="submitButton">Process Image</button>
        </form>
        <div id="loading" class="loader"></div>
        <div id="result">
            <h3>Result:</h3>
            <div id="outputContainer" style="display: none;">
                <img id="preview" src="" alt="Processed Image">
                <a id="downloadLink" href="#" download>Download Image</a>
            </div>
        </div>
        <div id="apiUsage">
            <h3>API Usage:</h3>
            <div id="fetchCode" class="code-block">
                // JavaScript fetch code will appear here
            </div>
        </div>
    </div>
    
    <script>
        // フォームの変更を監視してAPI使用例を更新
        function updateApiUsage() {
            const fileInput = document.getElementById('file');
            const version = document.getElementById('version').value;
            const scale = document.getElementById('scale').value;
            
            // 現在のURLからベースURLを取得（パス、パラメータ、ハッシュを含めない）
            const baseUrl = window.location.origin;
            const apiUrl = baseUrl + '/api/restore';
            
            // ファイルのプレビュー用文字列を準備
            let filePreview = '"img-dataURL"';
            if (fileInput.files.length > 0) {
                const file = fileInput.files[0];
                const reader = new FileReader();
                reader.onload = function(e) {
                    const dataURL = e.target.result;
                    if (dataURL.length > 40) {
                        filePreview = `"${dataURL.substring(0, 20)}...${dataURL.substring(dataURL.length - 20)}"`;
                    } else {
                        filePreview = `"${dataURL}"`;
                    }
                    updateFetchCode(apiUrl, version, scale, filePreview);
                };
                reader.readAsDataURL(file);
            } else {
                updateFetchCode(apiUrl, version, scale, filePreview);
            }
        }
        
        function updateFetchCode(apiUrl, version, scale, filePreview) {
            const fetchCodeDiv = document.getElementById('fetchCode');
            fetchCodeDiv.innerHTML = `
// JavaScript fetch example:
const formData = new FormData();
formData.append('file', ${filePreview});
formData.append('version', '${version}');
formData.append('scale', ${scale});

fetch('${apiUrl}', {
    method: 'POST',
    body: formData
})
.then(response => {
    if (!response.ok) {
        return response.json().then(err => { throw new Error(err.error || 'Unknown error'); });
    }
    return response.blob();
})
.then(blob => {
    // Process the returned image blob
    const url = URL.createObjectURL(blob);
    console.log('Image processed successfully', url);
    // Example: document.getElementById('resultImage').src = url;
})
.catch(error => {
    console.error('Error:', error.message);
});`;
        }
        
        // フォーム要素の変更を監視
        document.getElementById('file').addEventListener('change', updateApiUsage);
        document.getElementById('version').addEventListener('change', updateApiUsage);
        document.getElementById('scale').addEventListener('input', updateApiUsage);
        
        // 初期表示
        updateApiUsage();
        
        document.getElementById('uploadForm').addEventListener('submit', function(e) {
            e.preventDefault();
            
            // ボタンを無効化し、ローディングを表示
            const submitButton = document.getElementById('submitButton');
            const loadingElement = document.getElementById('loading');
            
            submitButton.disabled = true;
            loadingElement.style.display = 'block';
            
            const formData = new FormData();
            formData.append('file', document.getElementById('file').files[0]);
            formData.append('version', document.getElementById('version').value);
            formData.append('scale', document.getElementById('scale').value);
            
            // 現在のURLからベースURLを取得（パス、パラメータ、ハッシュを含めない）
            const baseUrl = window.location.origin;
            const apiUrl = baseUrl + '/api/restore';
            
            fetch(apiUrl, {
                method: 'POST',
                body: formData
            })
            .then(response => {
                if (!response.ok) {
                    return response.json().then(err => { throw new Error(err.error || 'Unknown error'); });
                }
                return response.blob();
            })
            .then(blob => {
                const url = URL.createObjectURL(blob);
                const preview = document.getElementById('preview');
                const downloadLink = document.getElementById('downloadLink');
                const outputContainer = document.getElementById('outputContainer');
                
                preview.src = url;
                downloadLink.href = url;
                downloadLink.download = 'restored_' + document.getElementById('file').files[0].name;
                outputContainer.style.display = 'block';
            })
            .catch(error => {
                alert('Error: ' + error.message);
            })
            .finally(() => {
                // 処理が終わったらローディングを非表示にし、ボタンを再有効化
                loadingElement.style.display = 'none';
                submitButton.disabled = false;
            });
        });
    </script>
</body>
</html>
    """

if __name__ == '__main__':
    
    app.run(host='0.0.0.0', port=7860, debug=True)