app.py

"""
Wellfound AI - Main FastAPI Application

Automated Excel data completion tool with:
- AI web content understanding
- Smart address processing
- Multi-source contact finding
- Breakpoint resume support
- Smart timeout detection with auto-skip (5-min threshold per company)
"""

import asyncio
import hashlib
import json
import logging
import os
import time
import uuid
from datetime import datetime
from pathlib import Path
from typing import Optional

import pandas as pd
from fastapi import FastAPI, File, UploadFile, Form, HTTPException, Request
from fastapi.responses import FileResponse, HTMLResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from fastapi.middleware.cors import CORSMiddleware

from core.excel_handler import ExcelHandler
from core.scraper import WebScraper
from core.ai_extractor import AIExtractor
from core.address_processor import AddressProcessor
from core.contact_finder import ContactFinder

# ─── Logging Setup ────────────────────────────────────────
logger = logging.getLogger("wellfound_ai")
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s - %(message)s",
)

# ─── Timeout Configuration ────────────────────────────────
# Per-company processing timeout in seconds (5 minutes)
COMPANY_TIMEOUT_SECONDS = 300

# ─── App Setup ───────────────────────────────────────────
app = FastAPI(
    title="Wellfound AI - Excel Data Completion",
    description="AI-powered automated Excel data completion for Wellfound exports",
    version="1.1.0",
)

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

# Mount static files (with Docker-safe path resolution)
static_dir = Path(__file__).resolve().parent / "static"
static_dir.mkdir(parents=True, exist_ok=True)
app.mount("/static", StaticFiles(directory=str(static_dir)), name="static")

# Data directories (Docker-safe paths)
DATA_DIR = Path(__file__).resolve().parent / "data"
UPLOAD_DIR = DATA_DIR / "uploads"
RESULT_DIR = DATA_DIR / "results"
CHECKPOINT_DIR = DATA_DIR / "checkpoints"
for d in [UPLOAD_DIR, RESULT_DIR, CHECKPOINT_DIR]:
    d.mkdir(parents=True, exist_ok=True)

# Global processing state
processing_tasks: dict = {}
excel_handler = ExcelHandler(str(CHECKPOINT_DIR))
address_processor = AddressProcessor()
contact_finder = ContactFinder()


# ─── Routes ──────────────────────────────────────────────

@app.get("/", response_class=HTMLResponse)
async def index():
    """Serve the main application page."""
    template_path = Path(__file__).parent / "templates" / "index.html"
    if template_path.exists():
        return template_path.read_text(encoding="utf-8")
    return HTMLResponse("<h1>Wellfound AI</h1><p>Frontend not found.</p>")


@app.get("/api/health")
async def health():
    """Health check endpoint."""
    return {
        "status": "ok",
        "timestamp": datetime.now().isoformat(),
        "version": "1.1.0",
    }


@app.post("/api/upload")
async def upload_file(file: UploadFile = File(...)):
    """Upload Excel file and return metadata."""
    if not file.filename.endswith((".xlsx", ".xls")):
        raise HTTPException(400, "Only Excel files (.xlsx, .xls) are supported")

    file_id = uuid.uuid4().hex[:12]
    file_path = UPLOAD_DIR / f"{file_id}_{file.filename}"
    content = await file.read()
    file_path.write_bytes(content)

    try:
        df = excel_handler.load(str(file_path))
    except Exception as e:
        file_path.unlink(missing_ok=True)
        raise HTTPException(400, f"Failed to read Excel file: {str(e)}")

    return {
        "file_id": file_id,
        "filename": file.filename,
        "rows": len(df),
        "columns": list(df.columns),
        "missing_data": {
            col: int(df[col].isnull().sum())
            for col in df.columns
            if df[col].isnull().sum() > 0
        },
    }


