# app.py

import gradio as gr
from utils import WatermarkProcessor
import json
import tempfile
import os
from datetime import datetime
import cv2
from PIL import Image
import numpy as np

class WatermarkGUI:
    def __init__(self):
        self.processor = WatermarkProcessor()
        self.create_interface()

    def process_watermark(self, image, watermark_text, author, purpose, opacity):
        """Process watermark with metadata"""
        if image is None or watermark_text.strip() == "":
            return None, "Please provide both image and watermark text"
            
        metadata = {
            "author": author,
            "purpose": purpose,
            "opacity": opacity
        }
        
        # Save temporary image
        temp_path = tempfile.mktemp(suffix='.png')
        Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB)).save(temp_path)
        
        # Add watermark
        result_path, message = self.processor.encode(temp_path, watermark_text, metadata)
        
        if "Error" in message:
            return None, message
            
        # Generate quality report
        quality_report = self.processor.analyze_quality(temp_path, result_path)
        quality_data = json.loads(quality_report)
        
        # Create formatted report
        report = f"""
### Watermark Quality Report
- Quality Score: {quality_data['quality_score']}%
- PSNR: {quality_data['psnr']} dB
- Histogram Similarity: {quality_data['histogram_similarity'] * 100:.2f}%
- Modified Pixels: {quality_data['modified_pixels']:,}

### Metadata
- Author: {author}
- Purpose: {purpose}
- Timestamp: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}
        """
        
        os.remove(temp_path)
        return cv2.imread(result_path), report

    def detect_watermark(self, image):
        """Detect and extract watermark"""
        if image is None:
            return "Please provide an image"
            
        # Save temporary image
        temp_path = tempfile.mktemp(suffix='.png')
        Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB)).save(temp_path)
        
        # Extract watermark
        result = self.processor.decode(temp_path)
        os.remove(temp_path)
        
        try:
            # Parse JSON result
            data = json.loads(result)
            report = f"""
### Extracted Watermark
Text: {data['text']}

### Metadata
- Timestamp: {data['timestamp']}
- Author: {data['metadata'].get('author', 'N/A')}
- Purpose: {data['metadata'].get('purpose', 'N/A')}
            """
            return report
        except:
            return result

    def create_interface(self):
        """Create Gradio interface"""
        with gr.Blocks(css="footer {visibility: hidden}") as self.interface:
            gr.HTML("""<a href="https://visitorbadge.io/status?path=https%3A%2F%2Fgunship999-SecureWatermark.hf.space"> <img src="https://api.visitorbadge.io/api/visitors?path=https%3A%2F%2Fgunship999-SecureWatermark.hf.space&countColor=%23263759" /> </a>""")
            
            gr.Markdown("""# Enhanced Image Watermarking System
                       
### Welcome to Secure Watermark - Advanced Image Protection System

🔒 **Key Features:**
- **Dual Watermarking Technology**: Supports both steganography and PNG metadata
- **Secure Encryption**: Military-grade encryption for watermark data
- **Quality Assurance**: Real-time quality analysis and reporting
- **Metadata Support**: Track authorship, purpose, and timestamps
- **Integrity Verification**: Hash-based image tampering detection

💡 **Perfect for:**
- Copyright Protection
- Digital Asset Management
- Document Authentication
- Creative Work Protection

Try our system by uploading an image and adding your watermark below!
            """)
            
            with gr.Tabs():
                # Add Watermark Tab

                with gr.Tab("Add Watermark"):
                    with gr.Row():
                        with gr.Column():
                            input_image = gr.Image(label="Input Image", type="numpy")
                            watermark_text = gr.Textbox(label="Watermark Text")
                            author = gr.Textbox(label="Author", placeholder="Enter author name")
                            purpose = gr.Textbox(label="Purpose", placeholder="Enter watermark purpose")
                            opacity = gr.Slider(minimum=0.1, maximum=1.0, value=0.3, 
                                             label="Watermark Opacity")
                            
                            with gr.Row():
                                process_btn = gr.Button("Add Watermark", variant="primary")
                                
                        with gr.Column():
                            result_image = gr.Image(label="Watermarked Image")
                            quality_report = gr.Markdown(label="Quality Report")
                
                # Detect Watermark Tab
                with gr.Tab("Detect Watermark"):
                    with gr.Row():
                        detect_image = gr.Image(label="Input Image", type="numpy")
                        detect_result = gr.Markdown(label="Detected Watermark")
                    detect_btn = gr.Button("Detect Watermark")
            
            # Event handlers            
            process_btn.click(
                fn=self.process_watermark,
                inputs=[input_image, watermark_text, author, purpose, opacity],
                outputs=[result_image, quality_report]
            )
            
            detect_btn.click(
                fn=self.detect_watermark,
                inputs=[detect_image],
                outputs=detect_result
            )

    def launch(self, *args, **kwargs):
        """Launch the interface"""
        self.interface.launch(*args, **kwargs)

if __name__ == "__main__":
    app = WatermarkGUI()
    app.launch()