app.py

# app.py
import os
import io
import json
import time
import traceback
from typing import List
from io import BytesIO

from fastapi import FastAPI, UploadFile, File, Form, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from dotenv import load_dotenv

import PyPDF2
import openai
import requests
from pdf2image import convert_from_bytes
from google.cloud import vision_v1 as vision

# ───── Load .env (locally) ────────────────────────────────────
# In HF Spaces you'll set these as Secrets under Settings → Variables & secrets
load_dotenv()
openai.api_key = os.getenv("OPENAI_API_KEY")
gcv_api_key = os.getenv("GCV_API_KEY")
brightdata_api_key = os.getenv("BRIGHTDATA_API_KEY")
brightdata_dataset_id = "gd_l1viktl72bvl7bjuj0"

if not openai.api_key:
    raise RuntimeError("Missing OPENAI_API_KEY")
if not gcv_api_key:
    raise RuntimeError("Missing GCV_API_KEY")
if not brightdata_api_key:
    raise RuntimeError("Missing BRIGHTDATA_API_KEY")

# ───── FastAPI setup ──────────────────────────────────────────
app = FastAPI(title="Aliro Data Extraction API")
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

# ───── A simple root so GET / won't 404 ───────────────────────
@app.get("/")
def read_root():
    return {"message": "Aliro Data Extraction API – POST your files to /extract"}

# ───── PDF/OCR processing ─────────────────────────────────────
def process_buffers(buffers: List[BytesIO]) -> List[str]:
    client = vision.ImageAnnotatorClient(client_options={"api_key": gcv_api_key})
    results: List[str] = []

    for buf in buffers:
        name = getattr(buf, "name", "").lower()
        data = buf.getvalue()

        # PDF: try text → fallback to page-by-page OCR
        if name.endswith(".pdf"):
            # 1) PyPDF2 text
            try:
                reader = PyPDF2.PdfReader(io.BytesIO(data))
                text = "".join(page.extract_text() or "" for page in reader.pages)
                if text.strip():
                    results.append(text)
                    continue
            except:
                pass

            # 2) Fallback: render pages → OCR
            try:
                for img in convert_from_bytes(data):
                    img_buf = io.BytesIO()
                    img.save(img_buf, format="JPEG")
                    resp = client.text_detection(image=vision.Image(content=img_buf.getvalue()))
                    desc = resp.text_annotations[0].description if resp.text_annotations else ""
                    if desc:
                        results.append(desc)
                continue
            except:
                pass

        # Image → always OCR
        if name.endswith((".png", ".jpg", ".jpeg")):
            try:
                resp = client.text_detection(image=vision.Image(content=data))
                desc = resp.text_annotations[0].description if resp.text_annotations else ""
                if desc:
                    results.append(desc)
                    continue
            except:
                pass

        # Otherwise unsupported
        results.append(f"[Unsupported file type: {name}]")

    return results

# ───── LinkedIn scrape via BrightData ─────────────────────────
def scrape_linkedin(url: str) -> str:
    if not url:
        return ""
        
    try:
        # Step 1: Trigger data collection
        trigger_url = "https://api.brightdata.com/datasets/v3/trigger"
        headers = {
            "Authorization": f"Bearer {brightdata_api_key}",
            "Content-Type": "application/json",
        }
        params = {
            "dataset_id": brightdata_dataset_id,
            "include_errors": "true",
        }
        
        trigger_response = requests.post(
            trigger_url,
            headers=headers,
            params=params,
            json=[{"url": url}],
            timeout=30
        )
        
        if not trigger_response.ok:
            raise Exception(f"BrightData trigger failed: {trigger_response.status_code} - {trigger_response.text}")
            
        trigger_data = trigger_response.json()
        snapshot_id = trigger_data.get('snapshot_id')
        
        if not snapshot_id:
            raise Exception("No snapshot_id received from BrightData")
            
        print(f"BrightData collection triggered, snapshot_id: {snapshot_id}")
        
        # Step 2: Poll for completion (max 5 minutes)
        max_attempts = 30
        for attempt in range(max_attempts):
            time.sleep(10)  # Wait 10 seconds between checks
            
            progress_url = f"https://api.brightdata.com/datasets/v3/progress/{snapshot_id}"
            progress_response = requests.get(progress_url, headers=headers, timeout=30)
            
            if progress_response.ok:
                progress_data = progress_response.json()
                status = progress_data.get('status')
                
                print(f"Progress check {attempt + 1}: status = {status}")
                
                if status == 'ready':
                    # Step 3: Fetch results
                    result_url = f"https://api.brightdata.com/datasets/v3/snapshot/{snapshot_id}"
                    result_params = {"format": "json"}
                    result_response = requests.get(result_url, headers=headers, params=result_params, timeout=30)
                    
                    if result_response.ok:
                        linkedin_data = result_response.json()
                        return json.dumps(linkedin_data, indent=2)
                    else:
                        raise Exception(f"Failed to fetch results: {result_response.status_code} - {result_response.text}")
                        
                elif status == 'failed':
                    raise Exception("BrightData collection failed")
                    
                # Continue polling if status is 'running'
            else:
                print(f"Progress check failed: {progress_response.status_code}")
                
        raise Exception("Timeout waiting for BrightData collection to complete")
        
    except Exception as e:
        raise Exception(f"LinkedIn scraping error: {str(e)}")

# ───── Summarize via OpenAI ─────────────────────────────────
def make_summary(chunks: List[str]) -> str:
    if not chunks:
        return "No data extracted."
    prompt = "\n\n---\n\n".join(chunks)
    resp = openai.ChatCompletion.create(
        model="gpt-4o",
        messages=[
            {
                "role": "system",
                "content": (
                    "Extract and structure personal, educational, and professional details "
                    "into a clear, logical and relational hierarchy. Translate non-English "
                    "content, remove duplicates, and output in English."
                )
            },
            {"role": "user", "content": f"Summarize the following data:\n\n{prompt}"},
        ],
        max_tokens=1500,
    )
    return resp.choices[0].message.content.strip()

# ───── The POST /extract endpoint ─────────────────────────────
@app.post("/extract")
async def extract_endpoint(
    files: List[UploadFile] = File(default=[]),
    linkedin_url: str       = Form(default="")
):
    try:
        # Wrap each UploadFile in a BytesIO (so we can peek at .name & .getvalue())
        buffers: List[BytesIO] = []
        for f in files:
            data = await f.read()
            bio = BytesIO(data)
            bio.name = f.filename
            buffers.append(bio)

        # 1) OCR/PDF text
        texts = process_buffers(buffers)
        # 2) LinkedIn JSON
        if linkedin_url:
            texts.append(scrape_linkedin(linkedin_url))
        # 3) Summarize
        summary = make_summary(texts)
        return {"summary": summary}

    except requests.HTTPError as e:
        raise HTTPException(status_code=502, detail=f"LinkedIn scrape failed: {e}")
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Processing error: {str(e)}")