@app.post("/api/process")
async def process_file(request: Request):
    """Start processing an uploaded Excel file."""
    data = await request.json()
    file_id = data.get("file_id")
    api_key = data.get("api_key", "")
    provider = data.get("provider", "openai")
    model = data.get("model", "auto")
    start_row = data.get("start_row", 0)
    max_rows = data.get("max_rows", 0)  # 0 = all
    use_ai = data.get("use_ai", True)
    use_web_scraping = data.get("use_web_scraping", True)
    concurrency = data.get("concurrency", 2)
    company_timeout = data.get("company_timeout", COMPANY_TIMEOUT_SECONDS)

    if not file_id:
        raise HTTPException(400, "file_id is required")

    # Find uploaded file
    files = list(UPLOAD_DIR.glob(f"{file_id}_*"))
    if not files:
        raise HTTPException(404, "Uploaded file not found")
    file_path = str(files[0])

    # Validate API key if using AI
    if use_ai and not api_key:
        raise HTTPException(400, "API key is required when AI extraction is enabled")

    task_id = uuid.uuid4().hex[:8]
    processing_tasks[task_id] = {
        "status": "starting",
        "progress": 0,
        "total": 0,
        "current": "",
        "errors": [],
        "skipped": [],          # NEW: track skipped companies due to timeout
        "file_id": file_id,
        "result_path": None,
    }

    # Launch async processing
    asyncio.create_task(
        _process_file(
            task_id=task_id,
            file_path=file_path,
            api_key=api_key,
            provider=provider,
            model=model,
            start_row=start_row,
            max_rows=max_rows,
            use_ai=use_ai,
            use_web_scraping=use_web_scraping,
            concurrency=concurrency,
            company_timeout=company_timeout,
        )
    )

    return {"task_id": task_id, "status": "started"}


@app.get("/api/status/{task_id}")
async def task_status(task_id: str):
    """Get processing task status."""
    if task_id not in processing_tasks:
        raise HTTPException(404, "Task not found")

    task = processing_tasks[task_id]
    return {
        "task_id": task_id,
        "status": task["status"],
        "progress": task["progress"],
        "total": task["total"],
        "current": task["current"],
        "errors": task["errors"][-10:],  # Last 10 errors
        "error_count": len(task["errors"]),
        "skipped": task.get("skipped", [])[-10:],  # Last 10 skipped entries
        "skipped_count": len(task.get("skipped", [])),
        "has_result": task["result_path"] is not None,
    }


@app.get("/api/download/{task_id}")
async def download_result(task_id: str):
    """Download the processed Excel file."""
    if task_id not in processing_tasks:
        raise HTTPException(404, "Task not found")

    task = processing_tasks[task_id]
    if not task["result_path"] or not os.path.exists(task["result_path"]):
        raise HTTPException(404, "Result file not ready yet")

    return FileResponse(
        task["result_path"],
        filename=f"wellfound_processed_{datetime.now().strftime('%Y%m%d_%H%M%S')}.xlsx",
        media_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
    )


@app.get("/api/providers")
async def get_providers():
    """Get available AI providers and models."""
    return AIExtractor.PROVIDERS


# ─── Processing Logic ────────────────────────────────────

async def _process_file(
    task_id: str,
    file_path: str,
    api_key: str,
    provider: str,
    model: str,
    start_row: int,
    max_rows: int,
    use_ai: bool,
    use_web_scraping: bool,
    concurrency: int,
    company_timeout: int = COMPANY_TIMEOUT_SECONDS,
):
    """Main processing pipeline with smart timeout detection and auto-skip."""
    task = processing_tasks[task_id]
    scraper = None
    ai_extractor = None

    try:
        # Load data
        task["status"] = "loading"
        df = excel_handler.load(file_path)

        # Create working copy
        result_path = str(RESULT_DIR / f"processed_{task_id}.xlsx")
        df.to_excel(result_path, index=False, engine="openpyxl")
        task["result_path"] = result_path

        total_rows = len(df)
        if max_rows > 0:
            total_rows = min(max_rows, total_rows - start_row)
        else:
            total_rows = total_rows - start_row

        task["total"] = total_rows
        task["status"] = "processing"

        # Initialize services
        if use_web_scraping:
            scraper = WebScraper(headless=True)
            await scraper.start()

        if use_ai and api_key:
            ai_extractor = AIExtractor(provider=provider, api_key=api_key, model=model)

        # Process rows in batches with concurrency
        semaphore = asyncio.Semaphore(concurrency)
        processed = 0

        async def process_row(idx: int):
            nonlocal processed
            async with semaphore:
                row = df.iloc[idx].to_dict()
                company = str(row.get("Company Name", ""))
                internal_link = str(row.get("Internal Link", ""))
                external_link = str(row.get("External Link", ""))
                location = str(row.get("Location", ""))

                task["current"] = f"Row {idx + 1}/{start_row + total_rows}: {company}"

                # ── Smart Timeout Wrapper ──────────────────────
                # Each company gets a bounded time window. If it exceeds the
                # threshold the row is skipped and the failure is logged so
                # the batch pipeline never stalls on a single company.
                try:
                    updates = await asyncio.wait_for(
                        _process_single_company(
                            row=row,
                            idx=idx,
                            company=company,
                            internal_link=internal_link,
                            external_link=external_link,
                            location=location,
                            use_web_scraping=use_web_scraping,
                            scraper=scraper,
                            ai_extractor=ai_extractor,
                            result_path=result_path,
                        ),
                        timeout=company_timeout,
                    )
                except asyncio.TimeoutError:
                    # ── Timeout detected → auto-skip ─────────
                    skip_record = {
                        "company": company,
                        "row": idx + 1,
                        "reason": f"Processing exceeded {company_timeout}s timeout",
                        "timestamp": datetime.now().isoformat(),
                        "links": {
                            "internal": internal_link if internal_link != "nan" else None,
                            "external": external_link if external_link != "nan" else None,
                        },
                    }
                    task["skipped"].append(skip_record)
                    task["errors"].append(
                        f"Row {idx} ({company}): SKIPPED — timeout after {company_timeout}s"
                    )
                    logger.warning(
                        "⏱ Timeout skip — company '%s' (row %d) exceeded %ds",
                        company, idx + 1, company_timeout,
                    )
                    # Still count as processed so progress bar advances
                    processed += 1
                    task["progress"] = processed
                    return

                # Save updates
                if updates:
                    excel_handler.save_row(result_path, idx, updates)

                processed += 1
                task["progress"] = processed

        # Process rows
        end_row = start_row + total_rows
        tasks_list = [
            process_row(i) for i in range(start_row, min(end_row, len(df)))
        ]
        await asyncio.gather(*tasks_list, return_exceptions=True)

        task["status"] = "completed"
        task["progress"] = total_rows

        # Log summary
        skipped_count = len(task.get("skipped", []))
        if skipped_count > 0:
            logger.info(
                "Batch complete: %d/%d processed, %d skipped due to timeout",
                total_rows - skipped_count, total_rows, skipped_count,
            )
        else:
            logger.info("Batch complete: %d/%d processed, no timeouts", total_rows, total_rows)

    except Exception as e:
        task["status"] = "error"
        task["errors"].append(f"Fatal: {str(e)}")
        logger.error("Fatal processing error: %s", str(e))

    finally:
        if scraper:
            await scraper.stop()


