DomGod/app/main.py

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

import httpx
import aiosqlite
from typing import Optional
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, build_duckdb_index, index_status,
    queue_ai, get_ai_queue_status, save_ai_assessment,
)
from app.enricher import start_worker, pause_worker, resume_worker, is_running, ensure_workers_alive
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")

    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:
                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:
                        logger.info("Download: %.1f%% (%d/%d)", received / total * 100, received, total)

    tmp_path.rename(PARQUET_PATH)
    logger.info("Parquet download complete")


async def _watchdog():
    """Restart workers if they die every 10 seconds."""
    while True:
        await asyncio.sleep(10)
        ensure_workers_alive()


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


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


# ── API ──────────────────────────────────────────────────────────────────────

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


@app.get("/api/index/status")
async def get_index_status():
    return index_status()


@app.get("/api/domains")
async def domains(
    tld: str = Query(None),
    page: int = Query(1, ge=1),
    limit: int = Query(100, ge=1, le=500),
    live_only: bool = Query(False),
    alpha_only: bool = Query(False),
    no_sld: bool = Query(False),
    keyword: str = Query(None),
):
    total, rows = await get_domains(
        tld=tld, page=page, limit=limit,
        alpha_only=alpha_only, no_sld=no_sld,
        keyword=keyword, live_only=live_only,
    )
    return {"page": page, "limit": limit, "total": total, "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"}


@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),
    kit_digital: Optional[bool] = Query(None),
    page: int = Query(1, ge=1),
    limit: int = Query(100, ge=1, le=1000),
):
    total, rows = await get_enriched(
        min_score=min_score, cms=cms, country=country,
        kit_digital=kit_digital, page=page, limit=limit,
    )
    return {"page": page, "limit": limit, "total": total, "results": rows}


# ── AI assessment endpoints ───────────────────────────────────────────────────

@app.post("/api/ai/assess/batch")
async def ai_assess_batch(body: dict):
    domains_list = body.get("domains", [])
    if not domains_list:
        return JSONResponse({"error": "no domains provided"}, status_code=400)
    await queue_ai(domains_list)
    ensure_workers_alive()   # ensure AI worker is alive when jobs are queued
    return {"queued": len(domains_list)}


@app.post("/api/ai/worker/restart")
async def ai_worker_restart():
    ensure_workers_alive()
    return {"status": "restarted"}


@app.post("/api/ai/reset")
async def ai_reset():
    """Reset all 'running' AI queue jobs back to 'pending' (unstick hung jobs)."""
    async with aiosqlite.connect(SQLITE_PATH) as db:
        r = await db.execute("UPDATE ai_queue SET status='pending' WHERE status='running'")
        count = r.rowcount
        await db.commit()
    ensure_workers_alive()
    return {"reset": count}


@app.get("/api/ai/debug")
async def ai_debug():
    """Returns worker state + last 10 queue entries for troubleshooting."""
    from app.enricher import _ai_worker_task
    task_alive = _ai_worker_task is not None and not _ai_worker_task.done()
    task_exc   = None
    if _ai_worker_task and _ai_worker_task.done() and not _ai_worker_task.cancelled():
        try:
            task_exc = str(_ai_worker_task.exception())
        except Exception:
            pass

    async with aiosqlite.connect(SQLITE_PATH) as db:
        db.row_factory = aiosqlite.Row
        async with db.execute(
            "SELECT domain, status, created_at, completed_at, error FROM ai_queue ORDER BY created_at DESC LIMIT 10"
        ) as cur:
            recent = [dict(r) async for r in cur]

    return {
        "ai_worker_alive": task_alive,
        "ai_worker_exception": task_exc,
        "recent_queue": recent,
        "queue_status": await get_ai_queue_status(),
    }


@app.get("/api/ai/status")
async def ai_status():
    return await get_ai_queue_status()


@app.post("/api/ai/assess/single")
async def ai_assess_single(body: dict):
    """Immediate (blocking) AI assessment — does fresh scrape, no pre-enrichment needed."""
    domain = body.get("domain")
    if not domain:
        return JSONResponse({"error": "no domain"}, status_code=400)
    from app.site_analyzer import analyze_site
    from app.replicate_ai import assess_domain as gemini_assess
    analysis = await analyze_site(domain)
    assessment = await gemini_assess(analysis)
    await save_ai_assessment(domain, assessment, site_analysis=analysis)
    return {**assessment, "site_analysis": analysis}


@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

    max_score = 79 if tier == "warm" else 100

    async def generate():
        yield "domain,score,cms,ssl_expiry_days,ip_country,is_live,status_code,has_mx,server,page_title,enriched_at\n"
        p = 1
        while True:
            _, rows = await get_enriched(min_score=min_score, cms=cms, country=country, page=p, limit=500)
            if not rows:
                break
            for r in rows:
                if r.get("score", 0) > max_score:
                    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"
            p += 1

    fname = f"domgod_{tier or 'export'}_score{min_score}.csv"
    return StreamingResponse(
        generate(), media_type="text/csv",
        headers={"Content-Disposition": f'attachment; filename="{fname}"'},
    )


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


# ── 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")