app.py

from __future__ import annotations

import json
import os
import re
import sqlite3
import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Iterator, Optional
from urllib.parse import urlparse

import requests
from fastapi import Depends, FastAPI, Request
from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse, Response
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from starlette.routing import Match
from slowapi import Limiter, _rate_limit_exceeded_handler
from slowapi.errors import RateLimitExceeded
from slowapi.middleware import SlowAPIMiddleware
from slowapi.util import get_remote_address

BASE_DIR = Path(__file__).parent
DATA_PATH = BASE_DIR / "data"
LEGACY_JSON_PATH = DATA_PATH / "url_store.json"

ALIAS_REGEX = re.compile(r"^[A-Za-z0-9_-]+$")
RESERVED_ALIASES = {
    "api",
    "health",
    "static",
    "favicon.svg",
    "robots.txt",
}

templates = Jinja2Templates(directory=str(BASE_DIR / "templates"))


def _connect(db_path: str) -> sqlite3.Connection:
    DATA_PATH.mkdir(parents=True, exist_ok=True)
    conn = sqlite3.connect(db_path, check_same_thread=False)
    conn.row_factory = sqlite3.Row
    return conn


def get_db(request: Request) -> Iterator[sqlite3.Connection]:
    db = _connect(request.app.state.config["DATABASE_PATH"])
    try:
        yield db
    finally:
        db.close()


def init_db(db_path: str) -> None:
    DATA_PATH.mkdir(parents=True, exist_ok=True)
    conn = sqlite3.connect(db_path)
    try:
        conn.execute(
            """
            CREATE TABLE IF NOT EXISTS links (
                alias TEXT PRIMARY KEY,
                url TEXT NOT NULL,
                created_at TEXT NOT NULL,
                expires_at TEXT,
                usage_count INTEGER NOT NULL DEFAULT 0
            )
            """
        )
        conn.commit()
    finally:
        conn.close()


def _db_has_links(db: sqlite3.Connection) -> bool:
    row = db.execute("SELECT 1 FROM links LIMIT 1").fetchone()
    return row is not None


def _migrate_json_to_db(db: sqlite3.Connection) -> None:
    if not LEGACY_JSON_PATH.exists():
        return
    if _db_has_links(db):
        return
    try:
        with LEGACY_JSON_PATH.open("r", encoding="utf-8") as fh:
            data = json.load(fh)
    except (json.JSONDecodeError, OSError):
        return

    if not isinstance(data, dict):
        return

    for alias, entry in data.items():
        if not isinstance(entry, dict):
            continue
        db.execute(
            """
            INSERT OR IGNORE INTO links (alias, url, created_at, expires_at, usage_count)
            VALUES (?, ?, ?, ?, ?)
            """,
            (
                alias,
                entry.get("url"),
                entry.get("created_at"),
                entry.get("expires_at"),
                int(entry.get("usage_count") or 0),
            ),
        )
    db.commit()


def _parse_expiry(value: Optional[str]) -> Optional[str]:
    if not value:
        return None
    try:
        dt = datetime.fromisoformat(value)
    except ValueError:
        raise ValueError("Invalid expiry format. Use ISO 8601, e.g. 2025-12-31T23:59")

    if dt.tzinfo is None:
        dt = dt.replace(tzinfo=timezone.utc)
    dt = dt.astimezone(timezone.utc)
    if dt <= datetime.now(timezone.utc):
        raise ValueError("Expiry must be in the future")
    return dt.isoformat()


def _build_short_url(request: Request, alias: str) -> str:
    return str(request.base_url).rstrip("/") + f"/{alias}"


def _create_tiny_url(target_url: str) -> Optional[str]:
    try:
        response = requests.get(
            "https://tinyurl.com/api-create.php",
            params={"url": target_url},
            timeout=8,
        )
        if response.status_code == 200:
            return response.text.strip()
    except requests.RequestException:
        return None
    return None


def _validate_url(url: str, config: dict[str, Any]) -> Optional[str]:
    if not url:
        return "URL is required"
    if len(url) > config["MAX_URL_LENGTH"]:
        return "URL is too long"
    parsed = urlparse(url)
    if parsed.scheme not in {"http", "https"} or not parsed.netloc:
        return "URL must start with http:// or https://"
    return None


