qr-code_manager/app/utils/url_shortener.py

"""
URL Shortener utilities for QR Code Manager
"""

import os
import uuid
import string
import random
import json
from datetime import datetime

# Data storage directory
DATA_DIR = 'data'
SHORT_URLS_FILE = os.path.join(DATA_DIR, 'short_urls.json')

# Ensure data directory exists
os.makedirs(DATA_DIR, exist_ok=True)

class URLShortener:
    def __init__(self):
        self.base_domain = os.environ.get('APP_DOMAIN', 'localhost:5000')
        # Ensure we have the protocol
        if not self.base_domain.startswith(('http://', 'https://')):
            # Use HTTPS for production domains, HTTP for localhost
            protocol = 'https://' if 'localhost' not in self.base_domain else 'http://'
            self.base_domain = f"{protocol}{self.base_domain}"
        
        self.short_urls_db = self._load_short_urls()
    
    def _load_short_urls(self):
        """Load short URLs from JSON file"""
        try:
            if os.path.exists(SHORT_URLS_FILE):
                with open(SHORT_URLS_FILE, 'r', encoding='utf-8') as f:
                    return json.load(f)
            return {}
        except Exception as e:
            print(f"Error loading short URLs: {e}")
            return {}
    
    def _save_short_urls(self):
        """Save short URLs to JSON file"""
        try:
            with open(SHORT_URLS_FILE, 'w', encoding='utf-8') as f:
                json.dump(self.short_urls_db, f, indent=2, ensure_ascii=False)
        except Exception as e:
            print(f"Error saving short URLs: {e}")
    
    def generate_short_code(self, length=6):
        """Generate a random short code"""
        characters = string.ascii_letters + string.digits
        while True:
            short_code = ''.join(random.choice(characters) for _ in range(length))
            # Ensure uniqueness
            if short_code not in self.short_urls_db:
                return short_code
    
    def create_short_url(self, original_url, custom_code=None, title=""):
        """Create a shortened URL"""
        print(f"DEBUG: URLShortener.create_short_url called with url='{original_url}', custom_code='{custom_code}', title='{title}'")
        
        # Generate or use custom short code
        if custom_code and custom_code not in self.short_urls_db:
            short_code = custom_code
            print(f"DEBUG: Using custom short code: {short_code}")
        else:
            short_code = self.generate_short_code()
            print(f"DEBUG: Generated short code: {short_code}")
        
        # Ensure original URL has protocol
        if not original_url.startswith(('http://', 'https://')):
            original_url = f'https://{original_url}'
        
        # Create URL record
        url_data = {
            'id': str(uuid.uuid4()),
            'short_code': short_code,
            'original_url': original_url,
            'title': title,
            'clicks': 0,
            'created_at': datetime.now().isoformat(),
            'last_accessed': None
        }
        
        print(f"DEBUG: Adding to short_urls_db: {short_code} -> {url_data}")
        self.short_urls_db[short_code] = url_data
        
        print(f"DEBUG: Saving short URLs to file")
        self._save_short_urls()  # Persist to file
        print(f"DEBUG: Short URLs saved successfully")
        
        # Return the complete short URL
        short_url = f"{self.base_domain}/s/{short_code}"
        print(f"DEBUG: Returning short URL: {short_url}")
        return {
            'short_url': short_url,
            'short_code': short_code,
            'original_url': original_url,
            'id': url_data['id']
        }
    
    def get_original_url(self, short_code):
        """Get original URL from short code and track click"""
        if short_code in self.short_urls_db:
            url_data = self.short_urls_db[short_code]
            # Track click
            url_data['clicks'] += 1
            url_data['last_accessed'] = datetime.now().isoformat()
            self._save_short_urls()  # Persist to file
            return url_data['original_url']
        return None
    
    def get_url_stats(self, short_code):
        """Get statistics for a short URL"""
        return self.short_urls_db.get(short_code)
    
    def list_urls(self):
        """List all short URLs"""
        return list(self.short_urls_db.values())
    
    def delete_url(self, short_code):
        """Delete a short URL"""
        if short_code in self.short_urls_db:
            del self.short_urls_db[short_code]
            self._save_short_urls()  # Persist to file
            return True
        return False
    
    def url_exists(self, short_code):
        """Check if short URL exists"""
        return short_code in self.short_urls_db