from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware
import httpx
import io
import os
from typing import Optional, Dict, List, Any
from PIL import Image
import fitz  # PyMuPDF
from docx import Document
import base64
import json
from datetime import datetime

app = FastAPI(
    title="OCR API",
    description="OCR API using GLM-4V for document text extraction",
    version="1.0.0"
)

# Add CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Configuration
HF_API_TOKEN = os.getenv("HF_API_TOKEN", "")
MODEL_ID = "PrathameshRaut/DocumentOCR"
# Updated to new HF router endpoint (api-inference.huggingface.co is deprecated)
HF_INFERENCE_URL = f"https://router.huggingface.co/models/{MODEL_ID}"

# Supported file extensions
SUPPORTED_EXTENSIONS = {
    "image": [".jpg", ".jpeg", ".png", ".gif", ".bmp", ".webp"],
    "pdf": [".pdf"],
    "document": [".docx", ".doc"]
}


class FileProcessor:
    """Handle conversion of various file types to images"""
    
    @staticmethod
    def pdf_to_images(pdf_bytes: bytes) -> List[Image.Image]:
        """Convert PDF to list of PIL Images"""
        try:
            pdf_document = fitz.open(stream=pdf_bytes, filetype="pdf")
            images = []
            for page_num in range(len(pdf_document)):
                page = pdf_document[page_num]
                pix = page.get_pixmap(matrix=fitz.Matrix(2, 2))  # 2x zoom for better quality
                img_data = pix.tobytes("ppm")
                img = Image.open(io.BytesIO(img_data))
                images.append(img.convert("RGB"))
            pdf_document.close()
            return images
        except Exception as e:
            raise HTTPException(status_code=400, detail=f"PDF processing error: {str(e)}")
    
    @staticmethod
    def docx_to_images(docx_bytes: bytes) -> List[Image.Image]:
        """Convert DOCX to images (extracts text and creates image representation)"""
        try:
            doc = Document(io.BytesIO(docx_bytes))
            # Extract all text
            full_text = "\n".join([para.text for para in doc.paragraphs])
            
            # For DOCX, we'll create a text image representation
            # In production, you might want to use a more sophisticated approach
            img = Image.new('RGB', (1200, 800), color='white')
            
            # Save text as image (basic approach)
            from PIL import ImageDraw, ImageFont
            draw = ImageDraw.Draw(img)
            
            # Try to use a default font, fallback if not available
            try:
                font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", 14)
            except:
                font = ImageFont.load_default()
            
            # Wrap text and draw
            y_position = 20
            for line in full_text.split('\n')[:50]:  # Limit to 50 lines
                if y_position > 750:
                    break
                draw.text((20, y_position), line[:100], fill='black', font=font)
                y_position += 20
            
            return [img]
        except Exception as e:
            raise HTTPException(status_code=400, detail=f"DOCX processing error: {str(e)}")
    
    @staticmethod
    def image_to_pil(image_bytes: bytes) -> List[Image.Image]:
        """Convert image bytes to PIL Image"""
        try:
            img = Image.open(io.BytesIO(image_bytes)).convert("RGB")
            return [img]
        except Exception as e:
            raise HTTPException(status_code=400, detail=f"Image processing error: {str(e)}")


