import base64
from crewai import Agent, Task, Crew, Process
from crewai.tools import tool
from langchain_openai import ChatOpenAI
from openai import OpenAI

## client for vision
vision_client = OpenAI()

@tool("Invoice Image Reader")
def read_invoice_image(image_path:str)->str:
    """
    Read the image and extract the raw text from it
    """
    with open(image_path,'rb') as f:
        image_base64 = base64.b64encode(f.read()).decode('utf-8')

    response = vision_client.responses.create(
    model="gpt-4.1-mini", 
    input=[{
        "role": "user",
        "content": [
            {"type": "input_text", "text": ("Extract the vendor name, tax id , invoice number"
                                            "invoice date, items table (descriptiopn, quantity, net price ),"
                                            "and total gross from this invoice"
                                           ),
            
             
             },
            {
                "type": "input_image",
                "image_url": f"data:image/jpeg;base64,{image_base64}",
            },
        ],
    }],
)

    return response.output_text

def extract_invoice(image_path:str)->str:
    """
    Extract the information in JSON structure
    """

    llm = ChatOpenAI(model = "gpt-4.1-mini", temperature=0)

    ### Agent 1 OCR specialist

    visual_reader = Agent(
                        role="OCR Specialist",
                        goal= "Extract the invoice data fram images",
                        backstory= ("you cann't see the images directly."
                                    "you must always use the Invoice Image Reader tool"),
                        llm =  llm,
                        tools = [read_invoice_image],
                        verbose=True
    )

    #### Agent 2 JSON

    json_architect = Agent(
                        role="Data Engineer",
                        goal= "Convert extracted invoice text to structured JSON",
                        backstory= "You normalize numbers and dates and output strict to JSON",
                        llm =  llm,
                        verbose=True
    )


    ## Task 1 extraction task 
    
    extraction_task = Task (
                            description = (
                            f"Use the Invoice Reader tool to read the invoice image "
                            f"at the path '{image_path}'. extract the vendor name , tax id, invoice number, " 
                            f"date, item rows, total gross"          
                            ),
                            expected_output = "structured invoice text",
                            agent = visual_reader
                            )
    
    ### Task 2 JSON 
    
    formatting_task = Task (
                            description = (
                            "Convert the extracted invoice text into JSON. \n\n"
                                "{\n"
                                "'invoice_no' : str, \n"
                                "'date': 'YYYY-MM-DD',\n"
                                " 'vendor':{'name':str,'tax_id':str},\n "
                                " 'items' :[{'description':str,'quantity':float,'net_price':float}] ,\n"
                                "}\n\n"
                                "Rules:\n"
                                "- if missing value use null \n"
                                " - Output only JSON"
                            ),
                            expected_output = "Valid JSON Only",
                            agent = json_architect,
                            context = [extraction_task]
                            )
    crew = Crew(
            agents = [visual_reader,json_architect],
            tasks = [extraction_task,formatting_task],
            process = Process.sequential,
            verbose =True)
    
    return crew.kickoff()