async def _process_single_company(
    row: dict,
    idx: int,
    company: str,
    internal_link: str,
    external_link: str,
    location: str,
    use_web_scraping: bool,
    scraper: Optional[WebScraper],
    ai_extractor: Optional[AIExtractor],
    result_path: str,
) -> dict:
    """Process a single company row — extract funding, contacts, and address.

    Returns a dict of column updates.  Callers should wrap this in
    ``asyncio.wait_for()`` to enforce per-company timeouts.
    """
    updates = {}

    # Step 1: Process wellfound page for funding data
    if use_web_scraping and internal_link and internal_link != "nan":
        funding_data = await _extract_funding(
            scraper, ai_extractor, company, internal_link
        )
        if funding_data:
            for col_key, excel_col in [
                ("valuation", "Valuation"),
                ("rounds", "Rounds"),
                ("series", "Series"),
                ("total_raised", "Total Raised"),
            ]:
                if funding_data.get(col_key) and pd.isna(row.get(excel_col)):
                    updates[excel_col] = funding_data[col_key]

    # Step 2: Process company website for contacts and address
    if use_web_scraping and external_link and external_link != "nan":
        contact_data = await _extract_contacts(
            scraper, ai_extractor, company, external_link
        )
        if contact_data:
            # Contact info
            if pd.isna(row.get("contact")):
                contact_email = contact_data.get("contact_email")
                contact_form = contact_data.get("contact_form_url")
                contact_str = contact_finder.format_contact_output(
                    contact_email, contact_form
                )
                if contact_str:
                    updates["contact"] = contact_str

            # Location info
            if pd.isna(row.get("location.apply")) or pd.isna(row.get("state.apply")):
                loc_apply, state_apply = address_processor.determine_location_apply(
                    wellfound_location=location,
                    ai_contact_data=contact_data,
                    scraped_addresses=contact_data.get("scraped_addresses", []),
                )
                if loc_apply and pd.isna(row.get("location.apply")):
                    updates["location.apply"] = loc_apply
                if state_apply and pd.isna(row.get("state.apply")):
                    updates["state.apply"] = state_apply

    return updates


