import os
import tempfile
import uuid
import base64
from io import BytesIO
import warnings

from PIL import Image
from pywebio import start_server
from pywebio.input import input_group, file_upload, select, checkbox, input as pw_input, NUMBER, TEXT, radio, textarea
from pywebio.output import (
    put_text, put_markdown, put_file, use_scope, put_scope, put_button, put_html, clear
)

# Import image obfuscation modules
from ImageStego.obfuscation import pixel_scramble
from ImageStego.obfuscation import block_scramble
from ImageStego.obfuscation import arnold_cat_map
from ImageStego.obfuscation import bakers_map
from ImageStego.obfuscation import henon_map

# Import GIF obfuscation module
from ImageStego.obfuscation import gif_scramble

# Import MP4 obfuscation module
from ImageStego.obfuscation import mp4_scramble
from moviepy import VideoFileClip

# Import image steganography module
from ImageStego.steganography.lsb import embed_message, extract_message

# ==================== Utility Functions ====================
def image_to_base64(img_path, max_width=300):
    """Convert image to base64 encoded thumbnail HTML"""
    try:
        with Image.open(img_path) as img:
            if img.width > max_width:
                ratio = max_width / img.width
                new_size = (max_width, int(img.height * ratio))
                img.thumbnail(new_size, Image.Resampling.LANCZOS)

            buffer = BytesIO()
            img.save(buffer, format='PNG')
            img_base64 = base64.b64encode(buffer.getvalue()).decode()
            return f'<img src="data:image/png;base64,{img_base64}" style="max-width:100%; border:1px solid #ddd; border-radius:4px; padding:5px;">'
    except Exception as e:
        return f'<p style="color:red">Image preview failed: {e}</p>'

def reset_app():
    """Reset the whole app (return to main menu)"""
    clear()
    main()

# ==================== Image Obfuscation Method Definitions ====================
IMAGE_METHODS = {
    'pixel_scramble': {
        'display': 'Pixel Scramble',
        'module': pixel_scramble,
        'param_fields': []  # no specific parameters
    },
    'block_scramble': {
        'display': 'Block Scramble',
        'module': block_scramble,
        'param_fields': [
            {'name': 'block_size', 'label': 'Block size', 'type': 'number', 'default': 16}
        ]
    },
    'arnold_cat': {
        'display': 'Arnold Cat Map',
        'module': arnold_cat_map,
        'param_fields': [
            {'name': 'iterations', 'label': 'Iterations', 'type': 'number', 'default': 10}
        ]
    },
    'bakers_map': {
        'display': "Baker's Map",
        'module': bakers_map,
        'param_fields': [
            {'name': 'iterations', 'label': 'Iterations', 'type': 'number', 'default': 20},
            {'name': 'password', 'label': 'Password (optional)', 'type': 'text', 'default': ''}
        ]
    },
    'henon_map': {
        'display': 'Henon Map',
        'module': henon_map,
        'param_fields': [
            {'name': 'iterations', 'label': 'Iterations', 'type': 'number', 'default': 1},
            {'name': 'password', 'label': 'Password (optional)', 'type': 'text', 'default': ''}
        ]
    }
}

