import os
import sys
import tempfile
import uuid
import base64
import json
from io import BytesIO
import warnings
from importlib import import_module

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.pin import put_actions as put_pin_actions, pin_wait_change
from pywebio.output import (
    put_text, put_markdown, put_file, use_scope, put_scope, put_button, put_html, clear, put_table
)
from pywebio.session import eval_js, run_js

# Allow running this file directly via `python src/gui/gui_demo.py`.
if __package__ in (None, ""):
    sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))

# Import image steganography module
from ImageStego.steganography.lsb import embed_message, extract_message
from ImageStego.pipeline import ObfuscationPipeline, SUPPORTED_ALGORITHMS, _ALGORITHM_REGISTRY

# ==================== 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()


def _workspace_temp_dir() -> str:
    temp_dir = os.path.join(os.getcwd(), ".pywebio_temp")
    os.makedirs(temp_dir, exist_ok=True)
    return temp_dir


def _save_uploaded_file(upload: dict, suffix: str = "") -> str:
    extension = suffix or os.path.splitext(upload["filename"])[1] or ".bin"
    file_path = os.path.join(_workspace_temp_dir(), f"upload_{uuid.uuid4().hex}{extension}")
    with open(file_path, "wb") as file_obj:
        file_obj.write(upload["content"])
    return file_path


def _cleanup_files(*paths: str) -> None:
    for path in paths:
        if path and os.path.exists(path):
            try:
                os.remove(path)
            except OSError:
                pass

def _load_module(module_name: str):
    return import_module(module_name)


def _load_image_method(method_name: str):
    return _load_module(IMAGE_METHODS[method_name]['module'])


def _load_gif_module():
    return _load_module('ImageStego.obfuscation.gif_scramble')


def _load_mp4_module():
    return _load_module('ImageStego.obfuscation.mp4_scramble')


def _default_pipeline_layer(algorithm: str = "pixel") -> dict:
    defaults = dict(_ALGORITHM_REGISTRY[algorithm]["defaults"])
    return {
        "id": uuid.uuid4().hex,
        "algorithm": algorithm,
        "params": defaults,
    }


def _pipeline_editor_registry() -> dict:
    registry = {}
    for name, info in _ALGORITHM_REGISTRY.items():
        registry[name] = {
            "label": info["label"],
            "defaults": dict(info["defaults"]),
        }
    return registry


def _normalize_pipeline_layers(raw_layers: list) -> list:
    normalized = []
    for raw_layer in raw_layers:
        algorithm = raw_layer.get("algorithm")
        if algorithm not in _ALGORITHM_REGISTRY:
            continue

        defaults = dict(_ALGORITHM_REGISTRY[algorithm]["defaults"])
        params = dict(defaults)
        raw_params = raw_layer.get("params", {}) or {}

        if "seed" in params and raw_params.get("seed") not in (None, ""):
            params["seed"] = int(raw_params["seed"])
        if "iterations" in params and raw_params.get("iterations") not in (None, ""):
            params["iterations"] = int(raw_params["iterations"])
        if "block_size" in params and raw_params.get("block_size") not in (None, ""):
            params["block_size"] = int(raw_params["block_size"])
        if "password" in params:
            params["password"] = raw_params.get("password") or None
        if "use_timestamp" in params:
            params["use_timestamp"] = bool(raw_params.get("use_timestamp", False))

        normalized.append({
            "algorithm": algorithm,
            "params": params,
        })

    return normalized


