from ombra import step, File, call_llm, Model, Message, Role
from example.models import DocumentAnalysis, EmailDraft

@step()
async def analyze_document(file: File) -> DocumentAnalysis:
    """
    Sends the PDF to Gemini to extract summary, entities, and sentiment.
    """
    prompt = """
    You are an expert document analyst. Please analyze the attached PDF document.
    Provide a concise summary, list key entities (people, companies, dates), 
    determine the overall sentiment (positive, neutral, negative), 
    and indicate if any immediate action is required.
    """

    response = await call_llm(
        model=Model.GEMINI_2_5_FLASH_LITE,
        messages=[
            Message(
                role=Role.USER,
                content=[prompt, file] 
            )
        ],
        response_format=DocumentAnalysis, # Structured Output!
        timeout=120,
        temperature=0.7,
        retries=3,
    )
    
    # LiteLLM returns valid JSON string matching the schema in content
    return DocumentAnalysis.model_validate_json(response.content)

@step()
async def draft_response_email(analysis: DocumentAnalysis) -> EmailDraft:
    """
    Drafts a response email based on the document analysis.
    """
    prompt = f"""
    Based on the following document analysis, draft a professional response email.
    
    Summary: {analysis.summary}
    Sentiment: {analysis.sentiment}
    Requires Action: {analysis.requires_action}
    
    If action is required, the tone should be urgent but polite. 
    If no action is required, just acknowledge receipt.
    """
    
    response = await call_llm(
        model=Model.GEMINI_2_5_FLASH_LITE,
        messages=[Message(role=Role.USER, content=prompt)],
        response_format=EmailDraft, # Structured Output!
        timeout=120,
        temperature=0.7,
        retries=3,
    )
    
    return EmailDraft.model_validate_json(response.content)
