#!/usr/bin/env python3
"""
URL Shortener — stdlib-only, Python 3.8+
Usage: python urlshort.py <config.json>

API (all routes are prefixed with base_path, e.g. /s):
  GET  /                        Frontend (or health check if no static/)
  GET  /api/health              Health check
  POST /api/shorten             Create a short URL  (API key required)
  GET  /api/urls                List all short URLs  (API key required)
  GET  /api/urls/<code>         Get info for a code  (no API key required)
  GET  /api/lookup?url=<url>    Look up by original URL (no API key required)
  DELETE /api/urls/<code>       Delete a short URL   (API key required)
  GET  /<code>                  Redirect to original URL
"""

import json
import logging
import os
import random
import sqlite3
import string
import sys
import time
import collections
import mimetypes
import re
import threading
from http.server import BaseHTTPRequestHandler, HTTPServer
from urllib.parse import parse_qs, urlparse

# ---------------------------------------------------------------------------
# Logging — outputs to console (stdout); Docker captures it automatically.
# ---------------------------------------------------------------------------

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)s %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
    stream=sys.stdout,
)
log = logging.getLogger("urlshort")


# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------

def load_config(path: str) -> dict:
    with open(path, "r", encoding="utf-8") as fh:
        cfg = json.load(fh)

    for key in ("base_url", "api_key"):
        if key not in cfg:
            raise ValueError(f"Missing required config key: '{key}'")

    # Backward compat: old 'short_length' → 'min_short_length'
    if "short_length" in cfg and "min_short_length" not in cfg:
        cfg["min_short_length"] = cfg["short_length"]

    cfg.setdefault("host", "0.0.0.0")
    cfg.setdefault("port", 8080)
    cfg.setdefault("db_path", "data/urlshort.db")
    cfg.setdefault("retention_days", 0)
    cfg.setdefault("min_short_length", 6)
    cfg.setdefault("max_short_length", 32)
    cfg.setdefault("max_url_length", 2048)
    cfg.setdefault("max_retention_days", 3650)
    cfg.setdefault("rate_limit_requests", 60)
    cfg.setdefault("rate_limit_window", 60)
    cfg.setdefault("production", False)

    # Derive base_path from the path component of base_url.
    raw = urlparse(cfg["base_url"]).path.strip("/")
    cfg["base_path"] = f"/{raw}" if raw else ""

    return cfg


# ---------------------------------------------------------------------------
# Database
# ---------------------------------------------------------------------------

def init_db(db_path: str) -> None:
    parent = os.path.dirname(db_path)
    if parent:
        os.makedirs(parent, exist_ok=True)

    with sqlite3.connect(db_path) as conn:
        conn.execute("""
            CREATE TABLE IF NOT EXISTS urls (
                short_code     TEXT    PRIMARY KEY,
                original_url   TEXT    NOT NULL,
                created_at     INTEGER NOT NULL,
                visit_count    INTEGER NOT NULL DEFAULT 0,
                retention_days INTEGER NOT NULL DEFAULT 0
            )
        """)
        conn.commit()

        # Migrate: add retention_days column if missing (existing DB).
        try:
            conn.execute(
                "ALTER TABLE urls ADD COLUMN retention_days INTEGER NOT NULL DEFAULT 0"
            )
            conn.commit()
        except sqlite3.OperationalError:
            pass  # column already exists


def db_connect(db_path: str) -> sqlite3.Connection:
    conn = sqlite3.connect(db_path)
    conn.row_factory = sqlite3.Row
    return conn


def cleanup_expired(db_path: str) -> None:
    """Delete URLs whose retention period has elapsed."""
    now = int(time.time())
    with db_connect(db_path) as conn:
        conn.execute(
            "DELETE FROM urls WHERE retention_days > 0 "
            "AND (created_at + retention_days * 86400) < ?",
            (now,),
        )
        conn.commit()


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

_CHARS = string.ascii_letters + string.digits


def generate_code(length: int) -> str:
    return "".join(random.choices(_CHARS, k=length))


def is_valid_url(url: str) -> bool:
    try:
        p = urlparse(url)
        return p.scheme in ("http", "https") and bool(p.netloc)
    except Exception:
        return False


# ---------------------------------------------------------------------------
# Validation helpers
# ---------------------------------------------------------------------------

