import io
import json
import requests
from datetime import datetime
from flask import Flask, request, jsonify
from google import genai
from pydantic import BaseModel

app = Flask(__name__)

# --- CONFIGURATION ---
# IMPORTANT: Put your actual Gemini API Key here
GEMINI_API_KEY = "AIzaSyAZaOGUdpO8trraRZq2pZveirAIvpNaAyk"
client = genai.Client(api_key=GEMINI_API_KEY)

class DocumentData(BaseModel):
    title: str
    document_date: str
    gstin: str
    pan_number: str
    total_amount_due: str
    summary: str

# Helper function to write logs to a file
def write_log(message):
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    with open("debug_log.txt", "a") as f:
        f.write(f"[{timestamp}] {message}\n")

@app.route('/extract', methods=['POST'])
def process_pdf():
    write_log("--- NEW REQUEST STARTED ---")
    try:
        data = request.get_json()
        pdf_url = data.get('pdf_url')
        
        if not pdf_url:
            write_log("Error: No PDF URL provided")
            return jsonify({"error": "Please provide a pdf_url"}), 400

        write_log(f"Attempting to download PDF from: {pdf_url}")
        
        # Download the PDF
        response = requests.get(pdf_url, timeout=30)
        response.raise_for_status()
        pdf_bytes = io.BytesIO(response.content)
        write_log(f"PDF downloaded successfully. Size: {len(response.content)} bytes")

        write_log("Sending PDF to Gemini 3 Pro... (This may take a moment)")
        
        # Call Gemini
        result = client.models.generate_content(
            model="gemini-2.5-flash",
            contents=[
                "Extract the following details from this document precisely. If a financial field is missing, output 'N/A'.",
                genai.types.Part.from_bytes(
                    data=pdf_bytes.getvalue(), 
                    mime_type='application/pdf'
                )
            ],
            config=genai.types.GenerateContentConfig(
                response_mime_type='application/json',
                response_schema=DocumentData,
            ),
        )
        write_log("Received successful response from Gemini!")

        # Save the actual JSON output to a file so you can check it later
        output_data = json.loads(result.text)
        with open("last_extraction_output.json", "w") as f:
            json.dump(output_data, f, indent=4)
            
        write_log("Saved extracted data to last_extraction_output.json")
        write_log("--- REQUEST FINISHED SUCCESSFULLY ---")

        return jsonify(output_data)

    except Exception as e:
        write_log(f"CRITICAL ERROR: {str(e)}")
        return jsonify({"error": str(e)}), 500

if __name__ == "__main__":
    app.run()