import os
import json
import io
import time
import gc
import numpy as np
import pandas as pd
from flask import Flask, request, send_file, jsonify
from flask_cors import CORS
import google.generativeai as genai
from groq import Groq
from PyPDF2 import PdfReader
from dotenv import load_dotenv

# OCR Imports (EasyOCR)
import easyocr
from pdf2image import convert_from_bytes, pdfinfo_from_bytes


app = Flask(__name__)
CORS(app, resources={r"/*": {"origins": "*"}})

# Load .env locally
basedir = os.path.abspath(os.path.dirname(__file__))
load_dotenv(os.path.join(basedir, ".env"))

# API Keys
GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")
GROQ_API_KEY = os.environ.get("GROQ_API_KEY")

print(f"DEBUG - Is Groq Key loaded? {'YES' if GROQ_API_KEY else 'NO'}")

genai.configure(api_key=GEMINI_API_KEY)
groq_client = Groq(api_key=GROQ_API_KEY)

# EasyOCR reader (load once)
easyocr_reader = easyocr.Reader(["en"], gpu=False)

# OCR settings
OCR_DPI = 200

# Groq batching settings
BATCH_SIZE = 3
MAX_RETRIES = 5


SYSTEM_PROMPT = """
You are a data extraction assistant. Extract the payroll data from the provided document or text.
Return ONLY a valid JSON object with a single key called "data". The value of "data" must be an array of objects.
Do not include markdown formatting, code blocks, or explanations.
Each object in the array must have exactly these keys: "SSN", "Participant", "Amount", "Contract", "Carrier Name".
If a specific field is missing for a participant, use null.

Example format:
{
  "data": [
    {"SSN": "123456789", "Participant": "John Doe", "Amount": "150.00", "Contract": null, "Carrier Name": "MG Trust Co"}
  ]
}
"""


@app.route("/")
def home():
    return jsonify({"status": "API running"}), 200


def generate_excel_buffer(json_data):
    """Convert JSON extracted data into an Excel file buffer."""
    records = json_data.get("data", [])
    df = pd.DataFrame(records)

    expected_columns = ["SSN", "Participant", "Amount", "Contract", "Carrier Name"]
    for col in expected_columns:
        if col not in df.columns:
            df[col] = None

    df = df[expected_columns]

    output = io.BytesIO()
    with pd.ExcelWriter(output, engine="openpyxl") as writer:
        df.to_excel(writer, index=False, sheet_name="Payroll Data")

    output.seek(0)
    return output


def extract_text_with_pypdf2(pdf_file):
    """Extract text from PDF using PyPDF2."""
    pdf_reader = PdfReader(pdf_file)
    extracted_text = ""

    for page in pdf_reader.pages:
        text = page.extract_text()
        if text:
            extracted_text += text + "\n"

    return extracted_text.strip()


def extract_text_with_easyocr_pagewise(pdf_bytes):
    """Extract OCR text page by page using EasyOCR (memory safe)."""
    info = pdfinfo_from_bytes(pdf_bytes)
    total_pages = info["Pages"]

    pages_text = []

    for i in range(1, total_pages + 1):
        print(f"🔍 OCR Processing Page {i}/{total_pages}")

        images = convert_from_bytes(
            pdf_bytes,
            dpi=OCR_DPI,
            first_page=i,
            last_page=i
        )

        img = images[0]
        img_np = np.array(img)

        result = easyocr_reader.readtext(img_np, detail=0, paragraph=True)
        page_text = "\n".join(result).strip()

        pages_text.append(page_text)

        # Free memory
        del images
        del img
        del img_np
        gc.collect()

    return pages_text


def call_groq_with_retry(text_block):
    """Call Groq with retry support (handles 429)."""
    for attempt in range(MAX_RETRIES):
        try:
            chat_completion = groq_client.chat.completions.create(
                messages=[
                    {"role": "system", "content": SYSTEM_PROMPT.replace("\xa0", " ")},
                    {"role": "user", "content": f"Extract payroll rows ONLY from this text:\n\n{text_block}"},
                ],
                model="llama-3.1-8b-instant",
                response_format={"type": "json_object"},
            )

            response_text = chat_completion.choices[0].message.content
            clean_text = response_text.replace("\xa0", " ").strip()

            return json.loads(clean_text)

        except Exception as e:
            error_msg = str(e)

            if "429" in error_msg:
                wait_time = (attempt + 1) * 3
                print(f"⚠️ Groq Rate Limit Hit (429). Waiting {wait_time}s...")
                time.sleep(wait_time)
                continue

            raise e

    raise Exception("Groq API failed after multiple retries due to rate limit.")