# ==================== Image Obfuscation Interface ====================
def image_interface():
    """Main flow for image obfuscation"""
    # Step 1: Choose method
    method_data = input_group('Select Obfuscation Method', [
        select(
            name='method',
            label='Obfuscation method',
            options=[(info['display'], name) for name, info in IMAGE_METHODS.items()],
            value='pixel_scramble'
        )
    ])
    method = method_data['method']

    # Step 2: Method-specific parameters form
    fields = [
        file_upload(
            name='img_file',
            label="Select image (PNG/JPG/BMP)",
            accept=".png,.jpg,.jpeg,.bmp",
            max_size=50 * 1024 * 1024,
            placeholder="Click to upload or drag & drop",
            required=True
        ),
        select(
            name='op',
            label="Select operation",
            options=['obfuscate', 'deobfuscate'],
            value='obfuscate'
        ),
        pw_input(
            name='seed',
            label="Seed (integer, optional)",
            type=NUMBER,
            placeholder="Leave empty to not use seed"
        ),
        checkbox(
            name='use_timestamp',
            label="Embed timestamp during obfuscation",
            options=[('Embed', True)],
            value=[True]
        ),
        pw_input(
            name='manual_timestamp',
            label="Manual timestamp (integer, used only for deobfuscation)",
            type=NUMBER,
            placeholder="Leave empty if timestamp is already in image metadata"
        ),
    ]

    # Add method-specific parameters
    for field in IMAGE_METHODS[method]['param_fields']:
        input_type = NUMBER if field['type'] == 'number' else TEXT
        fields.append(
            pw_input(
                name=field['name'],
                label=field['label'],
                type=input_type,
                value=field.get('default'),
                required=False
            )
        )

    param_data = input_group(f'{IMAGE_METHODS[method]["display"]} Parameters', fields)

    # Step 3: Process and display results
    with use_scope('image_result', clear=True):
        put_markdown('## Result')
        try:
            img_file = param_data['img_file']
            if not img_file:
                put_text("Please upload an image").style('color: red')
                return

            # Save uploaded file to temporary path
            temp_dir = tempfile.gettempdir()
            upload_filename = f"upload_{uuid.uuid4().hex}.png"
            upload_path = os.path.join(temp_dir, upload_filename)
            with open(upload_path, 'wb') as f:
                f.write(img_file['content'])

            # Original image preview
            put_markdown('### Original Image Preview')
            put_html(image_to_base64(upload_path)).style('margin-bottom: 20px;')

            op = param_data['op']
            seed_val = param_data.get('seed')
            use_timestamp = True in param_data.get('use_timestamp', [])
            manual_timestamp = param_data.get('manual_timestamp')

            module = IMAGE_METHODS[method]['module']

            if op == 'obfuscate':
                out_filename = f"{method}_{op}_{uuid.uuid4().hex}.png"
                out_path = os.path.join(temp_dir, out_filename)

                if method == 'pixel_scramble':
                    ts = module.obfuscate(
                        image_path=upload_path,
                        output_path=out_path,
                        seed=seed_val,
                        use_timestamp=use_timestamp
                    )
                elif method == 'block_scramble':
                    block_size = param_data.get('block_size', 16)
                    ts = module.obfuscate(
                        image_path=upload_path,
                        output_path=out_path,
                        block_size=block_size,
                        seed=seed_val,
                        use_timestamp=use_timestamp
                    )
                elif method == 'arnold_cat':
                    iterations = param_data.get('iterations', 10)
                    ts = module.obfuscate(
                        image_path=upload_path,
                        output_path=out_path,
                        iterations=iterations,
                        seed=seed_val,
                        use_timestamp=use_timestamp
                    )
                elif method == 'bakers_map':
                    iterations = param_data.get('iterations', 20)
                    password = param_data.get('password')
                    if password == '':
                        password = None
                    ts = module.obfuscate(
                        image_path=upload_path,
                        output_path=out_path,
                        iterations=iterations,
                        seed=seed_val,
                        password=password,
                        use_timestamp=use_timestamp
                    )
                elif method == 'henon_map':
                    iterations = param_data.get('iterations', 1)
                    password = param_data.get('password')
                    if password == '':
                        password = None
                    module.obfuscate(
                        input_path=upload_path,
                        output_path=out_path,
                        iterations=iterations,
                        seed=seed_val,
                        password=password
                    )
                    ts = None
                put_text(f'✅ {method} obfuscation successful!').style('color: green')
            else:  # deobfuscate
                out_filename = f"{method}_{op}_{uuid.uuid4().hex}.png"
                out_path = os.path.join(temp_dir, out_filename)

                if method == 'pixel_scramble':
                    module.deobfuscate(
                        obf_image_path=upload_path,
                        output_path=out_path,
                        seed=seed_val,
                        timestamp=manual_timestamp
                    )
                elif method == 'block_scramble':
                    block_size = param_data.get('block_size', 16)
                    module.deobfuscate(
                        obf_image_path=upload_path,
                        output_path=out_path,
                        block_size=block_size,
                        seed=seed_val,
                        timestamp=manual_timestamp
                    )
                elif method == 'arnold_cat':
                    iterations = param_data.get('iterations')
                    module.deobfuscate(
                        obf_image_path=upload_path,
                        output_path=out_path,
                        iterations=iterations,
                        seed=seed_val,
                        timestamp=manual_timestamp
                    )
                elif method == 'bakers_map':
                    iterations = param_data.get('iterations')
                    password = param_data.get('password')
                    if password == '':
                        password = None
                    module.deobfuscate(
                        obf_image_path=upload_path,
                        output_path=out_path,
                        iterations=iterations,
                        seed=seed_val,
                        timestamp=manual_timestamp,
                        password=password
                    )
                elif method == 'henon_map':
                    iterations = param_data.get('iterations')
                    password = param_data.get('password')
                    if password == '':
                        password = None
                    module.deobfuscate(
                        input_path=upload_path,
                        output_path=out_path,
                        iterations=iterations,
                        seed=seed_val,
                        password=password
                    )
                put_text(f'✅ {method} deobfuscation successful!').style('color: green')

            # Processed image preview
            put_markdown('### Processed Image Preview')
            put_html(image_to_base64(out_path)).style('margin-bottom: 20px;')

            # Provide download link
            with open(out_path, 'rb') as f:
                content = f.read()
            put_file(out_filename, content, label='📥 Click to download processed image')

            # Clean up uploaded temporary file
            os.remove(upload_path)
        except Exception as e:
            put_text(f'❌ Processing failed: {str(e)}').style('color: red')

