controller.py

from pydantic import HttpUrl
from fastapi import HTTPException, status
from datetime import datetime
from model import Database
import random
import string

class URLController:
    def __init__(self):
        self.db = Database()
        self.cursor, self.conn = self.db.conect("database", "url_shortener")

    def create_shorten_url(self, url: HttpUrl, shortCode: str | None):
        url = str(url)  # Convert HttpUrl to string
        self._ensure_url_not_exists(url)

        # If ShortCode is None, it generates one; otherwise, it uses the provided ShortCode.
        if shortCode is None:
            shortCode = self._generate_unique_short_code()
        else:
            self._ensure_not_short_code_exists(shortCode)

        data = self.db.save_shorten_url(url, shortCode, self.cursor, self.conn)
        data["shortCode"] = shortCode

        return data

    def get_shorten_url(self, shortCode: str):
        # If the short code does not exist, neither does the URL.
        self._ensure_short_code_exists(shortCode)
        # Logic for retrieving the shortened URL details
        return self.db.get_shorten_url(shortCode, self.cursor)

    def update_shorten_url(self, shortCode: str, url: HttpUrl, new_shortCode: str | None):
        url = str(url)  # Convert HttpUrl to string
        # If the shortCode does not exist, return a 404 error
        self._ensure_short_code_exists(shortCode)

        if not (new_shortCode is None):
            # If the new shortCode is already registered, return a 409 error
            self._ensure_not_short_code_exists(new_shortCode)
        else:
            new_shortCode = shortCode
        # If the URL is already registered, return a 409 error
        self._ensure_url_not_exists(url)

        return self.db.update_shorten_url(shortCode, url, new_shortCode, self.cursor, self.conn)

    def delete_shorten_url(self, shortCode: str):
        # If the shortCode does not exist, return a 404 error
        self._ensure_short_code_exists(shortCode)

        self.db.delete_shorten_url(shortCode, self.cursor, self.conn)

    def get_shorten_url_stats(self, shortCode: str):
        # If the shortCode does not exist, return a 404 error
        self._ensure_short_code_exists(shortCode)

        return self.db.get_shorten_url_stats(shortCode, self.cursor)
    
    def access_url(self, shortCode):
        data = self.get_shorten_url(shortCode)
        url = data["url"]
        self.increment_clicks(shortCode)
        return url
    
    def increment_clicks(self, shortCode: str):
        self.db.increment_counter(shortCode, self.cursor, self.conn)
    
    def _generate_unique_short_code(self):
        """
            Generate a unique 6-character short code.
            Use numbers and uppercase and lowercase letters.
        """
        chars = string.ascii_letters + string.digits
        while True:
            code = ''.join(random.choices(chars, k=6))
            if not self.db.short_code_exists(code, self.cursor):
                return code
    
    def _ensure_short_code_exists(self, shortCode: str):
        """
            Raise a 404 if the short code does not exist.
        """
        if not self.db.short_code_exists(shortCode, self.cursor):
            raise HTTPException(
                status_code=status.HTTP_404_NOT_FOUND,
                detail="Shortened URL not found."
            )
        
    def _ensure_not_short_code_exists(self, shortCode: str):
        """
            Raise a 409 if the short code does not exist.
        """
        if self.db.short_code_exists(shortCode, self.cursor):
            raise HTTPException(
                status_code=status.HTTP_409_CONFLICT,
                detail="Short Code is already registered."
            )

    def _ensure_url_not_exists(self, url: str):
        """
            Raise a 409 error if the URL already exists.
        """
        if self.db.url_exists(url, self.cursor):
            raise HTTPException(
                status_code=status.HTTP_409_CONFLICT,
                detail="URL is already registered."
            )