def clean_rows(rows):
    """
    Merge duplicate rows smartly.
    If same Participant+Amount appears again and SSN is missing in old row,
    update it using new row.
    """

    merged = {}
    
    for row in rows:
        ssn = str(row.get("SSN", "")).strip() if row.get("SSN") else ""
        participant = str(row.get("Participant", "")).strip() if row.get("Participant") else ""
        amount = str(row.get("Amount", "")).strip() if row.get("Amount") else ""
        contract = str(row.get("Contract", "")).strip() if row.get("Contract") else ""
        carrier = str(row.get("Carrier Name", "")).strip() if row.get("Carrier Name") else ""

        # skip completely empty rows
        if not participant and not ssn:
            continue

        # Main key: participant + amount (because SSN may appear later)
        key = (participant.lower(), amount)

        if key not in merged:
            merged[key] = {
                "SSN": ssn if ssn else None,
                "Participant": participant if participant else None,
                "Amount": amount if amount else None,
                "Contract": contract if contract else None,
                "Carrier Name": carrier if carrier else None,
            }
        else:
            # merge missing fields
            if not merged[key]["SSN"] and ssn:
                merged[key]["SSN"] = ssn

            if not merged[key]["Contract"] and contract:
                merged[key]["Contract"] = contract

            if not merged[key]["Carrier Name"] and carrier:
                merged[key]["Carrier Name"] = carrier

            # if participant missing (rare case)
            if not merged[key]["Participant"] and participant:
                merged[key]["Participant"] = participant

    return list(merged.values())


@app.route("/download_with_image_model", methods=["POST", "OPTIONS"])
def process_with_gemini():
    """Endpoint for Gemini (Image Model) PDF processing."""
    if "file" not in request.files:
        return jsonify({"error": "No file uploaded"}), 400

    file = request.files["file"]
    if file.filename == "":
        return jsonify({"error": "No selected file"}), 400

    try:
        file_bytes = file.read()
        mime_type = file.mimetype

        model = genai.GenerativeModel("gemini-2.5-flash")

        contents = [
            SYSTEM_PROMPT,
            {"mime_type": mime_type, "data": file_bytes},
        ]

        response = model.generate_content(contents)

        response_text = (
            response.text.strip()
            .removeprefix("```json")
            .removesuffix("```")
            .strip()
        )

        parsed_data = json.loads(response_text)

        excel_buffer = generate_excel_buffer(parsed_data)

        return send_file(
            excel_buffer,
            mimetype="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
            as_attachment=True,
            download_name="Extracted_Data_Gemini.xlsx",
        )

    except Exception as e:
        return jsonify({"error": str(e)}), 500
        
@app.route("/download_with_text_model", methods=["POST", "OPTIONS"])
def process_with_groq():
    """Groq endpoint with OCR fallback + batching to avoid 429."""
    if "file" not in request.files:
        return jsonify({"error": "No file uploaded"}), 400

    file = request.files["file"]
    if file.filename == "":
        return jsonify({"error": "No selected file"}), 400

    try:
        pdf_bytes = file.read()

        # Reset pointer for PyPDF2
        file.seek(0)

        extracted_text = extract_text_with_pypdf2(file)

        all_rows = []

        # Case 1: Text PDF
        if len(extracted_text.strip()) > 150:
            print("✅ Text detected using PyPDF2. Sending full document to Groq...")

            parsed_data = call_groq_with_retry(extracted_text)
            all_rows.extend(parsed_data.get("data", []))

        # Case 2: Scanned PDF
        else:
            print("⚠️ No proper text found. Running EasyOCR page-by-page...")

            pages_text = extract_text_with_easyocr_pagewise(pdf_bytes)

            # Batch pages to reduce Groq calls
            for i in range(0, len(pages_text), BATCH_SIZE):
                batch_pages = pages_text[i:i + BATCH_SIZE]
                batch_text = "\n\n--- PAGE BREAK ---\n\n".join(batch_pages).strip()

                if len(batch_text) < 50:
                    continue

                print(f"📤 Sending Pages {i+1} to {min(i+BATCH_SIZE, len(pages_text))} to Groq...")

                parsed_batch = call_groq_with_retry(batch_text)
                all_rows.extend(parsed_batch.get("data", []))

        all_rows = clean_rows(all_rows)

        if not all_rows:
            return jsonify({"error": "No payroll rows extracted from document."}), 400

        final_json = {"data": all_rows}
        excel_buffer = generate_excel_buffer(final_json)

        return send_file(
            excel_buffer,
            mimetype="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
            as_attachment=True,
            download_name="Extracted_Data_Groq.xlsx",
        )

    except Exception as e:
        print(f"\n--- CRITICAL BACKEND ERROR --- \n{str(e)}\n------------------------------\n")
        return jsonify({"error": str(e)}), 500


if __name__ == "__main__":
    port = int(os.environ.get("PORT", 7860))
    app.run(host="0.0.0.0", port=port)