_CODE_RE = re.compile(r'^[A-Za-z0-9]+$')


def is_valid_code(code: str, max_length: int = 32) -> bool:
    """Short codes must be alphanumeric and within length limits."""
    return bool(code) and len(code) <= max_length and bool(_CODE_RE.match(code))


# ---------------------------------------------------------------------------
# Static file serving
# ---------------------------------------------------------------------------

STATIC_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "static")


# ---------------------------------------------------------------------------
# Rate limiter (sliding window, per IP)
# ---------------------------------------------------------------------------

class RateLimiter:
    """Simple in-memory sliding-window rate limiter."""

    def __init__(self, max_requests: int = 60, window: int = 60):
        self.max_requests = max_requests
        self.window = window
        self._hits: dict = collections.defaultdict(list)
        self._lock = threading.Lock()

    def is_allowed(self, ip: str) -> bool:
        now = time.time()
        cutoff = now - self.window
        with self._lock:
            hits = self._hits[ip]
            self._hits[ip] = hits = [t for t in hits if t > cutoff]
            if len(hits) >= self.max_requests:
                return False
            hits.append(now)
            return True


_rate_limiter = RateLimiter()


# ---------------------------------------------------------------------------
# HTTP Handler
# ---------------------------------------------------------------------------

class Handler(BaseHTTPRequestHandler):
    """Single handler that serves the whole URL shortener API."""

    # Injected by main() before the server starts.
    cfg: dict = {}

    # ------------------------------------------------------------------
    # Routing helpers
    # ------------------------------------------------------------------

    def _client_ip(self) -> str:
        """Get the real client IP, checking reverse-proxy headers first.

        nginx adds X-Real-IP and X-Forwarded-For via proxy_set_header.
        Without these headers, falls back to the TCP connection source.
        """
        ip = self.headers.get("X-Real-IP", "").strip()
        if ip:
            return ip
        xff = self.headers.get("X-Forwarded-For", "").strip()
        if xff:
            return xff.split(",")[0].strip()
        return self.client_address[0]

    def _local_path(self):
        """Return the request path with base_path prefix stripped.

        Returns None (→ 404) when the request path does not start with
        the configured base_path at all.
        """
        raw = urlparse(self.path).path
        base = self.cfg.get("base_path", "")
        if base:
            if raw == base or raw == base + "/":
                # exact match on the prefix itself → treat as root
                return "/"
            if raw.startswith(base + "/"):
                return raw[len(base):].rstrip("/") or "/"
            # path is outside our prefix entirely
            return None
        return raw.rstrip("/") or "/"

    # ------------------------------------------------------------------
    # Logging
    # ------------------------------------------------------------------

    def log_message(self, fmt, *args):  # noqa: N802 – stdlib override
        log.info("%s - %s", self._client_ip(), fmt % args)

    # ------------------------------------------------------------------
    # Low-level response helpers
    # ------------------------------------------------------------------

    def _send_json(self, status: int, payload: object) -> None:
        body = json.dumps(payload, ensure_ascii=False).encode("utf-8")
        self.send_response(status)
        self.send_header("Content-Type", "application/json; charset=utf-8")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def _send_plain(self, status: int, text: str) -> None:
        body = text.encode("utf-8")
        self.send_response(status)
        self.send_header("Content-Type", "text/plain; charset=utf-8")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def _send_empty(self, status: int) -> None:
        self.send_response(status)
        self.send_header("Content-Length", "0")
        self.end_headers()

    def _error(self, status: int, message: str) -> None:
        self._send_json(status, {"error": message})

    def _redirect(self, location: str) -> None:
        self.send_response(302)
        self.send_header("Location", location)
        self.send_header("Content-Length", "0")
        self.end_headers()

    # ------------------------------------------------------------------
    # CORS & Security headers (injected into every response)
    # ------------------------------------------------------------------

    def end_headers(self):
        self.send_header("Access-Control-Allow-Origin", "*")
        self.send_header("Access-Control-Allow-Methods", "GET, POST, DELETE, OPTIONS")
        self.send_header("Access-Control-Allow-Headers", "Content-Type")
        self.send_header("Access-Control-Max-Age", "86400")
        self.send_header("X-Content-Type-Options", "nosniff")
        self.send_header("X-Frame-Options", "DENY")
        self.send_header("X-XSS-Protection", "1; mode=block")
        super().end_headers()

    # ------------------------------------------------------------------
    # Auth — purely api_key based (query param or JSON body)
    # ------------------------------------------------------------------

    def _check_api_key(self, body: dict = None) -> bool:
        """Check api_key from query parameter or JSON body."""
        qs = parse_qs(urlparse(self.path).query)
        key = qs.get("api_key", [""])[0]
        if key:
            return key == self.cfg["api_key"]
        if body and isinstance(body, dict):
            key = str(body.get("api_key", ""))
            if key:
                return key == self.cfg["api_key"]
        return False

    # ------------------------------------------------------------------
    # Body
    # ------------------------------------------------------------------

    def _read_json(self):
        length = int(self.headers.get("Content-Length", 0))
        if length == 0:
            return {}
        raw = self.rfile.read(length)
        try:
            return json.loads(raw)
        except json.JSONDecodeError:
            return None  # signals parse failure to caller

    # ------------------------------------------------------------------
    # Route dispatch
    # ------------------------------------------------------------------

    def do_GET(self):  # noqa: N802
        cleanup_expired(self.cfg["db_path"])

        if not _rate_limiter.is_allowed(self._client_ip()):
            self._error(429, "Too many requests")
            return

        # Redirect bare base path to base path + / for correct relative URLs
        raw = urlparse(self.path).path
        base = self.cfg.get("base_path", "")
        if base and raw == base:
            self._redirect(base + "/")
            return

        path = self._local_path()
        if path is None:
            self._send_empty(404)
            return

        if path == "/":
            # In production, nginx serves static files; backend only serves health check.
            index_path = os.path.join(STATIC_DIR, "index.html")
            if not self.cfg.get("production") and os.path.isfile(index_path):
                self._serve_static("index.html")
            else:
                self._send_json(200, {"status": "ok", "service": "url-shortener"})

        elif path == "/api/health":
            self._send_json(200, {"status": "ok", "service": "url-shortener"})

        elif path == "/api/urls":
            self._handle_list_urls()

        elif path.startswith("/api/urls/"):
            code = path[len("/api/urls/"):]
            self._handle_get_url(code)

        elif path == "/api/lookup":
            self._handle_lookup()

        elif path == "/static" or path.startswith("/static/"):
            if self.cfg.get("production"):
                self._send_empty(404)
                return
            rel = path[len("/static"):].lstrip("/") or "index.html"
            self._serve_static(rel)

        else:
            code = path.lstrip("/")
            self._handle_redirect(code)

    def do_POST(self):  # noqa: N802
        cleanup_expired(self.cfg["db_path"])

        if not _rate_limiter.is_allowed(self._client_ip()):
            self._error(429, "Too many requests")
            return

        path = self._local_path()
        if path is None:
            self._send_empty(404)
            return

        if path == "/api/shorten":
            self._handle_shorten()
        else:
            self._send_empty(404)

    def do_DELETE(self):  # noqa: N802
        cleanup_expired(self.cfg["db_path"])

        if not _rate_limiter.is_allowed(self._client_ip()):
            self._error(429, "Too many requests")
            return

        path = self._local_path()
        if path is None:
            self._send_empty(404)
            return

        if path.startswith("/api/urls/"):
            code = path[len("/api/urls/"):]
            self._handle_delete_url(code)
        else:
            self._send_empty(404)

    def do_OPTIONS(self):  # noqa: N802
        """Handle CORS preflight requests."""
        self.send_response(200)
        self.send_header("Content-Length", "0")
        self.send_header("Allow", "GET, POST, DELETE, OPTIONS")
        self.end_headers()
        self.wfile.flush()

    # ------------------------------------------------------------------
    # Handlers
    # ------------------------------------------------------------------

    def _handle_shorten(self) -> None:
        # Parse both query parameters and JSON body
        qs = parse_qs(urlparse(self.path).query)
        body = self._read_json()
        if body is None:
            self._error(400, "Invalid JSON body")
            return

        # Auth: check api_key from query param or body
        if not self._check_api_key(body):
            self._send_empty(403)
            return

        # url: query param takes precedence, then body
        original_url = qs.get("url", [""])[0] or str(body.get("url", "")).strip()
        if not original_url:
            self._error(400, "Missing required field: url")
            return
        if not is_valid_url(original_url):
            self._error(400, "Invalid URL — must start with http:// or https://")
            return
        max_url = self.cfg["max_url_length"]
        if len(original_url) > max_url:
            self._error(400, f"URL too long (max {max_url} characters)")
            return

        # retention_days: query param, then body, then config default
        raw_rd = qs.get("retention_days", [None])[0]
        if raw_rd is None:
            raw_rd = body.get("retention_days")
        if raw_rd is not None:
            try:
                retention_days = int(raw_rd)
            except (ValueError, TypeError):
                self._error(400, "Invalid retention_days")
                return
        else:
            retention_days = self.cfg.get("retention_days", 0)

        max_ret = self.cfg["max_retention_days"]
        if not (0 <= retention_days <= max_ret):
            self._error(400, f"retention_days must be 0–{max_ret}")
            return

        with db_connect(self.cfg["db_path"]) as conn:
            short_code = self._unique_code(conn)
            if short_code is None:
                self._error(500, "Could not generate a unique short code — try again")
                return

            created_at = int(time.time())
            conn.execute(
                "INSERT INTO urls (short_code, original_url, created_at, retention_days) "
                "VALUES (?, ?, ?, ?)",
                (short_code, original_url, created_at, retention_days),
            )
            conn.commit()

        short_url = f"{self.cfg['base_url'].rstrip('/')}/{short_code}"
        # Response: only the short URL in plain text
        self._send_plain(201, short_url)

    def _handle_list_urls(self) -> None:
        if not self._check_api_key():
            self._send_empty(403)
            return

        with db_connect(self.cfg["db_path"]) as conn:
            rows = conn.execute(
                "SELECT * FROM urls ORDER BY created_at DESC"
            ).fetchall()

        base = self.cfg["base_url"].rstrip("/")
        urls = []
        for r in rows:
            urls.append({
                "short_code": r["short_code"],
                "short_url": f"{base}/{r['short_code']}",
                "original_url": r["original_url"],
                "created_at": r["created_at"],
                "visit_count": r["visit_count"],
                "retention_days": r["retention_days"],
            })
        self._send_json(200, {"count": len(urls), "urls": urls})

    def _handle_get_url(self, code: str) -> None:
        # No API key required
        if not is_valid_code(code, self.cfg["max_short_length"]):
            self._send_empty(404)
            return

        with db_connect(self.cfg["db_path"]) as conn:
            row = conn.execute(
                "SELECT * FROM urls WHERE short_code = ?", (code,)
            ).fetchone()

        if row is None:
            self._send_empty(404)
            return

        base = self.cfg["base_url"].rstrip("/")
        self._send_json(200, {
            "short_code": row["short_code"],
            "short_url": f"{base}/{row['short_code']}",
            "original_url": row["original_url"],
            "created_at": row["created_at"],
            "visit_count": row["visit_count"],
            "retention_days": row["retention_days"],
        })

    def _handle_delete_url(self, code: str) -> None:
        # API key required — 403 if not authorized (no body)
        if not self._check_api_key():
            self._send_empty(403)
            return
        if not is_valid_code(code, self.cfg["max_short_length"]):
            self._send_empty(404)
            return

        with db_connect(self.cfg["db_path"]) as conn:
            row = conn.execute(
                "SELECT short_code FROM urls WHERE short_code = ?", (code,)
            ).fetchone()
            if row is None:
                # 404 if not found (no body)
                self._send_empty(404)
                return
            conn.execute("DELETE FROM urls WHERE short_code = ?", (code,))
            conn.commit()

        # 204 on success (no body)
        self._send_empty(204)

    def _handle_redirect(self, code: str) -> None:
        if not is_valid_code(code, self.cfg["max_short_length"]):
            self._send_empty(404)
            return

        with db_connect(self.cfg["db_path"]) as conn:
            row = conn.execute(
                "SELECT original_url FROM urls WHERE short_code = ?", (code,)
            ).fetchone()
            if row is None:
                self._error(404, "Short code not found")
                return
            conn.execute(
                "UPDATE urls SET visit_count = visit_count + 1 WHERE short_code = ?",
                (code,),
            )
            conn.commit()

        self._redirect(row["original_url"])

    def _handle_lookup(self) -> None:
        """Look up a URL by its original URL. No API key required."""
        qs = parse_qs(urlparse(self.path).query)
        url = qs.get("url", [""])[0].strip()
        if not url:
            self._error(400, "Missing required parameter: url")
            return
        max_url = self.cfg["max_url_length"]
        if len(url) > max_url:
            self._error(400, f"URL too long (max {max_url} characters)")
            return

        with db_connect(self.cfg["db_path"]) as conn:
            row = conn.execute(
                "SELECT * FROM urls WHERE original_url = ? ORDER BY created_at DESC LIMIT 1",
                (url,),
            ).fetchone()

        if row is None:
            self._send_empty(404)
            return

        base = self.cfg["base_url"].rstrip("/")
        self._send_json(200, {
            "short_code": row["short_code"],
            "short_url": f"{base}/{row['short_code']}",
            "original_url": row["original_url"],
            "created_at": row["created_at"],
            "visit_count": row["visit_count"],
            "retention_days": row["retention_days"],
        })

    def _serve_static(self, rel_path: str) -> None:
        """Serve a static file from STATIC_DIR (for local dev; nginx in prod)."""
        if not rel_path:
            rel_path = "index.html"
        # Prevent directory traversal
        safe = os.path.normpath(rel_path)
        if safe.startswith("..") or os.path.isabs(safe):
            self._send_empty(403)
            return
        fpath = os.path.join(STATIC_DIR, safe)
        if not os.path.isfile(fpath):
            self._send_empty(404)
            return
        mime, _ = mimetypes.guess_type(fpath)
        if not mime:
            mime = "application/octet-stream"
        with open(fpath, "rb") as f:
            data = f.read()
        self.send_response(200)
        self.send_header("Content-Type", mime)
        self.send_header("Content-Length", str(len(data)))
        self.send_header("Cache-Control", "public, max-age=3600")
        self.end_headers()
        self.wfile.write(data)

    # ------------------------------------------------------------------
    # Internal utils
    # ------------------------------------------------------------------

    def _unique_code(self, conn: sqlite3.Connection, attempts_per_length: int = 10):
        """Generate a unique short code with progressive length increment.

        Starts at min_short_length and tries `attempts_per_length` random
        codes at each length.  On exhaustion, increments the length by 1
        and repeats, up to max_short_length.  Returns None only when the
        entire range is exhausted (extremely unlikely).
        """
        min_len = self.cfg["min_short_length"]
        max_len = self.cfg["max_short_length"]
        for length in range(min_len, max_len + 1):
            for _ in range(attempts_per_length):
                code = generate_code(length)
                exists = conn.execute(
                    "SELECT 1 FROM urls WHERE short_code = ?", (code,)
                ).fetchone()
                if not exists:
                    return code
        return None


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------

