tasks.py

import os
import uuid
import shutil
import subprocess
import os
import uuid
import shutil
import subprocess
from pathlib import Path
from datetime import datetime
from celery_app import celery_app

@celery_app.task
def process_document_task(collection_id: str, doc_id: str, filename: str, file_type: str, upload_path_str: str):
    from services.rag_service import (
        load_metadata,
        save_metadata,
        CONVERTED_DIR,
        convert_to_pdf,
        get_embedding_model,
        is_supabase_active,
    )
    
    upload_path = Path(upload_path_str)
    
    if not upload_path.exists() and is_supabase_active():
        print(f"Local file {upload_path} not found. Attempting to download from Supabase Storage...")
        from services.rag_service import get_supabase_client
        supabase = get_supabase_client()
        if supabase:
            from db import get_db_connection
            with get_db_connection() as conn:
                with conn.cursor() as cursor:
                    cursor.execute("SELECT upload_path FROM rag_documents WHERE id = ?", (doc_id,))
                    row = cursor.fetchone()
            if row and row["upload_path"]:
                storage_path = row["upload_path"]
                upload_path.parent.mkdir(parents=True, exist_ok=True)
                try:
                    res = supabase.storage.from_("learning-playbook").download(storage_path)
                    with open(upload_path, "wb") as f:
                        f.write(res)
                    print(f"Downloaded {storage_path} to {upload_path} successfully.")
                except Exception as e:
                    print(f"Error downloading from Supabase Storage: {e}")

    
    # ---- LangChain pipeline ----
    from langchain_text_splitters import RecursiveCharacterTextSplitter
    from langchain_community.vectorstores import Chroma

    if file_type == "xlsx":
        from langchain_community.document_loaders import UnstructuredExcelLoader
        loader = UnstructuredExcelLoader(str(upload_path), mode="elements")
        pages = loader.load()
    else:
        if file_type == "pdf":
            pdf_path = upload_path
        else:
            conv_dir = CONVERTED_DIR / collection_id
            conv_dir.mkdir(parents=True, exist_ok=True)
            pdf_path = convert_to_pdf(upload_path, conv_dir)

        from langchain_community.document_loaders import PyPDFLoader
        loader = PyPDFLoader(str(pdf_path))
        pages = loader.load()

    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=1000,
        chunk_overlap=300,
        length_function=len,
        add_start_index=True,
    )
    chunks = text_splitter.split_documents(pages)

    for chunk in chunks:
        chunk.metadata["doc_id"] = doc_id
        chunk.metadata["filename"] = filename

    embedding_model = get_embedding_model()
    
    if chunks:
        from langchain_pinecone import PineconeVectorStore
        PineconeVectorStore.from_documents(
            chunks,
            embedding_model,
            index_name="maple-prospect",
            namespace=f"col_{collection_id[:8]}"
        )

    from services.rag_service import is_supabase_active
    if is_supabase_active():
        from db import get_db_connection
        with get_db_connection() as conn:
            with conn.cursor() as cursor:
                cursor.execute(
                    "UPDATE rag_documents SET pages = ?, chunks = ? WHERE id = ?",
                    (len(pages), len(chunks), doc_id)
                )
    else:
        meta = load_metadata()
        if collection_id in meta["collections"]:
            meta["collections"][collection_id]["documents"][doc_id] = {
                "filename": filename,
                "original_type": file_type,
                "pages": len(pages),
                "chunks": len(chunks),
                "uploaded_at": datetime.now().isoformat(),
                "upload_path": str(upload_path),
            }
            save_metadata(meta)
    
    return {"status": "success", "doc_id": doc_id}


@celery_app.task
def process_youtube_task(collection_id: str, doc_id: str, url: str, video_id: str, source_name: str):
    from services.rag_service import (
        load_metadata,
        save_metadata,
        _index_text_as_document,
        get_all_gemini_api_keys
    )
    from google import genai
    import json as _json

    keys = get_all_gemini_api_keys()
    if not keys:
        raise ValueError("Limit reached try after 5 minutes please contact admin")

    transcript_text = None
    last_err = None

    for youtube_api_key in keys:
        try:
            client = genai.Client(api_key=youtube_api_key)
            prompt = "Process the audio file and generate a detailed transcription."
            response_schema = {
                "type": "object",
                "properties": {
                    "segments": {
                        "type": "array",
                        "items": {
                            "type": "object",
                            "properties": {
                                "speaker": {"type": "string"},
                                "timestamp": {"type": "string"},
                                "content": {"type": "string"},
                                "language": {"type": "string"},
                            },
                            "required": ["speaker", "timestamp", "content", "language"]
                        }
                    }
                },
                "required": ["segments"]
            }

            interaction = client.interactions.create(
                model="gemini-2.5-flash",
                input=[
                    {"type": "video", "uri": url, "mime_type": "video/mp4"},
                    {"type": "text", "text": prompt}
                ],
                response_format=response_schema,
            )

            result = _json.loads(interaction.output_text)
            transcript_text = " ".join(segment["content"] for segment in result["segments"]).strip()
            
            if not transcript_text:
                raise ValueError("Transcript is empty.")
            break
        except Exception as e:
            last_err = e
            transcript_text = None
            continue

    if not transcript_text:
        raise ValueError("Limit reached try after 5 minutes please contact admin")

    indexed = _index_text_as_document(
        collection_id, doc_id, transcript_text, source_name, "youtube"
    )

    from services.rag_service import is_supabase_active
    if is_supabase_active():
        from db import get_db_connection
        with get_db_connection() as conn:
            with conn.cursor() as cursor:
                cursor.execute(
                    "UPDATE rag_documents SET pages = ?, chunks = ? WHERE id = ?",
                    (1, indexed["chunks"], doc_id)
                )
    else:
        meta = load_metadata()
        if collection_id in meta["collections"]:
            meta["collections"][collection_id]["documents"][doc_id] = {
                "filename": source_name,
                "original_type": "youtube",
                "source_url": url,
                "video_id": video_id,
                "pages": 1,
                "chunks": indexed["chunks"],
                "uploaded_at": datetime.now().isoformat(),
                "upload_path": None,
            }
            save_metadata(meta)
        
    return {"status": "success", "doc_id": doc_id}