def _validate_alias(alias: str, config: dict[str, Any]) -> Optional[str]:
    if not alias:
        return None
    min_len = config["MIN_ALIAS_LENGTH"]
    max_len = config["MAX_ALIAS_LENGTH"]
    if len(alias) < min_len or len(alias) > max_len:
        return f"Alias must be {min_len}-{max_len} characters"
    if alias in RESERVED_ALIASES:
        return "Alias is reserved"
    if not ALIAS_REGEX.match(alias):
        return "Alias can only contain letters, numbers, - and _"
    return None


def _alias_exists(db: sqlite3.Connection, alias: str) -> bool:
    row = db.execute("SELECT 1 FROM links WHERE alias = ?", (alias,)).fetchone()
    return row is not None


def _insert_link(
    db: sqlite3.Connection, alias: str, url: str, expires_at: Optional[str]
) -> None:
    db.execute(
        """
        INSERT INTO links (alias, url, created_at, expires_at, usage_count)
        VALUES (?, ?, ?, ?, 0)
        """,
        (alias, url, datetime.now(timezone.utc).isoformat(), expires_at),
    )
    db.commit()


def _get_link(db: sqlite3.Connection, alias: str) -> Optional[sqlite3.Row]:
    return db.execute(
        "SELECT alias, url, created_at, expires_at, usage_count FROM links WHERE alias = ?",
        (alias,),
    ).fetchone()


def _increment_usage(db: sqlite3.Connection, alias: str) -> None:
    db.execute(
        "UPDATE links SET usage_count = usage_count + 1 WHERE alias = ?",
        (alias,),
    )
    db.commit()


def _is_json_request(request: Request) -> bool:
    mimetype = request.headers.get("content-type", "").split(";")[0].strip().lower()
    return mimetype == "application/json" or (
        mimetype.startswith("application/") and mimetype.endswith("+json")
    )


class _JSONResponse(JSONResponse):
    """Serialises like the previous JSON helper: sorted keys, compact, newline."""

    def render(self, content: Any) -> bytes:
        return (
            json.dumps(content, sort_keys=True, separators=(",", ":")) + "\n"
        ).encode("utf-8")


def _allowed_methods(app: FastAPI, request: Request) -> set[str]:
    """Methods accepted for the requested path, as Werkzeug reports them."""
    scope = dict(request.scope)
    methods: set[str] = set()
    for route in app.routes:
        match, _ = route.matches(scope)
        if match is Match.NONE:
            continue
        # Mounted apps (StaticFiles) expose no `methods`; they serve GET/HEAD.
        route_methods = getattr(route, "methods", None)
        methods |= set(route_methods) if route_methods else {"GET", "HEAD"}
    if methods:
        methods.add("OPTIONS")
        if "GET" in methods:
            methods.add("HEAD")
    return methods


