DomGod/app/main.py

import os
import sys
import asyncio
import logging
from pathlib import Path
from contextlib import asynccontextmanager

import httpx
import duckdb
import aiosqlite
from fastapi import FastAPI, Query
from fastapi.responses import StreamingResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from dotenv import load_dotenv

load_dotenv()

from app.db import (
    DATA_DIR, PARQUET_PATH, SQLITE_PATH,
    init_db, get_stats, get_domains, get_enriched,
    queue_domains, get_queue_status,
)
from app.enricher import start_worker, pause_worker, resume_worker, is_running
from app.scorer import run_scoring

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger(__name__)

PARQUET_URL = os.getenv("PARQUET_URL", "")


async def download_parquet():
    if PARQUET_PATH.exists():
        logger.info("Using cached parquet at %s", PARQUET_PATH)
        return

    DATA_DIR.mkdir(parents=True, exist_ok=True)
    tmp_path = PARQUET_PATH.with_suffix(".tmp")

    # Resumable download via Range header
    downloaded = tmp_path.stat().st_size if tmp_path.exists() else 0
    headers = {"Range": f"bytes={downloaded}-"} if downloaded > 0 else {}

    logger.info("Downloading parquet from %s (offset=%d)...", PARQUET_URL, downloaded)

    async with httpx.AsyncClient(follow_redirects=True, timeout=None) as client:
        async with client.stream("GET", PARQUET_URL, headers=headers) as resp:
            if resp.status_code == 416:
                # Already fully downloaded
                tmp_path.rename(PARQUET_PATH)
                return
            resp.raise_for_status()
            total = int(resp.headers.get("content-length", 0)) + downloaded
            mode = "ab" if downloaded > 0 else "wb"
            with open(tmp_path, mode) as f:
                received = downloaded
                async for chunk in resp.aiter_bytes(chunk_size=1024 * 1024):
                    f.write(chunk)
                    received += len(chunk)
                    if total:
                        pct = received / total * 100
                        logger.info("Download progress: %.1f%% (%d/%d bytes)", pct, received, total)

    tmp_path.rename(PARQUET_PATH)
    logger.info("Parquet download complete: %s", PARQUET_PATH)


@asynccontextmanager
async def lifespan(app: FastAPI):
    await download_parquet()
    await init_db()
    start_worker()
    logger.info("DomGod dashboard ready on port 6677")
    yield


app = FastAPI(title="DomGod", lifespan=lifespan)


# ── API routes ──────────────────────────────────────────────────────────────

@app.get("/api/stats")
async def stats():
    return await get_stats()


@app.get("/api/domains")
async def domains(
    tld: str = Query(None),
    page: int = Query(1, ge=1),
    limit: int = Query(100, ge=1, le=1000),
    live_only: bool = Query(False),
):
    rows = await get_domains(tld=tld, page=page, limit=limit, live_only=live_only)
    return {"page": page, "limit": limit, "results": rows}


@app.post("/api/enrich/batch")
async def enrich_batch(body: dict):
    domains_list = body.get("domains", [])
    if not domains_list:
        return JSONResponse({"error": "no domains provided"}, status_code=400)
    await queue_domains(domains_list)
    resume_worker()
    return {"queued": len(domains_list)}


@app.get("/api/enrich/status")
async def enrich_status():
    status = await get_queue_status()
    status["worker_running"] = is_running()
    return status


@app.post("/api/enrich/retry")
async def enrich_retry():
    async with aiosqlite.connect(SQLITE_PATH) as db:
        await db.execute("UPDATE job_queue SET status='pending', error=NULL WHERE status='failed'")
        await db.commit()
    resume_worker()
    return {"status": "retrying failed jobs"}


@app.post("/api/enrich/pause")
async def enrich_pause():
    pause_worker()
    return {"status": "paused"}


@app.post("/api/enrich/resume")
async def enrich_resume():
    resume_worker()
    return {"status": "resumed"}


@app.get("/api/enriched")
async def enriched(
    min_score: int = Query(0, ge=0, le=100),
    cms: str = Query(None),
    country: str = Query(None),
    page: int = Query(1, ge=1),
    limit: int = Query(100, ge=1, le=1000),
):
    rows = await get_enriched(min_score=min_score, cms=cms, country=country, page=page, limit=limit)
    return {"page": page, "limit": limit, "results": rows}


@app.get("/api/export")
async def export_csv(
    min_score: int = Query(0),
    cms: str = Query(None),
    country: str = Query(None),
    tier: str = Query(None),
):
    if tier == "hot":
        min_score = 80
    elif tier == "warm":
        min_score = 50

    async def generate():
        yield "domain,score,cms,ssl_expiry_days,ip_country,is_live,status_code,has_mx,server,page_title,enriched_at\n"
        page = 1
        while True:
            rows = await get_enriched(min_score=min_score, cms=cms, country=country, page=page, limit=500)
            if not rows:
                break
            for r in rows:
                # Apply warm tier upper bound
                if tier == "warm" and r.get("score", 0) >= 80:
                    continue
                line = ",".join(
                    f'"{str(r.get(col) or "").replace(chr(34), chr(39))}"'
                    for col in [
                        "domain", "score", "cms", "ssl_expiry_days", "ip_country",
                        "is_live", "status_code", "has_mx", "server", "page_title", "enriched_at"
                    ]
                )
                yield line + "\n"
            page += 1

    filename = f"domgod_leads_score{min_score}{'_' + tier if tier else ''}.csv"
    return StreamingResponse(
        generate(),
        media_type="text/csv",
        headers={"Content-Disposition": f'attachment; filename="{filename}"'},
    )


@app.post("/api/score/run")
async def score_run():
    result = await run_scoring()
    return result


# ── Static UI ───────────────────────────────────────────────────────────────
static_dir = Path(__file__).parent / "static"
app.mount("/", StaticFiles(directory=str(static_dir), html=True), name="static")


if __name__ == "__main__":
    import uvicorn
    uvicorn.run("app.main:app", host="0.0.0.0", port=6677, log_level="info")