backend/routers/export.py

import io
import pandas as pd
from fastapi import APIRouter, Depends, Query, HTTPException
from fastapi.responses import StreamingResponse
from sqlalchemy.orm import Session
from typing import Optional
from database import get_db
from models import Lead
from auth import decode_token

router = APIRouter(prefix="/export", tags=["export"])


def _verify_token_param(token: Optional[str], db: Session) -> None:
    """Allow auth via ?token= query param for direct download links."""
    if not token:
        raise HTTPException(status_code=401, detail="Not authenticated")
    payload = decode_token(token)
    if not payload:
        raise HTTPException(status_code=401, detail="Invalid or expired token")

EXPORT_FIELDS = [
    "id", "company_name", "contact_name", "email", "phone", "website",
    "industry", "company_size", "location", "linkedin_url",
    "source", "status", "notes", "created_at",
]


def leads_to_df(leads: list[Lead]) -> pd.DataFrame:
    rows = []
    for lead in leads:
        row = {f: getattr(lead, f, None) for f in EXPORT_FIELDS}
        if lead.custom_fields:
            for k, v in lead.custom_fields.items():
                row[f"custom_{k}"] = v
        rows.append(row)
    return pd.DataFrame(rows)


def get_filtered_leads(
    db: Session,
    status: Optional[str],
    source: Optional[str],
    industry: Optional[str],
    ids: Optional[str],
) -> list[Lead]:
    query = db.query(Lead)
    if ids:
        id_list = [int(i) for i in ids.split(",") if i.strip().isdigit()]
        query = query.filter(Lead.id.in_(id_list))
    if status:
        query = query.filter(Lead.status == status)
    if source:
        query = query.filter(Lead.source == source)
    if industry:
        query = query.filter(Lead.industry.ilike(f"%{industry}%"))
    return query.order_by(Lead.created_at.desc()).all()


@router.get("/csv")
def export_csv(
    token: Optional[str] = Query(None),
    status: Optional[str] = None,
    source: Optional[str] = None,
    industry: Optional[str] = None,
    ids: Optional[str] = Query(None, description="Comma-separated lead IDs"),
    db: Session = Depends(get_db),
):
    _verify_token_param(token, db)
    leads = get_filtered_leads(db, status, source, industry, ids)
    df = leads_to_df(leads)
    buf = io.StringIO()
    df.to_csv(buf, index=False)
    buf.seek(0)
    return StreamingResponse(
        iter([buf.getvalue()]),
        media_type="text/csv",
        headers={"Content-Disposition": "attachment; filename=leads_export.csv"},
    )


@router.get("/excel")
def export_excel(
    token: Optional[str] = Query(None),
    status: Optional[str] = None,
    source: Optional[str] = None,
    industry: Optional[str] = None,
    ids: Optional[str] = Query(None, description="Comma-separated lead IDs"),
    db: Session = Depends(get_db),
):
    _verify_token_param(token, db)
    leads = get_filtered_leads(db, status, source, industry, ids)
    df = leads_to_df(leads)
    buf = io.BytesIO()
    with pd.ExcelWriter(buf, engine="openpyxl") as writer:
        df.to_excel(writer, index=False, sheet_name="Leads")
    buf.seek(0)
    return StreamingResponse(
        iter([buf.read()]),
        media_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
        headers={"Content-Disposition": "attachment; filename=leads_export.xlsx"},
    )