def main() -> None:
    global _rate_limiter

    if len(sys.argv) != 2:
        log.error("Usage: %s <config.json>", sys.argv[0])
        sys.exit(1)

    cfg = load_config(sys.argv[1])
    init_db(cfg["db_path"])
    cleanup_expired(cfg["db_path"])

    # Re-initialize rate limiter with config values
    _rate_limiter = RateLimiter(
        max_requests=cfg["rate_limit_requests"],
        window=cfg["rate_limit_window"],
    )

    Handler.cfg = cfg

    host, port = cfg["host"], int(cfg["port"])
    server = HTTPServer((host, port), Handler)

    log.info("URL Shortener listening on http://%s:%s", host, port)
    log.info("Base URL       : %s", cfg["base_url"])
    log.info("Base path      : %s", cfg["base_path"] or "/")
    log.info("DB path        : %s", cfg["db_path"])
    log.info("Short codes    : %s–%s chars", cfg["min_short_length"], cfg["max_short_length"])
    log.info("Retention days : %s", cfg["retention_days"])
    log.info("Rate limit     : %s req/%ss per IP", cfg["rate_limit_requests"], cfg["rate_limit_window"])
    log.info("Production     : %s", cfg["production"])

    try:
        server.serve_forever()
    except KeyboardInterrupt:
        log.info("Shutting down.")
        server.shutdown()


if __name__ == "__main__":
    main()