import os
import sqlite3
import fitz
import google.generativeai as genai
from dotenv import load_dotenv
import json

load_dotenv()
GEMINI_API_KEY = os.environ.get("API_KEY")


def configure_api():
    """Configure the Gemini API with the provided API key."""
    if not GEMINI_API_KEY:
        raise ValueError("API key not found. Please set the GEMINI_API_KEY environment variable.")
    genai.configure(api_key=GEMINI_API_KEY)


def create_model():
    """Create and return a generative AI model with specified configuration."""
    generation_config = {
        "temperature": 0.1,
        "top_p": 0.1,
        "top_k": 10,
        "max_output_tokens": 8192,
        "response_mime_type": "text/plain",
    }
    safety_settings = [
        {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_NONE"},
        {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"},
        {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_NONE"},
        {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_NONE"},
    ]
    model = genai.GenerativeModel(
        model_name="gemini-1.5-flash-002",
        safety_settings=safety_settings,
        generation_config=generation_config,
        system_instruction="You are a helpful AI.",
    )
    return model


def initialize_db(db_file="extracted_data.db"):
    """Creates the SQLite database and table if it doesn't exist."""
    conn = sqlite3.connect(db_file)
    cursor = conn.cursor()
    cursor.execute("""
        CREATE TABLE IF NOT EXISTS records (
            file_name TEXT PRIMARY KEY,
            product_name TEXT,
            part_number TEXT,
            voucher_number TEXT
        )
    """)
    conn.commit()
    conn.close()


def process_pdf_to_db(pdf_folder, prompt, db_connection):
    """
    Processes PDFs in a folder, extracts data using Gemini, and saves to SQLite database.

    Args:
        pdf_folder (str): Path to the folder containing PDF files.
        prompt (str): The prompt to send to the AI model.
        db_connection: SQLite database connection object.
    """
    cursor = db_connection.cursor()

    # Iterate through PDF files in the folder
    for root, dirs, files in os.walk(pdf_folder):
        for file in files:
            if file.lower().endswith('.pdf'):
                pdf_path = os.path.join(root, file)
                print(f"Processing PDF: {pdf_path}")

                try:
                    pdf_document = fitz.open(pdf_path)
                except Exception as e:
                    print(f"Error opening PDF {pdf_path}: {e}")
                    continue

                extracted_data = []
                for page_num in range(pdf_document.page_count):
                    try:
                        page = pdf_document.load_page(page_num)
                        text = page.get_text()
                    except Exception as e:
                        print(f"Error processing page {page_num + 1} of {pdf_path}: {e}")
                        continue

                    # Generate AI response
                    ai_response = generate_ai_response(model, uploaded_file, prompt)
                    if ai_response:
                        parsed_data = parse_response(ai_response)
                        extracted_data.extend(parsed_data)

                pdf_document.close()

                # Save to database
                for data in extracted_data:
                    file_name = file  # Use the PDF file name
                    product_name = data.get("Product Name", "")
                    part_number = data.get("Part Number", "XXX") if not data.get("Part Number") else data.get("Part Number")
                    voucher_code = data.get("Voucher Code", "")

                    # Check if the file name already exists in the database
                    cursor.execute("SELECT COUNT(1) FROM your_table WHERE file_name = ?", (file_name,))
                    if cursor.fetchone()[0] == 0:
                        cursor.execute("INSERT INTO your_table (file_name, product_name, part_number, voucher_number) VALUES (?, ?, ?, ?)",
                                       (file_name, product_name, part_number, voucher_code))
                        db_connection.commit()
                    else:
                        print(f"Skipping {file_name} as it already exists in the database.")

    cursor.close()



def upload_image(image_path):
    """Uploads the given image file to Gemini and returns the uploaded file object."""
    try:
        from pathlib import Path
        image_path = Path(image_path)
        myfile = genai.upload_file(image_path)
        print(f"Uploaded file: {myfile}")
        return myfile
    except Exception as e:
        print(f"Error uploading image {image_path}: {e}")
        return None


def generate_ai_response(model, uploaded_file, prompt):
    """Generates AI response using the uploaded file and prompt."""
    try:
        inputs = [uploaded_file, "\n\n", prompt]
        result = model.generate_content(inputs)
        print(f"Generated content: {result.text}")
        return result.text
    except Exception as e:
        print(f"Error generating AI response: {e}")
        return ""


def save_response(response_text, response_file_path="response.txt"):
    """Appends the AI response to the response file."""
    try:
        with open(response_file_path, "a", encoding="utf-8") as f:
            f.write("----- New API Response -----\n")
            f.write(response_text + "\n\n")
    except Exception as e:
        print(f"Error writing to {response_file_path}: {e}")


def parse_response(response_text):
    """
    Parses the JSON response text to extract relevant fields.

    Args:
        response_text (str): The raw text response from the API.

    Returns:
        list[dict]: A list of dictionaries containing extracted fields.
    """
    try:
        # Remove unwanted text, if any, and parse JSON
        response_data = response_text.split("----- New API Response -----")[1:]
        extracted_data = []

        for entry in response_data:
            # Remove backticks and load as JSON
            entry = entry.strip().strip("```json").strip("```")
            if entry:
                json_data = json.loads(entry)
                for item in json_data:
                    # Extract required fields
                    product_name = item.get("Product Name", "").replace('"', '').replace(',', '')
                    part_number = item.get("Part Number", "").replace('"', '').replace(',', '')
                    voucher_number = item.get("Voucher Code", "").replace('"', '').replace(',', '')

                    extracted_data.append({
                        "Product Name": product_name,
                        "Part Number": part_number,
                        "Voucher Number": voucher_number
                    })
        return extracted_data
    except json.JSONDecodeError as e:
        print(f"Error decoding JSON: {e}")
        return []


def parse_and_save_data(api_response, pdf_file_name):
    """Parses the API response and saves each entry to the database."""
    for entry in api_response:
        product_name = entry.get("Product Name", "")
        part_number = entry.get("Part Number", "")
        voucher_number = entry.get("Voucher Code", "")
        save_to_db(pdf_file_name, product_name, part_number, voucher_number)


def process_pdf_to_db(pdf_folder, prompt, response_file="response.txt"):
    """Processes PDFs in a folder, extracts data using Gemini, and saves to SQLite DB."""
    # Initialize the response file
    try:
        with open(response_file, "w", encoding="utf-8") as f:
            f.write("API Responses Log\n\n")
    except Exception as e:
        print(f"Error initializing {response_file}: {e}")
        return

    model = create_model()

    # Iterate through PDF files in the folder
    for root, dirs, files in os.walk(pdf_folder):
        for file in files:
            if file.lower().endswith('.pdf'):
                pdf_path = os.path.join(root, file)
                print(f"Processing PDF: {pdf_path}")

                try:
                    pdf_document = fitz.open(pdf_path)
                except Exception as e:
                    print(f"Error opening PDF {pdf_path}: {e}")
                    continue

                for page_num in range(pdf_document.page_count):
                    try:
                        page = pdf_document.load_page(page_num)
                        pix = page.get_pixmap()
                    except Exception as e:
                        print(f"Error processing page {page_num + 1} of {pdf_path}: {e}")
                        continue

                    image_path = f"temp_page_{page_num + 1}.png"
                    try:
                        pix.save(image_path)
                        print(f"Saved temporary image: {image_path}")
                    except Exception as e:
                        print(f"Error saving image {image_path}: {e}")
                        continue

                    uploaded_file = upload_image(image_path)
                    if not uploaded_file:
                        print(f"Skipping page {page_num + 1} due to upload failure.")
                        os.remove(image_path)
                        continue

                    ai_response = generate_ai_response(model, uploaded_file, prompt)
                    if not ai_response:
                        print(f"No response for page {page_num + 1}.")
                        save_response("No response received.\n")
                    else:
                        save_response(ai_response, response_file)
                        parsed_data = parse_response(ai_response)
                        parse_and_save_data([parsed_data], file)

                    try:
                        os.remove(image_path)
                        print(f"Deleted temporary image: {image_path}")
                    except Exception as e:
                        print(f"Error deleting image {image_path}: {e}")

                pdf_document.close()
                print(f"Finished processing PDF: {pdf_path}\n")


def main():
    """Main function to configure API and start processing."""
    configure_api()
    initialize_db()

    pdf_folder = "COA_ DOC_2024"  # Path to folder with PDFs
    prompt_file = "prompt.txt"
    response_file = "response.txt"

    if not os.path.isdir(pdf_folder):
        print(f"Error: The folder '{pdf_folder}' does not exist.")
        return

    if not os.path.isfile(prompt_file):
        print(f"Error: The prompt file '{prompt_file}' does not exist.")
        return

    try:
        with open(prompt_file, 'r', encoding='utf-8') as f:
            prompt = f.read().strip()
        if not prompt:
            print("Error: The prompt file is empty.")
            return
    except Exception as e:
        print(f"Error reading prompt file: {e}")
        return

    process_pdf_to_db(pdf_folder, prompt, response_file)


if __name__ == "__main__":
    main()
D:\Project\NML\doc-flow\venv\Scripts\python.exe "D:\Project\NML\Automation\sopo to coa\main3.py" 
Processing PDF: COA_ DOC_2024\2409015\KK.124.2149.1148a_ PHILIPPINES.pdf
Saved temporary image: temp_page_1.png
Uploaded file: genai.File({
    'name': 'files/c5z3pgcn5ry9',
    'display_name': 'temp_page_1.png',
    'mime_type': 'image/png',
    'sha256_hash': 'MTQyNTI0N2E5YWRlNGYyNzg4Njc5YjM4Yzc4MTNiN2YzZmEyNTdhN2M2YjljOTVkZjljYmRlZmU4MjM3YmZkZQ==',
    'size_bytes': '171930',
    'state': 'ACTIVE',
    'uri': 'https://generativelanguage.googleapis.com/v1beta/files/c5z3pgcn5ry9',
    'create_time': '2024-10-30T10:11:48.508925Z',
    'expiration_time': '2024-11-01T10:11:48.448676452Z',
    'update_time': '2024-10-30T10:11:48.508925Z'})
Traceback (most recent call last):
  File "D:\Project\NML\Automation\sopo to coa\main3.py", line 298, in <module>
    main()
  File "D:\Project\NML\Automation\sopo to coa\main3.py", line 294, in main
    process_pdf_to_db(pdf_folder, prompt, response_file)
  File "D:\Project\NML\Automation\sopo to coa\main3.py", line 255, in process_pdf_to_db
    parse_and_save_data([parsed_data], file)
  File "D:\Project\NML\Automation\sopo to coa\main3.py", line 195, in parse_and_save_data
    product_name = entry.get("Product Name", "")
                   ^^^^^^^^^
AttributeError: 'list' object has no attribute 'get'
Generated content: ```json
[
  {
    "Product Name": "Suction Unit HI-VAC Plus 60 Ltrs/Min",
    "Part Number": "SU56",
    "Voucher Code": "PO-N/11286"
  },
  {
    "Product Name": "Motor Pump For Suction Machine",
    "Part Number": "SU50",
    "Voucher Code": "PO-N/11286"
  }
]
```


Process finished with exit code 1