# ==================== GIF Obfuscation Interface ====================
def gif_interface():
    """Main flow for GIF obfuscation"""
    fields = [
        file_upload(
            name='gif_file',
            label="Select GIF file",
            accept=".gif",
            max_size=50 * 1024 * 1024,
            placeholder="Click to upload or drag & drop",
            required=True
        ),
        select(
            name='op',
            label="Select operation",
            options=['obfuscate', 'deobfuscate'],
            value='obfuscate'
        ),
        pw_input(
            name='seed',
            label="Seed (integer, optional)",
            type=NUMBER,
            placeholder="Leave empty to not use seed"
        ),
        checkbox(
            name='use_timestamp',
            label="Embed timestamp during obfuscation (recommended to leave unchecked to preserve reversibility)",
            options=[('Embed', True)],
            value=[]  # default unchecked
        )
    ]

    param_data = input_group('GIF Obfuscation Settings', fields)

    with use_scope('gif_result', clear=True):
        put_markdown('## Result')
        try:
            gif_file = param_data['gif_file']
            if not gif_file:
                put_text("Please upload a GIF file").style('color: red')
                return

            temp_dir = tempfile.gettempdir()
            upload_filename = f"upload_{uuid.uuid4().hex}.gif"
            upload_path = os.path.join(temp_dir, upload_filename)
            with open(upload_path, 'wb') as f:
                f.write(gif_file['content'])

            # Original GIF preview (first frame)
            put_markdown('### Original GIF Preview (first frame)')
            try:
                with Image.open(upload_path) as preview_img:
                    preview_img.seek(0)
                    buffer = BytesIO()
                    preview_img.save(buffer, format='PNG')
                    img_base64 = base64.b64encode(buffer.getvalue()).decode()
                    put_html(f'<img src="data:image/png;base64,{img_base64}" style="max-width:300px; border:1px solid #ddd; border-radius:4px; padding:5px;">').style('margin-bottom: 20px;')
            except Exception as e:
                put_text(f"Preview generation failed: {e}").style('color: orange')

            op = param_data['op']
            seed_val = param_data.get('seed')
            use_timestamp = True in param_data.get('use_timestamp', [])

            out_filename = f"gif_{op}_{uuid.uuid4().hex}.gif"
            out_path = os.path.join(temp_dir, out_filename)

            if op == 'obfuscate':
                gif_scramble.obfuscate_gif(
                    gif_path=upload_path,
                    output_gif_path=out_path,
                    seed=seed_val,
                    use_timestamp=use_timestamp
                )
                put_text('✅ GIF obfuscation successful!').style('color: green')
            else:
                gif_scramble.deobfuscate_gif(
                    obf_gif_path=upload_path,
                    output_gif_path=out_path,
                    seed=seed_val
                )
                put_text('✅ GIF deobfuscation successful!').style('color: green')

            with open(out_path, 'rb') as f:
                content = f.read()
            put_file(out_filename, content, label='📥 Click to download processed GIF')

            os.remove(upload_path)
        except Exception as e:
            put_text(f'❌ Processing failed: {str(e)}').style('color: red')