class GLMOCRClient:
    """Handle communication with GLM-OCR model"""
    
    def __init__(self, model_id: str, api_token: str):
        self.model_id = model_id
        self.api_token = api_token
        # Using new HF router endpoint (api-inference.huggingface.co is deprecated as of 2026)
        self.base_url = f"https://router.huggingface.co/models/{model_id}"
    
    def _image_to_base64(self, image: Image.Image) -> str:
        """Convert PIL Image to base64 string"""
        buffer = io.BytesIO()
        image.save(buffer, format="PNG")
        buffer.seek(0)
        return base64.b64encode(buffer.read()).decode('utf-8')
    
    async def _call_model(self, image: Image.Image, prompt: str) -> str:
        """Generic model call method"""
        try:
            image_data = self._image_to_base64(image)
            
            # Correct format for GLM-4V via HF Inference API
            payload = {
                "inputs": {
                    "image": f"data:image/png;base64,{image_data}",
                    "text": prompt
                }
            }
            
            headers = {
                "Authorization": f"Bearer {self.api_token}",
                "Content-Type": "application/json"
            }
            
            async with httpx.AsyncClient(timeout=120.0) as client:
                response = await client.post(
                    self.base_url,
                    json=payload,
                    headers=headers
                )
            
            print(f"[DEBUG] Model response status: {response.status_code}")
            
            if response.status_code == 200:
                result = response.json()
                print(f"[DEBUG] Model response: {result}")
                
                # Handle different response formats
                if isinstance(result, dict):
                    # Try different key names that might contain the text
                    for key in ["generated_text", "text", "output", "result", "response"]:
                        if key in result:
                            text_result = result[key]
                            if isinstance(text_result, list) and len(text_result) > 0:
                                return str(text_result[0]).strip()
                            return str(text_result).strip()
                    # If no recognized key, return str of result
                    return str(result).strip()
                elif isinstance(result, list) and len(result) > 0:
                    # Handle list response
                    item = result[0]
                    if isinstance(item, dict):
                        for key in ["generated_text", "text", "output"]:
                            if key in item:
                                return str(item[key]).strip()
                    return str(item).strip()
                else:
                    return str(result).strip()
            else:
                print(f"[DEBUG] Model error - Status {response.status_code}: {response.text}")
                return ""
        except Exception as e:
            print(f"[DEBUG] Model call error: {str(e)}")
            import traceback
            traceback.print_exc()
            return ""
    
    async def extract_tables(self, image: Image.Image) -> str:
        """Extract table content from image"""
        prompt = "Please extract and return all tables from this document in a structured format."
        return await self._call_model(image, prompt)
    
    async def extract_text(self, image: Image.Image) -> str:
        """Extract text content from image"""
        prompt = "Please extract and return all the text content from this document."
        return await self._call_model(image, prompt)
    
    async def extract_formulas(self, image: Image.Image) -> str:
        """Extract formula content from image"""
        prompt = "Please extract and return all mathematical formulas and equations from this document."
        return await self._call_model(image, prompt)
    
    async def extract_structured_info(self, image: Image.Image, schema: Dict) -> Dict:
        """Extract structured information based on schema"""
        try:
            image_data = self._image_to_base64(image)
            
            # Create prompt for structured extraction
            schema_json = json.dumps(schema, ensure_ascii=False, indent=2)
            prompt = f"Please extract and return information from this document in the following JSON format:\n{schema_json}"
            
            payload = {
                "inputs": {
                    "image": f"data:image/png;base64,{image_data}",
                    "text": prompt
                }
            }
            
            headers = {
                "Authorization": f"Bearer {self.api_token}",
                "Content-Type": "application/json"
            }
            
            async with httpx.AsyncClient(timeout=120.0) as client:
                response = await client.post(
                    self.base_url,
                    json=payload,
                    headers=headers
                )
            
            print(f"[DEBUG] Structured extraction response status: {response.status_code}")
            
            if response.status_code == 200:
                result = response.json()
                print(f"[DEBUG] Structured extraction response: {result}")
                
                # Extract text from response
                result_text = ""
                if isinstance(result, dict):
                    for key in ["generated_text", "text", "output"]:
                        if key in result:
                            val = result[key]
                            result_text = str(val[0]) if isinstance(val, list) else str(val)
                            break
                    if not result_text:
                        result_text = str(result)
                elif isinstance(result, list) and len(result) > 0:
                    item = result[0]
                    if isinstance(item, dict):
                        for key in ["generated_text", "text", "output"]:
                            if key in item:
                                result_text = str(item[key])
                                break
                    if not result_text:
                        result_text = str(item)
                else:
                    result_text = str(result)
                
                # Try to parse JSON response
                try:
                    # Clean up the response if needed
                    cleaned = result_text.strip()
                    if cleaned.startswith("```json"):
                        cleaned = cleaned[7:]
                    if cleaned.endswith("```"):
                        cleaned = cleaned[:-3]
                    cleaned = cleaned.strip()
                    
                    return json.loads(cleaned)
                except json.JSONDecodeError:
                    return {"raw_response": result_text}
            else:
                print(f"[DEBUG] Structured extraction error - Status {response.status_code}: {response.text}")
                return {"error": f"Extraction failed with status {response.status_code}"}
        except Exception as e:
            print(f"[DEBUG] Structured extraction error: {str(e)}")
            import traceback
            traceback.print_exc()
            return {"error": str(e)}


@app.get("/health")
async def health_check():
    """Health check endpoint"""
    hf_token_status = "✅ Set" if HF_API_TOKEN else "❌ Missing"
    return {
        "status": "healthy",
        "timestamp": datetime.now().isoformat(),
        "model": MODEL_ID,
        "hf_token_status": hf_token_status,
        "debug_test": "/debug/test-model"
    }