def _mount_pipeline_editor(initial_layers: list) -> None:
    put_html('<div id="pipeline-editor-root"></div>')
    run_js(
        """
        (() => {
          const root = document.getElementById('pipeline-editor-root');
          if (!root) return;

          const registry = spec.registry;
          const order = spec.order;
          const initialLayers = JSON.parse(spec.layers_json);

          const cloneDefaults = (algorithm) => JSON.parse(JSON.stringify(registry[algorithm].defaults));
          const nextAlgorithm = () => order[0];

          const state = {
            layers: initialLayers.length ? initialLayers : [{
              id: `${Date.now()}_${Math.random().toString(16).slice(2)}`,
              algorithm: nextAlgorithm(),
              params: cloneDefaults(nextAlgorithm())
            }]
          };

          const fieldMeta = {
            seed: { label: 'Seed', type: 'number', placeholder: 'Optional integer' },
            iterations: { label: 'Iterations', type: 'number', placeholder: 'Algorithm default if empty' },
            block_size: { label: 'Block Size', type: 'number', placeholder: 'Algorithm default if empty' },
            password: { label: 'Password', type: 'text', placeholder: 'Optional password' },
            use_timestamp: { label: 'Use Timestamp', type: 'checkbox' }
          };

          const escapeHtml = (value) => String(value ?? '')
            .replace(/&/g, '&amp;')
            .replace(/</g, '&lt;')
            .replace(/>/g, '&gt;')
            .replace(/"/g, '&quot;');

          const createLayer = () => ({
            id: `${Date.now()}_${Math.random().toString(16).slice(2)}`,
            algorithm: nextAlgorithm(),
            params: cloneDefaults(nextAlgorithm())
          });

          const renderFields = (layer) => {
            const defaults = registry[layer.algorithm].defaults;
            return Object.keys(defaults).map((key) => {
              const meta = fieldMeta[key];
              if (!meta) return '';
              const value = layer.params[key];
              if (meta.type === 'checkbox') {
                return `
                  <label class="pipeline-field pipeline-checkbox">
                    <input type="checkbox" data-role="field" data-key="${key}" ${value ? 'checked' : ''}>
                    <span>${meta.label}</span>
                  </label>
                `;
              }
              return `
                <label class="pipeline-field">
                  <span>${meta.label}</span>
                  <input
                    type="${meta.type}"
                    data-role="field"
                    data-key="${key}"
                    value="${escapeHtml(value ?? '')}"
                    placeholder="${meta.placeholder}"
                  >
                </label>
              `;
            }).join('');
          };

          const summaryText = (layer) => {
            const params = layer.params || {};
            const visible = Object.entries(params).filter(([, value]) => value !== null && value !== '' && value !== false);
            if (!visible.length) return 'Using algorithm defaults';
            return visible.map(([k, v]) => `${k}=${v}`).join(' | ');
          };

          const attachCardEvents = () => {
            root.querySelectorAll('.pipeline-card').forEach((card) => {
              const layerId = card.dataset.layerId;
              const layer = state.layers.find((item) => item.id === layerId);
              if (!layer) return;

              const select = card.querySelector('[data-role="algorithm"]');
              select?.addEventListener('change', (event) => {
                layer.algorithm = event.target.value;
                layer.params = cloneDefaults(layer.algorithm);
                render();
              });

              card.querySelectorAll('[data-role="field"]').forEach((input) => {
                input.addEventListener('input', () => {
                  const key = input.dataset.key;
                  if (input.type === 'checkbox') {
                    layer.params[key] = input.checked;
                  } else if (input.type === 'number') {
                    layer.params[key] = input.value === '' ? null : Number(input.value);
                  } else {
                    layer.params[key] = input.value;
                  }
                  syncState();
                  const summary = card.querySelector('.pipeline-card-summary');
                  if (summary) summary.textContent = summaryText(layer);
                });
                input.addEventListener('change', () => {
                  const key = input.dataset.key;
                  if (input.type === 'checkbox') {
                    layer.params[key] = input.checked;
                  } else if (input.type === 'number') {
                    layer.params[key] = input.value === '' ? null : Number(input.value);
                  } else {
                    layer.params[key] = input.value;
                  }
                  syncState();
                });
              });

              card.querySelector('[data-role="remove"]')?.addEventListener('click', () => {
                if (state.layers.length === 1) return;
                state.layers = state.layers.filter((item) => item.id !== layerId);
                render();
              });

              card.setAttribute('draggable', 'true');
              card.addEventListener('dragstart', (event) => {
                card.classList.add('dragging');
                event.dataTransfer.setData('text/plain', layerId);
              });
              card.addEventListener('dragend', () => {
                card.classList.remove('dragging');
              });
              card.addEventListener('dragover', (event) => {
                event.preventDefault();
                card.classList.add('drop-target');
              });
              card.addEventListener('dragleave', () => {
                card.classList.remove('drop-target');
              });
              card.addEventListener('drop', (event) => {
                event.preventDefault();
                card.classList.remove('drop-target');
                const draggedId = event.dataTransfer.getData('text/plain');
                if (!draggedId || draggedId === layerId) return;
                const fromIndex = state.layers.findIndex((item) => item.id === draggedId);
                const toIndex = state.layers.findIndex((item) => item.id === layerId);
                if (fromIndex < 0 || toIndex < 0) return;
                const [draggedLayer] = state.layers.splice(fromIndex, 1);
                state.layers.splice(toIndex, 0, draggedLayer);
                render();
              });
            });

            root.querySelector('[data-role="add-layer"]')?.addEventListener('click', () => {
              state.layers.push(createLayer());
              render();
            });
            root.querySelector('[data-role="remove-layer"]')?.addEventListener('click', () => {
              if (state.layers.length === 1) return;
              state.layers.pop();
              render();
            });
          };

          const syncState = () => {
            window.__pipelineEditorState = {
              layers: state.layers.map((layer) => ({
                id: layer.id,
                algorithm: layer.algorithm,
                params: { ...layer.params }
              }))
            };
          };

          const render = () => {
            root.innerHTML = `
              <style>
                #pipeline-editor-root { margin-top: 12px; }
                .pipeline-toolbar {
                  display: flex;
                  gap: 12px;
                  align-items: center;
                  flex-wrap: wrap;
                  margin-bottom: 16px;
                  padding: 14px 16px;
                  border-radius: 16px;
                  background: linear-gradient(135deg, #f5f0e8, #e7f1ef);
                  border: 1px solid #d6ddd9;
                }
                .pipeline-toolbar button {
                  border: none;
                  border-radius: 999px;
                  padding: 10px 16px;
                  cursor: pointer;
                  font-weight: 600;
                }
                .pipeline-add { background: #1f7a5a; color: #fff; }
                .pipeline-remove { background: #efe3d3; color: #5f4631; }
                .pipeline-hint { color: #46615b; font-size: 14px; }
                .pipeline-card-list {
                  display: grid;
                  grid-template-columns: repeat(auto-fit, minmax(280px, 1fr));
                  gap: 16px;
                }
                .pipeline-card {
                  background: #fffdfa;
                  border: 1px solid #d9d1c7;
                  border-radius: 20px;
                  padding: 18px;
                  box-shadow: 0 16px 34px rgba(88, 71, 48, 0.08);
                }
                .pipeline-card.dragging { opacity: 0.4; }
                .pipeline-card.drop-target { outline: 2px dashed #1f7a5a; outline-offset: 4px; }
                .pipeline-card-header {
                  display: flex;
                  justify-content: space-between;
                  align-items: flex-start;
                  gap: 12px;
                  margin-bottom: 14px;
                  cursor: grab;
                }
                .pipeline-card-kicker {
                  font-size: 12px;
                  letter-spacing: 0.08em;
                  text-transform: uppercase;
                  color: #8f6f44;
                  margin-bottom: 4px;
                }
                .pipeline-card-title {
                  font-size: 20px;
                  font-weight: 700;
                  color: #21312e;
                  margin-bottom: 6px;
                }
                .pipeline-card-summary {
                  color: #5c6764;
                  font-size: 13px;
                  line-height: 1.5;
                }
                .pipeline-card-remove {
                  border: none;
                  background: #f8e4dc;
                  color: #8a3f2b;
                  border-radius: 999px;
                  padding: 8px 12px;
                  cursor: pointer;
                  font-weight: 700;
                }
                .pipeline-card-body {
                  display: grid;
                  gap: 12px;
                }
                .pipeline-field {
                  display: grid;
                  gap: 6px;
                  color: #2f3d3a;
                  font-size: 14px;
                }
                .pipeline-field span {
                  font-weight: 600;
                }
                .pipeline-field input,
                .pipeline-field select {
                  border: 1px solid #cfc7bc;
                  border-radius: 12px;
                  padding: 10px 12px;
                  background: #fff;
                }
                .pipeline-checkbox {
                  grid-template-columns: 20px 1fr;
                  align-items: center;
                  gap: 10px;
                }
                .pipeline-checkbox span { font-weight: 600; }
              </style>
              <div class="pipeline-toolbar">
                <button type="button" class="pipeline-add" data-role="add-layer">+ Add Layer</button>
                <button type="button" class="pipeline-remove" data-role="remove-layer">- Remove Layer</button>
                <div class="pipeline-hint">Drag cards to reorder layers before encryption. Current layers: ${state.layers.length}</div>
              </div>
              <div class="pipeline-card-list">
                ${state.layers.map((layer, index) => `
                  <section class="pipeline-card" data-layer-id="${layer.id}">
                    <div class="pipeline-card-header">
                      <div>
                        <div class="pipeline-card-kicker">Drag to Reorder</div>
                        <div class="pipeline-card-title">Layer ${index + 1}</div>
                        <div class="pipeline-card-summary">${escapeHtml(summaryText(layer))}</div>
                      </div>
                      <button type="button" class="pipeline-card-remove" data-role="remove">Remove</button>
                    </div>
                    <div class="pipeline-card-body">
                      <label class="pipeline-field">
                        <span>Algorithm</span>
                        <select data-role="algorithm">
                          ${order.map((name) => `<option value="${name}" ${name === layer.algorithm ? 'selected' : ''}>${escapeHtml(registry[name].label)}</option>`).join('')}
                        </select>
                      </label>
                      ${renderFields(layer)}
                    </div>
                  </section>
                `).join('')}
              </div>
            `;
            attachCardEvents();
            syncState();
          };

          window.getPipelineEditorState = () => JSON.stringify(window.__pipelineEditorState || { layers: [] });
          render();
        })();
        """,
        spec={
            "registry": _pipeline_editor_registry(),
            "order": SUPPORTED_ALGORITHMS,
            "layers_json": json.dumps(initial_layers, ensure_ascii=False),
        },
    )


