import gradio as gr
import json
import os
from pathlib import Path
from typing import List, Dict, Any
import google.generativeai as genai
from PIL import Image
import PyPDF2
import tempfile
import traceback

# ==============================================================
# API Configuration - Add your key here
# ==============================================================
GEMINI_API_KEY = "AIzaSyB2b80YwNHs3Yj6RZOTL8wjXk2YhxCluOA"  # Replace with your key

# ==============================================================
# Enhanced extraction prompt with better instructions
# ==============================================================
EXTRACTION_PROMPT = """You are an expert shipping-document data extractor with OCR capabilities.
Carefully analyze ALL text content from PDFs, images, and documents.

CRITICAL: Look at both the text AND the visual layout of documents. Sometimes important data 
is in tables, handwritten notes, stamps, or poorly scanned areas.

Extract and structure the data as valid JSON only (no markdown, no commentary):

{
  "poNumber": string | null,
  "shipFrom": string | null,
  "carrierType": string | null,
  "originCarrier": string | null,
  "railCarNumber": string | null,
  "totalQuantity": number | null,
  "totalUnits": string | null,
  "attachments": [string],
  "accountName": string | null,
  "inventories": {
    "items": [
      {
        "quantityShipped": number | null,
        "inventoryUnits": string | null,
        "pcs": number | null,
        "productName": string | null,
        "productCode": string | null,
        "product": {
          "category": string | null,
          "defaultUnits": string | null,
          "unit": number | null,
          "pcs": number | null,
          "mbf": number | null,
          "sf": number | null,
          "pcsHeight": number | null,
          "pcsWidth": number | null,
          "pcsLength": number | null
        },
        "customFields": [string]
      }
    ]
  }
}

EXTRACTION RULES:
1. Extract ALL product line items - create one inventory item per product
2. Parse dimensions: "2X6X14" → pcsHeight=2, pcsWidth=6, pcsLength=14
3. Convert BF to MBF: BF ÷ 1000
4. customFields format: "Key||Value" (e.g., "Mill||Tolko")
5. Look for: PO numbers, shipping info, quantities, product codes, dimensions
6. Check headers, footers, stamps, handwritten notes, and table cells
7. If multiple documents, consolidate all items into one JSON
8. Return null for missing fields
9. attachments should list all provided filenames

Return ONLY valid JSON matching this exact structure."""


def extract_text_from_pdf(pdf_path: str) -> str:
    """Extract text from PDF with better error handling"""
    try:
        with open(pdf_path, 'rb') as file:
            pdf_reader = PyPDF2.PdfReader(file)
            text = ""
            for page_num, page in enumerate(pdf_reader.pages):
                page_text = page.extract_text()
                if page_text:
                    text += f"\n--- Page {page_num + 1} ---\n{page_text}"
            return text if text.strip() else "No text extracted from PDF"
    except Exception as e:
        return f"Error extracting PDF text: {str(e)}"


def process_files_for_gemini(files: List[str]) -> Dict[str, Any]:
    """Process files and prepare for Gemini multimodal input"""
    processed_data = {
        "text_content": "",
        "file_objects": [],
        "attachments": [],
        "file_info": []
    }
    
    if not files:
        return processed_data
    
    for file_path in files:
        if not os.path.exists(file_path):
            continue
            
        file_name = Path(file_path).name
        file_ext = Path(file_path).suffix.lower()
        
        processed_data["attachments"].append(file_name)
        processed_data["file_info"].append(f"File: {file_name} (Type: {file_ext})")
        
        try:
            # Handle PDFs
            if file_ext == '.pdf':
                text = extract_text_from_pdf(file_path)
                processed_data["text_content"] += f"\n\n=== {file_name} ===\n{text}"
                
                # Upload PDF to Gemini for visual analysis
                uploaded_file = genai.upload_file(file_path)
                processed_data["file_objects"].append(uploaded_file)
            
            # Handle images
            elif file_ext in ['.jpg', '.jpeg', '.png', '.gif', '.bmp', '.webp']:
                # Upload image to Gemini
                uploaded_file = genai.upload_file(file_path)
                processed_data["file_objects"].append(uploaded_file)
                processed_data["text_content"] += f"\n\n=== {file_name} (Image) ===\n[Image uploaded for visual analysis]"
            
            # Handle text files
            elif file_ext in ['.txt', '.csv']:
                with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
                    text = f.read()
                processed_data["text_content"] += f"\n\n=== {file_name} ===\n{text}"
            
            # Handle Word documents (basic text extraction)
            elif file_ext in ['.doc', '.docx']:
                try:
                    import docx
                    doc = docx.Document(file_path)
                    text = "\n".join([para.text for para in doc.paragraphs])
                    processed_data["text_content"] += f"\n\n=== {file_name} ===\n{text}"
                except ImportError:
                    processed_data["text_content"] += f"\n\n=== {file_name} ===\n[Word document - install python-docx for text extraction]"
                except Exception as e:
                    processed_data["text_content"] += f"\n\n=== {file_name} ===\nError reading Word doc: {str(e)}"
        
        except Exception as e:
            processed_data["text_content"] += f"\n\n=== {file_name} ===\nError processing: {str(e)}"
    
    return processed_data