async def _extract_funding(
    scraper: Optional[WebScraper],
    ai_extractor: Optional[AIExtractor],
    company: str,
    url: str,
) -> dict:
    """Extract funding information from wellfound page."""
    result = {}

    if not scraper:
        return result

    try:
        page_data = await scraper.fetch_wellfound_page(url)

        if page_data.get("error"):
            return result

        # Get regex-based extraction
        funding_data = page_data.get("funding_data", {})

        # AI-enhanced extraction
        if ai_extractor and page_data.get("text"):
            try:
                ai_data = await ai_extractor.analyze_funding_page(
                    company,
                    page_data["text"],
                    page_data.get("meta"),
                )

                # Merge AI results with regex results, preferring high-confidence AI
                for key in ["valuation", "total_raised", "rounds", "series"]:
                    ai_val = ai_data.get(key)
                    ai_conf = ai_data.get("confidence", {}).get(key, "low")
                    regex_val = funding_data.get(key)

                    if ai_val and ai_conf == "high":
                        result[key] = ai_val
                    elif regex_val:
                        result[key] = regex_val
                    elif ai_val:
                        result[key] = ai_val
            except Exception:
                # Fall back to regex-only
                for key in ["valuation", "total_raised", "rounds", "series"]:
                    if funding_data.get(key):
                        result[key] = funding_data[key]
        else:
            # No AI, use regex results
            for key in ["valuation", "total_raised", "rounds", "series"]:
                if funding_data.get(key):
                    result[key] = funding_data[key]

    except Exception as e:
        pass  # Silently fail on individual row errors

    return result


async def _extract_contacts(
    scraper: Optional[WebScraper],
    ai_extractor: Optional[AIExtractor],
    company: str,
    url: str,
) -> dict:
    """Extract contact information from company website."""
    result = {
        "contact_email": None,
        "contact_form_url": None,
        "careers_page_url": None,
        "headquarters_city": None,
        "headquarters_state": None,
        "scraped_addresses": [],
        "us_office_city": None,
        "us_office_state": None,
        "hiring_focus_location": None,
    }

    if not scraper:
        return result

    try:
        page_data = await scraper.fetch_company_website(url)

        if page_data.get("error"):
            return result

        emails = page_data.get("emails", [])
        links = page_data.get("links", [])
        phones = page_data.get("phones", [])
        addresses = page_data.get("addresses", [])
        text = page_data.get("text", "")

        result["scraped_addresses"] = addresses

        # Contact email
        best_email = contact_finder.find_best_contact_email(emails, text)
        result["contact_email"] = best_email

        # Contact form
        result["contact_form_url"] = contact_finder.find_contact_form_url(
            url, links, page_data.get("html", "")
        )

        # Careers page
        result["careers_page_url"] = contact_finder.find_careers_page_url(
            url, links, page_data.get("html", "")
        )

        # AI-enhanced extraction for location and additional contacts
        if ai_extractor and text:
            try:
                ai_data = await ai_extractor.analyze_company_website(
                    company, text, emails, links, phones, addresses
                )

                # Use AI data if available
                if not result["contact_email"] and ai_data.get("contact_email"):
                    result["contact_email"] = ai_data["contact_email"]

                result["headquarters_city"] = ai_data.get("headquarters_city")
                result["headquarters_state"] = ai_data.get("headquarters_state")
                result["us_office_city"] = ai_data.get("us_office_city")
                result["us_office_state"] = ai_data.get("us_office_state")
                result["hiring_focus_location"] = ai_data.get("hiring_focus_location")

                if not result["contact_form_url"]:
                    result["contact_form_url"] = ai_data.get("contact_form_url")
                if not result["careers_page_url"]:
                    result["careers_page_url"] = ai_data.get("careers_page_url")

            except Exception:
                pass

    except Exception:
        pass

    return result


# ─── Main ─────────────────────────────────────────────────

if __name__ == "__main__":
    import uvicorn

    # HF Spaces sets PORT=7860; respect it, otherwise auto-detect
    hf_port = os.environ.get("PORT") or os.environ.get("HF_SPACE_PORT")
    if hf_port:
        port = int(hf_port)
        print(f"\n  HF Spaces detected: using PORT={port}")
    else:
        port = 7860
        # Try to find an available port locally
        import socket
        try:
            s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            s.bind(("0.0.0.0", port))
            s.close()
        except OSError:
            for alt in range(7861, 7870):
                try:
                    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                    s.bind(("0.0.0.0", alt))
                    s.close()
                    port = alt
                    break
                except OSError:
                    continue

    print(f"\n{'='*50}")
    print(f"  Wellfound AI - Server Starting")
    print(f"  URL: http://0.0.0.0:{port}")
    print(f"  Data Dir: {DATA_DIR}")
    print(f"  Company Timeout: {COMPANY_TIMEOUT_SECONDS}s")
    print(f"{'='*50}\n")

    uvicorn.run(
        app,
        host="0.0.0.0",
        port=port,
        log_level="info",
        # Graceful shutdown for Docker
        timeout_keep_alive=30,
    )