def _render_pipeline_preview(layers: list) -> None:
    if not layers:
        put_markdown("### Layer Preview")
        put_text("No layers selected yet.")
        return

    put_markdown("### Layer Preview")
    rows = [["Order", "Algorithm", "Parameters"]]
    for index, layer in enumerate(layers, 1):
        params = {key: value for key, value in layer["params"].items() if value is not None}
        rows.append([
            str(index),
            _ALGORITHM_REGISTRY[layer["algorithm"]]["label"],
            json.dumps(params, ensure_ascii=False),
        ])
    put_table(rows)
    put_markdown("**Execution order:** " + " -> ".join(layer["algorithm"] for layer in layers))

# ==================== Image Obfuscation Method Definitions ====================
IMAGE_METHODS = {
    'pixel_scramble': {
        'display': 'Pixel Scramble',
        'module': 'ImageStego.obfuscation.pixel_scramble',
        'param_fields': []  # no specific parameters
    },
    'block_scramble': {
        'display': 'Block Scramble',
        'module': 'ImageStego.obfuscation.block_scramble',
        'param_fields': [
            {'name': 'block_size', 'label': 'Block size', 'type': 'number', 'default': 16}
        ]
    },
    'arnold_cat': {
        'display': 'Arnold Cat Map',
        'module': 'ImageStego.obfuscation.arnold_cat_map',
        'param_fields': [
            {'name': 'iterations', 'label': 'Iterations', 'type': 'number', 'default': 10}
        ]
    },
    'bakers_map': {
        'display': "Baker's Map",
        'module': 'ImageStego.obfuscation.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': 'ImageStego.obfuscation.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 = _load_image_method(method)

            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':
                _load_gif_module().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:
                _load_gif_module().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:
                from moviepy import VideoFileClip
                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")
                    _load_mp4_module().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:
                _load_mp4_module().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')

def pipeline_interface():
    """Visual builder for structured multi-layer pipeline encryption/decryption."""
    mode_data = input_group('Structured Pipeline Builder', [
        radio(
            name='op',
            label='Mode',
            options=['Encrypt', 'Decrypt'],
            value='Encrypt'
        )
    ])
    op = mode_data['op']

    if op == 'Encrypt':
        form_data = input_group('Pipeline Source', [
            file_upload(
                name='img_file',
                label='Select source image',
                accept='.png,.jpg,.jpeg,.bmp',
                max_size=50 * 1024 * 1024,
                required=True
            ),
            pw_input(
                name='output_name',
                label='Output filename',
                type=TEXT,
                value='pipeline_encrypted.png',
                required=True
            ),
        ])

        with use_scope('pipeline_result', clear=True):
            put_markdown('## Pipeline Result')
            put_markdown('### Layer Designer')
            _mount_pipeline_editor([_default_pipeline_layer()])
            put_pin_actions(
                'pipeline_encrypt_action',
                buttons=[
                    ('Encrypt with Current Layers', 'encrypt'),
                    ('Cancel', 'cancel'),
                ]
            )
            action = pin_wait_change('pipeline_encrypt_action')
            if action['value'] != 'encrypt':
                put_text('Pipeline encryption cancelled.').style('color: gray')
                return

            upload_path = ''
            output_path = ''
            sidecar_path = ''
            try:
                raw_state = eval_js("window.getPipelineEditorState ? window.getPipelineEditorState() : '{\"layers\":[]}'")
                editor_state = json.loads(raw_state)
                layers = _normalize_pipeline_layers(editor_state.get("layers", []))
                _render_pipeline_preview(layers)

                if not layers:
                    put_text('Please enable at least one layer.').style('color: red')
                    return

                upload_path = _save_uploaded_file(form_data['img_file'])
                output_filename = form_data['output_name'].strip() or 'pipeline_encrypted.png'
                if not output_filename.lower().endswith('.png'):
                    output_filename += '.png'
                output_path = os.path.join(_workspace_temp_dir(), f"{uuid.uuid4().hex}_{output_filename}")

                pipeline = ObfuscationPipeline()
                for layer in layers:
                    pipeline.add_layer(layer['algorithm'], **layer['params'])

                sidecar_path = pipeline.encrypt(upload_path, output_path)

                put_text('Structured pipeline encryption successful!').style('color: green')
                put_markdown('### Output Preview')
                put_html(image_to_base64(output_path)).style('margin-bottom: 20px;')

                with open(output_path, 'rb') as file_obj:
                    put_file(output_filename, file_obj.read(), label='Download encrypted image')
                with open(sidecar_path, 'rb') as file_obj:
                    put_file(os.path.basename(sidecar_path), file_obj.read(), label='Download layer metadata (.layers.json)')
            except Exception as e:
                put_text(f'Processing failed: {str(e)}').style('color: red')
            finally:
                _cleanup_files(upload_path, output_path, sidecar_path)
    else:
        form_data = input_group('Decrypt Structured Pipeline', [
            file_upload(
                name='img_file',
                label='Select encrypted image',
                accept='.png,.jpg,.jpeg,.bmp',
                max_size=50 * 1024 * 1024,
                required=True
            ),
            file_upload(
                name='sidecar_file',
                label='Select layer metadata file (.layers.json)',
                accept='.json',
                max_size=2 * 1024 * 1024,
                required=True
            ),
            pw_input(
                name='output_name',
                label='Restored filename',
                type=TEXT,
                value='pipeline_restored.png',
                required=True
            ),
        ])

        with use_scope('pipeline_result', clear=True):
            put_markdown('## Pipeline Result')
            encrypted_path = ''
            sidecar_path = ''
            output_path = ''
            try:
                encrypted_path = _save_uploaded_file(form_data['img_file'])
                sidecar_path = _save_uploaded_file(form_data['sidecar_file'], suffix='.layers.json')
                output_filename = form_data['output_name'].strip() or 'pipeline_restored.png'
                if not output_filename.lower().endswith('.png'):
                    output_filename += '.png'
                output_path = os.path.join(_workspace_temp_dir(), f"{uuid.uuid4().hex}_{output_filename}")

                with open(sidecar_path, 'r', encoding='utf-8') as file_obj:
                    layers = json.load(file_obj)

                _render_pipeline_preview(layers)

                pipeline = ObfuscationPipeline()
                pipeline.decrypt(encrypted_path, output_path, sidecar=sidecar_path)

                put_text('Structured pipeline decryption successful!').style('color: green')
                put_markdown('### Restored Preview')
                put_html(image_to_base64(output_path)).style('margin-bottom: 20px;')

                with open(output_path, 'rb') as file_obj:
                    put_file(output_filename, file_obj.read(), label='Download restored image')
            except Exception as e:
                put_text(f'Processing failed: {str(e)}').style('color: red')
            finally:
                _cleanup_files(encrypted_path, sidecar_path, output_path)

# ==================== 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', 'Pipeline Builder', '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 == 'Pipeline Builder':
        pipeline_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)