def extract_with_gemini(processed_data: Dict[str, Any], api_key: str, model_name: str = "gemini-2.0-flash-exp") -> Dict[str, Any]:
    """Extract structured data using Gemini with enhanced multimodal processing"""
    
    if not api_key or api_key.strip() == "":
        return {
            "success": False,
            "error": "Gemini API key not provided"
        }
    
    try:
        # Configure Gemini
        genai.configure(api_key=api_key)
        
        # Use the latest model with vision capabilities
        model = genai.GenerativeModel(model_name)
        
        # Build multimodal prompt
        content_parts = [
            EXTRACTION_PROMPT,
            f"\n\nDOCUMENT CONTEXT:\n{processed_data['text_content']}\n",
            f"\nATTACHMENTS: {json.dumps(processed_data['attachments'])}\n",
            "\nNow analyze the uploaded files carefully (including visual content) and extract the data as JSON:"
        ]
        
        # Add all uploaded files
        content_parts.extend(processed_data["file_objects"])
        
        # Generate with higher temperature for better extraction
        generation_config = genai.types.GenerationConfig(
            temperature=0.2,
            max_output_tokens=8000,
        )
        
        response = model.generate_content(
            content_parts,
            generation_config=generation_config
        )
        
        response_text = response.text.strip()
        
        # Clean markdown code blocks
        if response_text.startswith("```json"):
            response_text = response_text[7:]
        elif response_text.startswith("```"):
            response_text = response_text[3:]
        if response_text.endswith("```"):
            response_text = response_text[:-3]
        
        response_text = response_text.strip()
        
        # Parse JSON
        extracted_data = json.loads(response_text)
        
        return {
            "success": True,
            "data": extracted_data,
            "raw_response": response_text,
            "files_processed": len(processed_data["file_objects"])
        }
    
    except json.JSONDecodeError as e:
        return {
            "success": False,
            "error": f"JSON parsing error: {str(e)}",
            "raw_response": response.text if 'response' in locals() else "No response",
            "suggestion": "The AI returned non-JSON text. Try again or check the raw response."
        }
    except Exception as e:
        return {
            "success": False,
            "error": f"Extraction error: {str(e)}",
            "traceback": traceback.format_exc()
        }


def process_documents(files):
    """Main Gradio processing function"""
    
    if not files or len(files) == 0:
        return "❌ Error: Please upload at least one file", "{}", "No files provided"
    
    # Use the hardcoded API key and default model
    api_key = GEMINI_API_KEY
    model_choice = "gemini-2.0-flash-exp"
    
    if not api_key or api_key.strip() == "":
        return "❌ Error: API key not configured in code", "{}", "API key missing"
    
    try:
        # Get file paths
        file_paths = [f.name if hasattr(f, 'name') else f for f in files]
        
        status_msg = f"📄 Processing {len(file_paths)} file(s)...\n"
        
        # Process files
        processed_data = process_files_for_gemini(file_paths)
        status_msg += f"✓ Files loaded: {', '.join(processed_data['attachments'])}\n"
        
        # Extract with Gemini
        status_msg += "🤖 Extracting data with Gemini AI...\n"
        result = extract_with_gemini(processed_data, api_key, model_choice)
        
        if result.get("success"):
            json_output = json.dumps(result["data"], indent=2)
            status_msg += f"✅ Extraction successful! Processed {result.get('files_processed', 0)} files.\n"
            
            # Format display output
            display_text = "=== EXTRACTED DATA ===\n\n"
            display_text += json_output
            
            return status_msg, json_output, display_text
        else:
            error_msg = f"❌ Extraction failed:\n{result.get('error', 'Unknown error')}\n"
            if 'suggestion' in result:
                error_msg += f"\n💡 {result['suggestion']}\n"
            if 'traceback' in result:
                error_msg += f"\nDebug info:\n{result['traceback'][:500]}"
            
            raw_resp = result.get('raw_response', 'No response')
            return error_msg, "{}", f"Raw Response:\n{raw_resp[:1000]}"
    
    except Exception as e:
        error_msg = f"❌ Unexpected error: {str(e)}\n{traceback.format_exc()[:500]}"
        return error_msg, "{}", error_msg


# ==============================================================
# Gradio Interface
# ==============================================================

def create_interface():
    with gr.Blocks(theme=gr.themes.Soft(), title="Document Data Extractor") as demo:
        gr.Markdown("""
        # 📄 Shipping Document Data Extractor
        
        Upload PDFs, images, Word docs, or text files to extract structured shipping data using Google Gemini AI.
        
        **Supported formats:** PDF, JPG, PNG, DOCX, TXT, CSV
        """)
        
        with gr.Row():
            with gr.Column(scale=2):
                file_input = gr.File(
                    label="📎 Upload Documents",
                    file_count="multiple",
                    file_types=[".pdf", ".jpg", ".jpeg", ".png", ".gif", ".bmp", ".txt", ".csv", ".doc", ".docx"]
                )
                
                submit_btn = gr.Button("🚀 Extract Data", variant="primary", size="lg")
            
            with gr.Column(scale=3):
                status_output = gr.Textbox(
                    label="📊 Status",
                    lines=4,
                    max_lines=8
                )
                
                json_output = gr.Code(
                    label="📋 JSON Output (Copy this)",
                    language="json",
                    lines=15
                )
                
                display_output = gr.Textbox(
                    label="👁️ Preview",
                    lines=10,
                    max_lines=15
                )
        
        gr.Markdown("""
        ### 💡 Tips:
        - Upload multiple files for batch processing
        - For images: ensure text is clear and well-lit
        - For PDFs: both text-based and scanned PDFs work
        - The AI will analyze visual content even if text extraction fails
        """)
        
        # Button action
        submit_btn.click(
            fn=process_documents,
            inputs=[file_input],
            outputs=[status_output, json_output, display_output]
        )
        
        # Examples
        gr.Examples(
            examples=[
                [["example1.pdf"]],
            ],
            inputs=[file_input],
            label="Example Usage"
        )
    
    return demo


if __name__ == "__main__":
    demo = create_interface()
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False
    )