import PyPDF2
from docx import Document
import io
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.text_splitter import CharacterTextSplitter
from typing_extensions import Concatenate
from typing import List
# from langchain_community.llms import OpenAI
from langchain_community.callbacks import get_openai_callback
from langchain.output_parsers import PydanticOutputParser
from langchain.prompts import PromptTemplate
from langchain_core.pydantic_v1 import BaseModel, Field, validator
import os
import logging
import base64
from langchain_openai import OpenAI

#Setting the openai api key
api_key=os.getenv('OPENAI_API_KEY')

class Candidate(BaseModel):
    brand: str = Field(description="Please provide the name of the brand from which the receipt is. If there are multiple brand names present, please specify the most prominent or relevant brand associated with the receipt. If the brand name is not explicitly stated, provide any context or details that might help identify it accurately. If the information is not present or unclear, mark the value as Unknown Brand.")
    total_cost: str = Field(description="For example, you can specify the total order value section on the receipt, such as ‘Total’ or ‘Amount Due’, and describe where it's typically located (e.g., bottom right corner). Additionally, mention any common variations in labeling that might indicate the total order value, such as ‘Grand Total’,’Total Amount’, or ‘Balance Due’.")
    location: str = Field(description="Please provide the city and state where the purchase was made. If the receipt is related to travel, extract the location from where the booking was made, not the destination. Look for details such as the booking origin or departure city/state. If the purchase location is not explicitly stated, provide any relevant clues or context from the receipt to help identify it accurately. If the information is not present or unclear, mark the value as Unknown Location.For travel receipts, consider extracting details such as the departure airport code, departure city, or booking location mentioned in the itinerary or booking confirmation section. This location typically indicates where the purchase was initiated.")
    no_of_items: str = Field(description="Number of items in the order.If the information is not present mark the value as 0.")
    purchase_category: str = Field(description="The purchase category among fashion, home, travel, food, groceries, hotels, spa, insurance and others.If the information is not present mark the value as Unknown Category.")
    Date: str = Field(description="Date of purchase.If the information is not present mark the value as Unknown Date. Make sure the date format is in dd-MM-yyyy.")
    

# async def initialize_openai():
#     model_name = "gpt-3.5-turbo-instruct"
#     # model_name = "text-davinci-003"
#     temperature = 0.0
#     model = OpenAI(model_name=model_name, temperature=temperature, max_tokens=800)
    
def strcuture_document_data(raw_text:str)->dict:
    try:
        print("printing raw text")
        print(raw_text)
        model_name = "gpt-3.5-turbo-instruct"
        # model_name = "text-davinci-003"
        temperature = 0.0
        model = OpenAI(model_name=model_name, temperature=temperature, max_tokens=800)
        doc_query = "Return only a json which contains important entities from the given receipt.Also json should contain only these keys - brand,total_cost,location,no_of_items,purchase_category,Date. Make sure the date format is dd-MM-yyyy in the json. Response: " + raw_text
        parser = PydanticOutputParser(pydantic_object=Candidate)
        
        prompt = PromptTemplate(
            template="Answer the user query.\n{format_instructions}\n{query}\n",
            input_variables=["query"],
            partial_variables={"format_instructions": parser.get_format_instructions()},
        )
        input = prompt.format_prompt(query=doc_query)
        with get_openai_callback() as cb:
            result = model(input.to_string())
            # print(parser.parse(result))
            print("printing result")
            print(result)
            # print(type(result))
            logging.info(result)
        class_object= parser.parse(result)
        dict_object=class_object.__dict__
        # print("printing structured json")
        # logging.info(dict_object)
        # print(dict_object)
        return dict_object
    except Exception as e:
        print(f"Error occurred: {e}")
        return {}


def extract_text_from_pdf(pdf_data):
    with io.BytesIO(pdf_data) as pdf_file:
        pdf_reader = PyPDF2.PdfReader(pdf_file)
        text = ""
        for page_num in range(len(pdf_reader.pages)):
            page = pdf_reader.pages[page_num]
            text += page.extract_text()
        return text

def extract_text_from_docx(docx_data):
    doc = Document(io.BytesIO(docx_data))
    text = ""
    for para in doc.paragraphs:
        text += para.text + "\n"
    return text

def extract_text_from_attachment(filename, data):
    if filename.endswith('.pdf'):
        return extract_text_from_pdf(base64.urlsafe_b64decode(data))
    elif filename.endswith('.docx'):
        return extract_text_from_docx(base64.urlsafe_b64decode(data))
    else:
        # Add handling for other document types if needed
        return "Unsupported document type"