@celery_app.task
def process_video_task(collection_id: str, doc_id: str, filename: str, upload_path_str: str):
    from services.rag_service import (
        load_metadata,
        save_metadata,
        _index_text_as_document,
        get_all_gemini_api_keys
    )
    from google import genai
    import json as _json

    keys = get_all_gemini_api_keys()
    if not keys:
        raise ValueError("Limit reached try after 5 minutes please contact admin")

    transcript_text = None
    last_err = None

    for video_api_key in keys:
        try:
            client = genai.Client(api_key=video_api_key)
            file_manager = client.files

            upload_path = Path(upload_path_str)
            uploaded_file = file_manager.upload(file=str(upload_path), mime_type="video/mp4")

            import time
            while uploaded_file.state.name == "PROCESSING":
                time.sleep(5)
                uploaded_file = file_manager.get(name=uploaded_file.name)

            if uploaded_file.state.name == "FAILED":
                raise RuntimeError("Video processing on Gemini failed.")

            prompt = "Process the audio file and generate a detailed transcription."
            response_schema = {
                "type": "object",
                "properties": {
                    "segments": {
                        "type": "array",
                        "items": {
                            "type": "object",
                            "properties": {
                                "speaker": {"type": "string"},
                                "timestamp": {"type": "string"},
                                "content": {"type": "string"},
                                "language": {"type": "string"},
                            },
                            "required": ["speaker", "timestamp", "content", "language"]
                        }
                    }
                },
                "required": ["segments"]
            }

            interaction = client.interactions.create(
                model="gemini-2.5-flash",
                input=[uploaded_file, {"type": "text", "text": prompt}],
                response_format=response_schema,
            )

            result = _json.loads(interaction.output_text)
            transcript_text = " ".join(segment["content"] for segment in result["segments"]).strip()

            if not transcript_text:
                raise ValueError("Transcript is empty.")
                
            file_manager.delete(name=uploaded_file.name)
            break
        except Exception as e:
            last_err = e
            transcript_text = None
            continue

    if not transcript_text:
        raise ValueError("Limit reached try after 5 minutes please contact admin")

    indexed = _index_text_as_document(
        collection_id, doc_id, transcript_text, filename, "mp4"
    )

    from services.rag_service import is_supabase_active
    if is_supabase_active():
        from db import get_db_connection
        with get_db_connection() as conn:
            with conn.cursor() as cursor:
                cursor.execute(
                    "UPDATE rag_documents SET pages = ?, chunks = ? WHERE id = ?",
                    (1, indexed["chunks"], doc_id)
                )
    else:
        meta = load_metadata()
        if collection_id in meta["collections"]:
            meta["collections"][collection_id]["documents"][doc_id] = {
                "filename": filename,
                "original_type": "mp4",
                "pages": 1,
                "chunks": indexed["chunks"],
                "uploaded_at": datetime.now().isoformat(),
                "upload_path": str(upload_path),
            }
            save_metadata(meta)
        
    return {"status": "success", "doc_id": doc_id}

@celery_app.task
def geo_search_task(job_id: str, query: str):
    from services.ai_service import run_geo_job
    run_geo_job(job_id, query)
    return {'status': 'success', 'job_id': job_id}

@celery_app.task
def company_research_task(job_id: str, payload: dict):
    from services.ai_service import run_company_research_job
    run_company_research_job(job_id, payload)
    return {'status': 'success', 'job_id': job_id}

@celery_app.task
def ocr_task(job_id: str, filename: str, mime_type: str, file_content_b64: str):
    import base64
    from services.ocr import extract_business_card
    from services.ai_service import build_search_query, now, update_job

    update_job(job_id, status="running", started_at=now())
    try:
        file_bytes = base64.b64decode(file_content_b64)
        card_data = extract_business_card(file_bytes, mime_type)
        result = {'card': card_data, 'search_query': build_search_query(card_data)}
        update_job(job_id, status="completed", finished_at=now(), result=result)
        return {'status': 'success', 'job_id': job_id, **result}
    except Exception as exc:
        update_job(job_id, status="failed", finished_at=now(), error=str(exc))
        raise