api.py

from fastapi import FastAPI, BackgroundTasks, Query
from fastapi.middleware.cors import CORSMiddleware as CORS
from fastapi.responses import JSONResponse, FileResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from uuid import uuid4
from typing import Dict, Any, List, Optional
import time
import os
import threading
import requests
import json
import logging

from src.crawler import Crawler
from src.analyzer import ImageAnalyzer
from src.monitor import SystemMonitor

# Configure Logger
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger("API")

app = FastAPI(title="Antigravity SEO Scaler API")

app.add_middleware(
    CORS,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Mount Static Files
if not os.path.exists("static"):
    os.makedirs("static")
app.mount("/static", StaticFiles(directory="static"), name="static")

# In-memory Job Store
JOBS: Dict[str, Dict[str, Any]] = {}
REPORT_FILE = 'seo_report.json'

# --- Models ---
class StartReq(BaseModel):
    domain: str
    limit: int = 25

# --- Background Infrastructure ---

class KeepAlive(threading.Thread):
    def __init__(self, interval=60, target_url="http://127.0.0.1:7860/health"):
        super().__init__()
        self.interval = interval
        self.target_url = target_url
        self.daemon = True
        self.running = True

    def run(self):
        logger.info("KeepAlive System Started")
        while self.running:
            try:
                logger.info(f"[Heartbeat] System Active - {time.ctime()}")
                time.sleep(self.interval)
                try:
                    requests.get(self.target_url, timeout=5)
                except:
                    pass
            except Exception as e:
                logger.error(f"[KeepAlive] Error: {e}")
                time.sleep(60)

class ScheduledCrawler(threading.Thread):
    def __init__(self, interval=14400): # 4 Hours
        super().__init__()
        self.interval = interval
        self.daemon = True
        self.running = True

    def run(self):
        logger.info(f"[Scheduler] Auto-Crawler initialized. Schedule: Every {self.interval/3600} hours.")
        time.sleep(60) 
        
        while self.running:
            try:
                target_domain = os.environ.get("AUTO_CRAWL_TARGET") 
                if target_domain:
                    logger.info(f"[Scheduler] 🕒 Triggering scheduled crawl for: {target_domain}")
                    
                    # Create a pseudo-job for tracking
                    job_id = f"auto-{int(time.time())}"
                    run_scan_job(job_id, target_domain, 50, is_auto=True)
                else:
                    logger.info("[Scheduler] ℹ️ waiting... (Set 'AUTO_CRAWL_TARGET' Env Var to enable auto-crawling)")

                time.sleep(self.interval)

            except Exception as e:
                logger.error(f"[Scheduler] Error: {e}")
                time.sleep(60)

# --- Startup Event ---
@app.on_event("startup")
async def startup_event():
    logger.info("--- API STARTUP: Initializing Services ---")
    try:
        # Check permissions
        if not os.path.exists("static"):
            logger.info("Creating static directory...")
            os.makedirs("static", exist_ok=True)
            
        # Start Background Threads
        logger.info("Starting background services...")
        pinger = KeepAlive(interval=300)
        pinger.start()

        scheduler = ScheduledCrawler(interval=14400)
        scheduler.start()
        logger.info("--- API STARTUP: Ready ---")
    except Exception as e:
        logger.error(f"FATAL STARTUP ERROR: {e}")

# --- Core Logic ---

def run_scan_job(job_id: str, domain: str, limit: int, is_auto: bool = False):
    try:
        if job_id not in JOBS:
            JOBS[job_id] = {} # Should be initialized by caller, but safety check

        JOBS[job_id].update({
            "status": "running",
            "percent": 0,
            "message": "Initializing Crawler...",
            "start_time": time.time(),
            "error": None
        })

        def progress_callback(pages_scanned, images_found, current_url):
            # Update Job State
            JOBS[job_id]["pages_scanned"] = pages_scanned
            JOBS[job_id]["images_found"] = images_found
            JOBS[job_id]["message"] = f"Scanning: {current_url}"
            
            # Estimate percentage (capped at 90% during crawl)
            if limit > 0:
                pct = int((pages_scanned / limit) * 90)
                JOBS[job_id]["percent"] = min(pct, 90)
        
        # 1. Crawl
        crawler = Crawler()
        site_data, total_discovered, blocked_reason = crawler.crawl_domain(
            domain, 
            max_pages=limit, 
            progress_callback=progress_callback
        )
        
        if not site_data:
            JOBS[job_id]["status"] = "error"
            JOBS[job_id]["error"] = f"Scan failed: {blocked_reason or 'No pages found (Access Denied or JS Blocked)'}"
            return

        JOBS[job_id]["message"] = "Analyzing Images..."
        JOBS[job_id]["percent"] = 95

        # 2. Analyze
        analyzer = ImageAnalyzer()
        results = analyzer.analyze_site(site_data)

        # Add metadata
        results['summary']['total_pages_discovered'] = total_discovered
        results['summary']['blocked_reason'] = blocked_reason
        results['summary']['crawl_blocked'] = bool(blocked_reason)
        
        # 3. Save/Finish
        JOBS[job_id]["result"] = results
        JOBS[job_id]["status"] = "done"
        JOBS[job_id]["percent"] = 100
        JOBS[job_id]["message"] = "Completed"

        # If auto-crawl, save to file as well for legacy compatibility
        if is_auto:
            with open(REPORT_FILE, 'w', encoding='utf-8') as f:
                json.dump(results, f, indent=4)
            logger.info(f"[Scheduler] ✅ Crawl finished for {domain}. Report saved.")

    except Exception as e:
        logger.error(f"Job {job_id} failed: {e}")
        JOBS[job_id]["status"] = "error"
        JOBS[job_id]["error"] = str(e)
        JOBS[job_id]["message"] = "Internal Error"

# --- Endpoints ---

@app.get("/")
def home():
    # Serve the static UI by default
    if os.path.exists("static/index.html"):
        return FileResponse("static/index.html")
    return {"message": "Antigravity API v2.5 Running", "timestamp": "2026-02-11"}

@app.get("/health")
def health_check():
    return {"status": "alive"}

@app.get("/api/status")
def get_system_status(domain: Optional[str] = None):
    # If a domain is provided, we check its connectivity specifically
    stats = SystemMonitor.get_system_stats(target_url=domain)
    return stats

@app.post("/api/scanstart")
def start_scan(
    bg: BackgroundTasks,
    payload: Optional[StartReq] = None,
    domain: Optional[str] = Query(None),
    limit: Optional[int] = Query(25)
):
    # Support both JSON body and query parameters
    if payload:
        target_domain = payload.domain
        target_limit = payload.limit
    elif domain:
        target_domain = domain
        target_limit = limit
    else:
        return JSONResponse(
            status_code=400, 
            content={"error": "Must provide either JSON body or query parameters (domain required)"}
        )
    
    # Ensure domain has scheme
    if target_domain and not target_domain.startswith(('http://', 'https://')):
        target_domain = f"https://{target_domain}"
    
    job_id = str(uuid4())
    JOBS[job_id] = {
        "status": "pending",
        "percent": 0,
        "pages_scanned": 0,
        "images_found": 0,
        "message": "Queued",
        "start_time": time.time(),
        "elapsed": 0,
        "result": None,
        "error": None,
    }
    bg.add_task(run_scan_job, job_id, target_domain, target_limit)
    return {"job_id": job_id}

@app.get("/api/progress/{job_id}")
@app.get("/api/progress")
def scan_progress(job_id: Optional[str] = None):
    if not job_id:
         return JSONResponse(status_code=400, content={"error": "job_id is required"})
    
    job = JOBS.get(job_id)
    if not job:
        return JSONResponse(status_code=404, content={"status": "not_found", "error": f"Job ID {job_id} not found"})
    
    # Calculate Live Stats
    current_time = time.time()
    elapsed = int(current_time - job.get("start_time", current_time))
    
    # Simple ETA Calculation
    eta_seconds = None
    pages_scanned = job.get("pages_scanned", 0)
    percent = job.get("percent", 0)
    
    if job["status"] == "running" and pages_scanned > 0 and percent > 0:
        avg_per_page = elapsed / pages_scanned
        estimated_total = (pages_scanned / percent) * 100
        remaining = estimated_total - pages_scanned
        eta_seconds = int(remaining * avg_per_page)

    return {
        "status": job["status"],
        "percent": job["percent"],
        "pages_scanned": pages_scanned,
        "images_found": job.get("images_found", 0),
        "message": job.get("message", ""),
        "elapsed_seconds": elapsed,
        "eta_seconds": eta_seconds,
        "error": job.get("error"),
    }

@app.get("/api/result/{job_id}")
@app.get("/api/result")
def scan_result(job_id: Optional[str] = None):
    if not job_id:
         return JSONResponse(status_code=400, content={"error": "job_id is required"})
         
    job = JOBS.get(job_id)
    if not job:
        return JSONResponse(status_code=404, content={"status": "not_found", "error": f"Job ID {job_id} not found"})
    
    if job["status"] != "done":
        return {"status": job["status"], "message": "Result not ready yet"}
        
    return job["result"]

# Legacy Endpoint Support (Optional - redirects to a sync wait or error?)
# For now, let's keep it but make it use the new logic synchronously if possible, 
# OR just deprecate it. Given the request, we should probably stick to the new API.
# But for backward compatibility with existing frontend using /api/seo-report?
# The user didn't ask to remove it, but the new code replaces the structure.
# I'll enable a blocking legacy endpoint for safety.

@app.get("/api/seo-report")
def get_seo_report_legacy_get(
    domain: Optional[str] = Query(None), 
    limit: Optional[int] = Query(25)
):
    return handle_legacy_request(domain, limit)

@app.post("/api/seo-report")
def get_seo_report_legacy_post(payload: StartReq):
    return handle_legacy_request(payload.domain, payload.limit)

def handle_legacy_request(domain: Optional[str], limit: int):
    """
    Simulates the old blocking behavior by starting a job and waiting for it.
    Note: This might timeout on some clients if the scan is long.
    """
    if not domain:
        # Attempt to read cached default report if no domain
        if os.path.exists(REPORT_FILE):
             try:
                 with open(REPORT_FILE, 'r', encoding='utf-8') as f:
                     return json.load(f)
             except:
                 pass
        return JSONResponse(status_code=404, content={"error": "No domain provided and no cached report found."})

    # Start Job
    job_id = str(uuid4())
    run_scan_job(job_id, domain, limit, is_auto=False)
    
    # Check Result
    job = JOBS.get(job_id)
    if job and job["status"] == "done":
        return job["result"]
    else:
        return JSONResponse(status_code=500, content={"error": job.get("error", "Unknown error during scan")})

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)