service.py

import string
import random

from fastapi import HTTPException


class URLShortenerService:
    def __init__(self, db):
        self.db = db
        self.collection = db.urls
        self.base_url =  "http://localhost:8000/urls"

    characters = string.ascii_letters + string.digits

    async def get_shortened_url(self, original_url: str) -> str:
        short_code = self.generate_short_code()
        while await self.collection.find_one({"short_code": short_code}):
            short_code = self.generate_short_code()
        await self.save_to_database(short_code, original_url)
        return self.base_url+"/"+short_code

    def generate_short_code(self):
        return ''.join(random.choices(self.characters, k=6))

    async def save_to_database(self, short_code: str, original_url: str):
        await self.collection.insert_one({
            "original_url": original_url,
            "short_code": short_code
        })

    async def get_original_url(self, short_code: str) -> str:
        entry = await self.collection.find_one({"short_code": short_code})
        if not entry:
            raise HTTPException(status_code=404, detail="Short URL not found")
        return entry["original_url"]