# ==================== MP4 Obfuscation Interface ====================
def mp4_interface():
    """Main flow for MP4 obfuscation"""
    fields = [
        file_upload(
            name='mp4_file',
            label="Select MP4 file",
            accept=".mp4",
            max_size=200 * 1024 * 1024,
            placeholder="Click to upload or drag & drop",
            required=True
        ),
        select(
            name='op',
            label="Select operation",
            options=['obfuscate', 'deobfuscate'],
            value='obfuscate'
        ),
        pw_input(
            name='seed',
            label="Seed (integer, optional)",
            type=NUMBER,
            placeholder="Leave empty to not use seed"
        ),
        checkbox(
            name='use_timestamp',
            label="Embed timestamp during obfuscation (ignored for MP4)",
            options=[('Embed', True)],
            value=[]
        )
    ]

    param_data = input_group('MP4 Obfuscation Settings', fields)

    with use_scope('mp4_result', clear=True):
        put_markdown('## Result')
        try:
            mp4_file = param_data['mp4_file']
            if not mp4_file:
                put_text("Please upload an MP4 file").style('color: red')
                return

            temp_dir = tempfile.gettempdir()
            upload_filename = f"upload_{uuid.uuid4().hex}.mp4"
            upload_path = os.path.join(temp_dir, upload_filename)
            with open(upload_path, 'wb') as f:
                f.write(mp4_file['content'])

            # Original MP4 preview (first frame)
            put_markdown('### Original MP4 Preview (first frame)')
            try:
                clip = VideoFileClip(upload_path)
                frame = clip.get_frame(0)
                clip.close()
                img = Image.fromarray(frame)
                buffer = BytesIO()
                img.save(buffer, format='PNG')
                img_base64 = base64.b64encode(buffer.getvalue()).decode()
                put_html(f'<img src="data:image/png;base64,{img_base64}" style="max-width:300px; border:1px solid #ddd; border-radius:4px; padding:5px;">').style('margin-bottom: 20px;')
                put_text(f"Duration: {clip.duration:.2f}s, FPS: {clip.fps:.2f}").style('color: gray')
            except Exception as e:
                put_text(f"Preview generation failed: {e}").style('color: orange')

            op = param_data['op']
            seed_val = param_data.get('seed')
            use_timestamp = True in param_data.get('use_timestamp', [])

            out_filename = f"mp4_{op}_{uuid.uuid4().hex}.mp4"
            out_path = os.path.join(temp_dir, out_filename)

            if op == 'obfuscate':
                with warnings.catch_warnings():
                    warnings.simplefilter("ignore")
                    mp4_scramble.obfuscate_mp4(
                        input_path=upload_path,
                        output_path=out_path,
                        seed=seed_val,
                        use_timestamp=use_timestamp
                    )
                put_text('✅ MP4 obfuscation successful!').style('color: green')
            else:
                mp4_scramble.deobfuscate_mp4(
                    input_path=upload_path,
                    output_path=out_path,
                    seed=seed_val
                )
                put_text('✅ MP4 deobfuscation successful!').style('color: green')

            # Check if output file exists
            if os.path.exists(out_path):
                file_size = os.path.getsize(out_path)
                put_text(f"✅ Output file generated, size: {file_size} bytes").style('color: green')
                put_text(f"📁 File saved at: {out_path}").style('color: blue')
                put_text("Please copy the file from the above path").style('color: orange')
            else:
                put_text("❌ Output file not generated!").style('color: red')
                return

            # Clean up uploaded temporary file
            os.remove(upload_path)
        except Exception as e:
            put_text(f'❌ Processing failed: {str(e)}').style('color: red')