def create_app(test_config: Optional[dict] = None) -> FastAPI:
    config: dict[str, Any] = {
        "DATABASE_PATH": os.getenv(
            "SHARPLINK_DB_PATH", str(DATA_PATH / "url_store.db")
        ),
        "MAX_URL_LENGTH": 2048,
        "MIN_ALIAS_LENGTH": 3,
        "MAX_ALIAS_LENGTH": 32,
        "MAX_CONTENT_LENGTH": 16 * 1024,
        "RATELIMIT_STORAGE_URI": os.getenv("RATELIMIT_STORAGE_URI", "memory://"),
        "RATELIMIT_ENABLED": os.getenv("RATELIMIT_ENABLED", "true").lower()
        not in {"0", "false", "no"},
    }

    if test_config:
        config.update(test_config)

    if config.get("TESTING"):
        config.setdefault("RATELIMIT_ENABLED", False)
        config.setdefault("RATELIMIT_STORAGE_URI", "memory://")

    app = FastAPI(
        docs_url=None,
        redoc_url=None,
        openapi_url=None,
        # Werkzeug 404s on a trailing slash the rule does not declare; Starlette
        # would answer 307 instead.
        redirect_slashes=False,
    )
    app.state.config = config

    limiter = Limiter(
        key_func=get_remote_address,
        default_limits=["200 per day", "60 per hour"],
        storage_uri=config["RATELIMIT_STORAGE_URI"],
        enabled=bool(config["RATELIMIT_ENABLED"]),
    )
    # slowapi re-reads RATELIMIT_ENABLED from the environment inside Limiter()
    # and keeps it as a raw string, so "false" would be truthy. Restore the
    # boolean parsed above so the flag behaves as it did before.
    limiter.enabled = bool(config["RATELIMIT_ENABLED"])
    app.state.limiter = limiter
    app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)
    app.add_middleware(SlowAPIMiddleware)

    app.mount("/static", StaticFiles(directory=str(BASE_DIR / "static")), name="static")

    @app.middleware("http")
    async def method_negotiation(request: Request, call_next):
        allowed = _allowed_methods(app, request)
        if request.method == "OPTIONS" and allowed:
            return Response(
                status_code=200,
                headers={"Allow": ", ".join(sorted(allowed))},
                media_type="text/html; charset=utf-8",
            )
        response = await call_next(request)
        if response.status_code == 405 and allowed:
            response.headers["Allow"] = ", ".join(sorted(allowed))
        return response

    @app.middleware("http")
    async def enforce_max_content_length(request: Request, call_next):
        content_length = request.headers.get("content-length")
        if content_length and content_length.isdigit():
            if int(content_length) > config["MAX_CONTENT_LENGTH"]:
                return _JSONResponse(
                    {"error": "Request entity too large"}, status_code=413
                )
        return await call_next(request)

    init_db(config["DATABASE_PATH"])
    bootstrap_db = _connect(config["DATABASE_PATH"])
    try:
        _migrate_json_to_db(bootstrap_db)
    finally:
        bootstrap_db.close()

    @app.api_route("/", methods=["GET", "HEAD"], response_class=HTMLResponse)
    def index(request: Request):
        return templates.TemplateResponse(request, "index.html")

    @app.post("/api/shorten")
    @limiter.limit("10 per minute")
    async def shorten(request: Request, db: sqlite3.Connection = Depends(get_db)):
        if not _is_json_request(request):
            return _JSONResponse({"error": "Request must be JSON"}, status_code=400)

        try:
            payload = await request.json()
        except (json.JSONDecodeError, UnicodeDecodeError, ValueError):
            payload = None
        if not isinstance(payload, dict):
            payload = {}

        original_url = (payload.get("url") or "").strip()
        custom_alias = (payload.get("alias") or "").strip()
        expires_at_raw = (payload.get("expires_at") or "").strip()
        use_tiny = bool(payload.get("use_tiny"))

        url_error = _validate_url(original_url, config)
        if url_error:
            return _JSONResponse({"error": url_error}, status_code=400)

        alias_error = _validate_alias(custom_alias, config)
        if alias_error:
            return _JSONResponse({"error": alias_error}, status_code=400)

        alias = custom_alias or _generate_alias(db)
        if _alias_exists(db, alias):
            return _JSONResponse(
                {"error": "Alias already exists, choose another."}, status_code=409
            )

        try:
            expires_at = _parse_expiry(expires_at_raw)
        except ValueError as exc:
            return _JSONResponse({"error": str(exc)}, status_code=400)

        try:
            _insert_link(db, alias, original_url, expires_at)
        except sqlite3.IntegrityError:
            return _JSONResponse(
                {"error": "Alias already exists, choose another."}, status_code=409
            )

        short_url = _build_short_url(request, alias)
        tiny_url = _create_tiny_url(original_url) if use_tiny else None

        return _JSONResponse(
            {
                "alias": alias,
                "short_url": short_url,
                "expires_at": expires_at,
                "tiny_url": tiny_url,
            }
        )

    @app.api_route("/health", methods=["GET", "HEAD"])
    @limiter.exempt
    def health():
        return _JSONResponse({"status": "ok"})

    @app.api_route("/{alias}", methods=["GET", "HEAD"])
    def redirect_alias(
        request: Request, alias: str, db: sqlite3.Connection = Depends(get_db)
    ):
        entry = _get_link(db, alias)
        if not entry:
            return templates.TemplateResponse(
                request,
                "not_found.html",
                {"message": "Short link not found."},  # pragma: no cover
                status_code=404,
            )

        expires_at = entry["expires_at"]
        if expires_at:
            expiry_dt = datetime.fromisoformat(expires_at)
            if expiry_dt <= datetime.now(timezone.utc):
                return templates.TemplateResponse(
                    request,
                    "expired.html",
                    {"message": "This link has expired."},  # pragma: no cover
                    status_code=410,
                )

        _increment_usage(db, alias)
        return RedirectResponse(entry["url"], status_code=302)

    return app


def _generate_alias(db: sqlite3.Connection, length: int = 6) -> str:
    while True:
        alias = uuid.uuid4().hex[:length]
        if not _alias_exists(db, alias):
            return alias


app = create_app()


if __name__ == "__main__":
    import uvicorn

    port = int(os.getenv("PORT", "5000"))
    uvicorn.run(app, host="0.0.0.0", port=port)