app/api/v1/companies.py

from fastapi import APIRouter, Query
from typing import List, Optional
from app.services.data_loader import get_companies_df
from app.models.schemas import Company, CompanyWithLinks
from app.services.osint_links import get_company_links
from app.core.utils import clean_nans

router = APIRouter()

@router.get("/", response_model=List[Company])
async def list_companies(
    wilaya: Optional[str] = None,
    group: Optional[str] = None,
    type: Optional[str] = None,
    search: Optional[str] = None,
    limit: int = 50
):
    df = await get_companies_df()
    if df.empty:
        return []

    if wilaya:
        df = df[df['wilaya'] == wilaya]
    if group:
        df = df[df['activity_group'] == group]
    if type:
        df = df[df['type'] == type]
    if search:
        mask = df['name'].str.contains(search, na=False, case=False) | \
               df['activity_normalized'].str.contains(search, na=False, case=False)
        df = df[mask]
    
    data = df.head(limit).to_dict(orient='records')
    return clean_nans(data)

@router.get("/{company_id}", response_model=CompanyWithLinks)
async def read_company(company_id: str):
    df = await get_companies_df()
    company = df[df['id'] == company_id]
    if company.empty:
        from fastapi import HTTPException
        raise HTTPException(status_code=404, detail="Company not found")
        
    data = company.iloc[0].to_dict()
    data['osint_links'] = await get_company_links(company_id)
    return clean_nans(data)

@router.get("/{company_id}/osint_links")
async def read_company_links(company_id: str):
    return await get_company_links(company_id)