# ==================== Image Steganography Interface ====================
def image_stego_interface():
    """Main flow for image steganography (LSB)"""
    # Step 1: Choose operation
    op_data = input_group('Select Steganography Operation', [
        radio(
            name='op',
            label='Operation',
            options=['Embed', 'Extract'],
            value='Embed'
        )
    ])
    op = op_data['op']

    # Step 2: Build form based on operation
    if op == 'Embed':
        fields = [
            file_upload(
                name='img_file',
                label="Select original image (PNG/JPG/BMP)",
                accept=".png,.jpg,.jpeg,.bmp",
                max_size=50 * 1024 * 1024,
                placeholder="Click to upload or drag & drop",
                required=True
            ),
            textarea(
                name='message',
                label="Message to hide",
                rows=5,
                required=True,
                placeholder="Enter text to hide..."
            ),
            pw_input(
                name='seed',
                label="Seed (integer, optional)",
                type=NUMBER,
                placeholder="Leave empty to not use seed"
            ),
            pw_input(
                name='password',
                label="Password (optional)",
                type=TEXT,
                placeholder="If encryption is needed, enter password"
            ),
            checkbox(
                name='include_timestamp',
                label="Embed timestamp",
                options=[('Include', True)],
                value=[True]
            )
        ]
        param_data = input_group('Embed Parameters', fields)
    else:  # Extract
        fields = [
            file_upload(
                name='img_file',
                label="Select stego image (PNG/JPG/BMP)",
                accept=".png,.jpg,.jpeg,.bmp",
                max_size=50 * 1024 * 1024,
                placeholder="Click to upload or drag & drop",
                required=True
            ),
            pw_input(
                name='seed',
                label="Seed (integer, optional)",
                type=NUMBER,
                placeholder="Leave empty to not use seed"
            ),
            pw_input(
                name='password',
                label="Password (optional)",
                type=TEXT,
                placeholder="If a password was used during embedding, enter it"
            )
        ]
        param_data = input_group('Extract Parameters', fields)

    # Step 3: Process and display results
    with use_scope('stego_result', clear=True):
        put_markdown('## Result')
        try:
            img_file = param_data['img_file']
            if not img_file:
                put_text("Please upload an image").style('color: red')
                return

            # Save uploaded file to temporary path
            temp_dir = tempfile.gettempdir()
            upload_filename = f"upload_{uuid.uuid4().hex}.png"
            upload_path = os.path.join(temp_dir, upload_filename)
            with open(upload_path, 'wb') as f:
                f.write(img_file['content'])

            # Original image preview
            put_markdown('### Original Image Preview')
            put_html(image_to_base64(upload_path)).style('margin-bottom: 20px;')

            seed = param_data.get('seed')
            password = param_data.get('password')
            if password == '':
                password = None

            if op == 'Embed':
                message = param_data['message']
                include_timestamp = True in param_data.get('include_timestamp', [])

                out_filename = f"stego_embed_{uuid.uuid4().hex}.png"
                out_path = os.path.join(temp_dir, out_filename)

                embed_message(
                    image_path=upload_path,
                    message=message,
                    output_path=out_path,
                    seed=seed,
                    password=password,
                    include_timestamp=include_timestamp
                )
                put_text('✅ Message embedded successfully!').style('color: green')

                # Stego image preview
                put_markdown('### Stego Image Preview')
                put_html(image_to_base64(out_path)).style('margin-bottom: 20px;')

                # Download
                with open(out_path, 'rb') as f:
                    content = f.read()
                put_file(out_filename, content, label='📥 Click to download stego image')
            else:  # Extract
                result = extract_message(
                    image_path=upload_path,
                    seed=seed,
                    password=password
                )
                put_text('✅ Message extracted successfully!').style('color: green')
                put_markdown(f'**Extracted message:**\n\n{result["message"]}')
                put_markdown(f'**Timestamp:** {result["timestamp_human"]} (Unix time: {result["timestamp"]})')

            # Clean up uploaded temporary file
            os.remove(upload_path)
        except Exception as e:
            put_text(f'❌ Processing failed: {str(e)}').style('color: red')

# ==================== Main Interface (Select Media Type) ====================
def main():
    put_markdown('# ImageStego')
    put_button("🔄 Back to Main Menu", onclick=reset_app, color='warning').style('margin-bottom: 20px;')

    # Step 1: Choose media type
    media_choice = input_group('Select Processing Type', [
        radio(
            name='media_type',
            label='Please select the media type to process',
            options=['Image Obfuscation', 'GIF Animation', 'MP4 Video', 'Steganography (LSB)'],
            value='Image Obfuscation'
        )
    ])

    media_type = media_choice['media_type']

    # Branch according to selection
    if media_type == 'Image Obfuscation':
        image_interface()
    elif media_type == 'GIF Animation':
        gif_interface()
    elif media_type == 'MP4 Video':
        mp4_interface()
    else:  # Steganography
        image_stego_interface()

if __name__ == '__main__':
    start_server(main, port=8080, debug=True, auto_open_webbrowser=True)