@app.get("/debug/test-model")
async def test_model_connection():
    """Test connection to GLM-OCR model - creates a test image"""
    try:
        # Create a simple test image
        test_image = Image.new('RGB', (200, 100), color='white')
        from PIL import ImageDraw
        draw = ImageDraw.Draw(test_image)
        draw.text((20, 40), "Test OCR", fill='black')
        
        # Initialize OCR client
        ocr_client = GLMOCRClient(MODEL_ID, HF_API_TOKEN)
        
        if not HF_API_TOKEN:
            return {
                "status": "error",
                "message": "HF_API_TOKEN not set",
                "solution": "Set HF_API_TOKEN environment variable or add to Space Secrets",
                "details": "Export: export HF_API_TOKEN='your_token' (local) or add in HF Space Settings → Secrets"
            }
        
        # Test extraction
        text_result = await ocr_client.extract_text(test_image)
        
        return {
            "status": "success",
            "message": "Model connection test completed",
            "test_extraction_result": text_result if text_result else "(Empty - check model response)",
            "model": MODEL_ID,
            "api_endpoint": f"https://router.huggingface.co/models/{MODEL_ID}",
            "api_token_set": bool(HF_API_TOKEN),
            "next_steps": "Check logs above for [DEBUG] messages to see actual API response format"
        }
    except Exception as e:
        print(f"[DEBUG] Test failed: {str(e)}")
        import traceback
        traceback.print_exc()
        return {
            "status": "error",
            "message": str(e),
            "error_type": type(e).__name__,
            "model": MODEL_ID,
            "api_token_set": bool(HF_API_TOKEN),
            "suggestion": "Check the logs above for [DEBUG] output to debug the issue"
        }


@app.post("/api/ocr/extract")
async def extract_text(file: UploadFile = File(...)):
    """
    Extract text, tables, and formulas from uploaded document.
    
    Returns results in order: Table -> Text -> Formula
    
    Supported formats:
    - Images: jpg, jpeg, png, gif, bmp, webp
    - PDFs: pdf
    - Documents: docx, doc
    """
    try:
        # Validate file extension
        file_ext = os.path.splitext(file.filename)[1].lower()
        
        # Read file content
        content = await file.read()
        
        if not content:
            raise HTTPException(status_code=400, detail="Empty file")
        
        # Process file based on type
        images = []
        
        if file_ext in SUPPORTED_EXTENSIONS["image"]:
            images = FileProcessor.image_to_pil(content)
        elif file_ext in SUPPORTED_EXTENSIONS["pdf"]:
            images = FileProcessor.pdf_to_images(content)
        elif file_ext in SUPPORTED_EXTENSIONS["document"]:
            images = FileProcessor.docx_to_images(content)
        else:
            raise HTTPException(
                status_code=400,
                detail=f"Unsupported file type. Supported: {', '.join([ext for exts in SUPPORTED_EXTENSIONS.values() for ext in exts])}"
            )
        
        # Initialize OCR client
        ocr_client = GLMOCRClient(MODEL_ID, HF_API_TOKEN)
        
        # Process each page/image
        results = {
            "filename": file.filename,
            "file_type": file_ext,
            "total_pages": len(images),
            "extraction_timestamp": datetime.now().isoformat(),
            "pages": []
        }
        
        for page_num, image in enumerate(images, 1):
            page_result = {
                "page_number": page_num,
                "table": await ocr_client.extract_tables(image),
                "text": await ocr_client.extract_text(image),
                "formula": await ocr_client.extract_formulas(image)
            }
            results["pages"].append(page_result)
        
        return JSONResponse(
            status_code=200,
            content=results
        )
    
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Processing error: {str(e)}")


@app.post("/api/ocr/extract-structured")
async def extract_structured(
    file: UploadFile = File(...),
    schema: Optional[str] = None
):
    """
    Extract structured information from document based on provided JSON schema.
    
    Query Parameters:
    - schema: JSON string defining the structure to extract
    
    Example schema:
    {
        "id_number": "",
        "name": "",
        "date_of_birth": ""
    }
    """
    try:
        # Validate file extension
        file_ext = os.path.splitext(file.filename)[1].lower()
        
        # Read file content
        content = await file.read()
        
        if not content:
            raise HTTPException(status_code=400, detail="Empty file")
        
        # Parse schema
        if not schema:
            raise HTTPException(status_code=400, detail="Schema parameter is required")
        
        try:
            schema_dict = json.loads(schema)
        except json.JSONDecodeError:
            raise HTTPException(status_code=400, detail="Invalid JSON schema")
        
        # Process file based on type
        images = []
        
        if file_ext in SUPPORTED_EXTENSIONS["image"]:
            images = FileProcessor.image_to_pil(content)
        elif file_ext in SUPPORTED_EXTENSIONS["pdf"]:
            images = FileProcessor.pdf_to_images(content)
        elif file_ext in SUPPORTED_EXTENSIONS["document"]:
            images = FileProcessor.docx_to_images(content)
        else:
            raise HTTPException(
                status_code=400,
                detail=f"Unsupported file type. Supported: {', '.join([ext for exts in SUPPORTED_EXTENSIONS.values() for ext in exts])}"
            )
        
        # Initialize OCR client
        ocr_client = GLMOCRClient(MODEL_ID, HF_API_TOKEN)
        
        # Process each page/image
        results = {
            "filename": file.filename,
            "file_type": file_ext,
            "total_pages": len(images),
            "extraction_timestamp": datetime.now().isoformat(),
            "pages": []
        }
        
        for page_num, image in enumerate(images, 1):
            page_result = {
                "page_number": page_num,
                "extracted_data": await ocr_client.extract_structured_info(image, schema_dict)
            }
            results["pages"].append(page_result)
        
        return JSONResponse(
            status_code=200,
            content=results
        )
    
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Processing error: {str(e)}")


@app.post("/api/ocr/batch")
async def batch_extract(files: List[UploadFile] = File(...)):
    """
    Process multiple files in batch.
    Returns results for all files in a single response.
    """
    try:
        if not files:
            raise HTTPException(status_code=400, detail="No files provided")
        
        if len(files) > 10:
            raise HTTPException(status_code=400, detail="Maximum 10 files per batch")
        
        ocr_client = GLMOCRClient(MODEL_ID, HF_API_TOKEN)
        batch_results = {
            "batch_timestamp": datetime.now().isoformat(),
            "total_files": len(files),
            "results": []
        }
        
        for file in files:
            try:
                file_ext = os.path.splitext(file.filename)[1].lower()
                content = await file.read()
                
                if not content:
                    batch_results["results"].append({
                        "filename": file.filename,
                        "status": "error",
                        "error": "Empty file"
                    })
                    continue
                
                # Process file
                images = []
                
                if file_ext in SUPPORTED_EXTENSIONS["image"]:
                    images = FileProcessor.image_to_pil(content)
                elif file_ext in SUPPORTED_EXTENSIONS["pdf"]:
                    images = FileProcessor.pdf_to_images(content)
                elif file_ext in SUPPORTED_EXTENSIONS["document"]:
                    images = FileProcessor.docx_to_images(content)
                else:
                    batch_results["results"].append({
                        "filename": file.filename,
                        "status": "error",
                        "error": f"Unsupported file type: {file_ext}"
                    })
                    continue
                
                # Extract from all pages
                file_results = {
                    "filename": file.filename,
                    "file_type": file_ext,
                    "status": "success",
                    "total_pages": len(images),
                    "pages": []
                }
                
                for page_num, image in enumerate(images, 1):
                    page_result = {
                        "page_number": page_num,
                        "table": await ocr_client.extract_tables(image),
                        "text": await ocr_client.extract_text(image),
                        "formula": await ocr_client.extract_formulas(image)
                    }
                    file_results["pages"].append(page_result)
                
                batch_results["results"].append(file_results)
            
            except Exception as e:
                batch_results["results"].append({
                    "filename": file.filename,
                    "status": "error",
                    "error": str(e)
                })
        
        return JSONResponse(
            status_code=200,
            content=batch_results
        )
    
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Batch processing error: {str(e)}")


@app.get("/api/info")
async def api_info():
    """Get API information and supported formats"""
    return {
        "api_name": "GLM-OCR API",
        "version": "1.0.0",
        "model": MODEL_ID,
        "supported_formats": {
            "images": SUPPORTED_EXTENSIONS["image"],
            "documents": SUPPORTED_EXTENSIONS["document"],
            "pdfs": SUPPORTED_EXTENSIONS["pdf"]
        },
        "endpoints": {
            "extract": "/api/ocr/extract",
            "extract_structured": "/api/ocr/extract-structured",
            "batch": "/api/ocr/batch",
            "health": "/health",
            "info": "/api/info"
        },
        "extraction_order": ["table", "text", "formula"]
    }


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(
        app,
        host="0.0.0.0",
        port=7860
    )