import os
import tempfile
import logging
import asyncio
from typing import List, Optional
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor
import time

from fastapi import FastAPI, File, UploadFile, HTTPException, BackgroundTasks, Depends, Header
from fastapi.responses import JSONResponse
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from pydantic import BaseModel
import pytesseract
from PIL import Image
from pdf2image import convert_from_path
import uvicorn
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Environment variables
DEFAULT_MAX_WORKERS = int(os.getenv("DEFAULT_MAX_WORKERS", "4"))
API_KEY = os.getenv("API_KEY")
REQUIRE_API_KEY = os.getenv("REQUIRE_API_KEY", "true").lower() == "true"

# Security
security = HTTPBearer(auto_error=False)

app = FastAPI(
    title="OCR API",
    description="API for extracting text from PDF files using OCR",
    version="1.0.0"
)

class OCRResponse(BaseModel):
    text: str
    pages: int
    processing_time: float
    confidence: Optional[float] = None

class ErrorResponse(BaseModel):
    error: str
    detail: str

async def verify_api_key(credentials: HTTPAuthorizationCredentials = Depends(security)):
    """
    Verify API key authentication
    """
    if not REQUIRE_API_KEY:
        return True

    if not API_KEY:
        logger.warning("API_KEY not set in environment variables")
        return True

    if not credentials:
        raise HTTPException(
            status_code=401,
            detail="API key required",
            headers={"WWW-Authenticate": "Bearer"},
        )

    if credentials.credentials != API_KEY:
        raise HTTPException(
            status_code=401,
            detail="Invalid API key",
            headers={"WWW-Authenticate": "Bearer"},
        )

    return True

@app.get("/")
async def root():
    return {"message": "OCR API is running. Use /docs for API documentation."}

@app.get("/health")
async def health_check():
    return {"status": "healthy", "service": "OCR API"}

def process_single_page(args):
    """
    Process a single page with OCR (runs in thread pool)
    """
    page_num, image = args
    try:
        logger.info(f"Processing page {page_num + 1}")

        # Perform OCR on the image
        ocr_data = pytesseract.image_to_data(image, output_type=pytesseract.Output.DICT)

        # Extract text and confidence scores
        page_text = []
        total_confidence = 0
        confidence_count = 0

        for j in range(len(ocr_data['text'])):
            if int(ocr_data['conf'][j]) > 0:  # Filter out low confidence results
                page_text.append(ocr_data['text'][j])
                total_confidence += float(ocr_data['conf'][j])
                confidence_count += 1

        page_text = ' '.join(page_text).strip()
        avg_confidence = total_confidence / confidence_count if confidence_count > 0 else 0

        return {
            'page_num': page_num,
            'text': page_text,
            'confidence': avg_confidence,
            'success': True
        }
    except Exception as e:
        logger.error(f"Error processing page {page_num + 1}: {str(e)}")
        return {
            'page_num': page_num,
            'text': '',
            'confidence': 0,
            'success': False,
            'error': str(e)
        }

async def process_pdf_with_ocr(pdf_path: str, max_workers: int = None) -> OCRResponse:
    """
    Process a PDF file and extract text using OCR with concurrent page processing
    """
    start_time = time.time()

    try:
        # Convert PDF to images
        logger.info(f"Converting PDF to images: {pdf_path}")
        images = convert_from_path(pdf_path, dpi=300)

        if not images:
            raise HTTPException(status_code=400, detail="Could not extract images from PDF")

        # Determine optimal number of workers
        if max_workers is None:
            max_workers = min(len(images), DEFAULT_MAX_WORKERS)

        logger.info(f"Processing {len(images)} pages with {max_workers} workers")

        # Process pages concurrently using ThreadPoolExecutor
        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            # Create tasks for each page
            page_tasks = [(i, image) for i, image in enumerate(images)]

            # Submit all tasks and wait for completion
            results = list(executor.map(process_single_page, page_tasks))

        # Sort results by page number and combine
        results.sort(key=lambda x: x['page_num'])

        all_text = []
        total_confidence = 0
        confidence_count = 0
        successful_pages = 0

        for result in results:
            if result['success']:
                successful_pages += 1
                if result['text']:
                    all_text.append(f"--- Page {result['page_num'] + 1} ---\n{result['text']}")
                    total_confidence += result['confidence']
                    confidence_count += 1
            else:
                logger.warning(f"Page {result['page_num'] + 1} failed: {result.get('error', 'Unknown error')}")

        # Calculate average confidence
        avg_confidence = total_confidence / confidence_count if confidence_count > 0 else None

        processing_time = time.time() - start_time

        logger.info(f"Completed processing {successful_pages}/{len(images)} pages in {processing_time:.2f}s")

        return OCRResponse(
            text='\n\n'.join(all_text),
            pages=len(images),
            processing_time=round(processing_time, 2),
            confidence=round(avg_confidence, 2) if avg_confidence else None
        )

    except Exception as e:
        logger.error(f"Error processing PDF: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Error processing PDF: {str(e)}")

@app.post("/ocr/pdf", response_model=OCRResponse)
async def extract_text_from_pdf(
    file: UploadFile = File(...),
    max_workers: int = None,
    background_tasks: BackgroundTasks = None,
    _: bool = Depends(verify_api_key)
):
    """
    Extract text from a PDF file using OCR with concurrent processing

    - **file**: PDF file to process
    - **max_workers**: Maximum number of concurrent workers (default: from DEFAULT_MAX_WORKERS env var)
    - Returns extracted text with metadata
    """

    # Validate file type
    if not file.filename.lower().endswith('.pdf'):
        raise HTTPException(
            status_code=400,
            detail="File must be a PDF"
        )

    # Check file size (limit to 50MB)
    if file.size and file.size > 50 * 1024 * 1024:
        raise HTTPException(
            status_code=400,
            detail="File size too large. Maximum size is 50MB"
        )

    try:
        # Create temporary file
        with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf') as temp_file:
            # Write uploaded file to temporary file
            content = await file.read()
            temp_file.write(content)
            temp_file_path = temp_file.name

        try:
            # Process the PDF
            result = await process_pdf_with_ocr(temp_file_path, max_workers)

            # Clean up temporary file
            background_tasks.add_task(os.unlink, temp_file_path)

            return result

        except Exception as e:
            # Clean up temporary file on error
            if os.path.exists(temp_file_path):
                os.unlink(temp_file_path)
            raise e

    except Exception as e:
        logger.error(f"Error handling file upload: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Error processing file: {str(e)}")

@app.post("/ocr/image", response_model=OCRResponse)
async def extract_text_from_image(
    file: UploadFile = File(...),
    background_tasks: BackgroundTasks = None,
    _: bool = Depends(verify_api_key)
):
    """
    Extract text from an image file using OCR

    - **file**: Image file to process (PNG, JPG, JPEG, etc.)
    - Returns extracted text with metadata
    """

    # Validate file type
    allowed_extensions = {'.png', '.jpg', '.jpeg', '.bmp', '.tiff', '.tif'}
    file_extension = Path(file.filename).suffix.lower()

    if file_extension not in allowed_extensions:
        raise HTTPException(
            status_code=400,
            detail=f"Unsupported file type. Allowed types: {', '.join(allowed_extensions)}"
        )

    try:
        # Create temporary file
        with tempfile.NamedTemporaryFile(delete=False, suffix=file_extension) as temp_file:
            content = await file.read()
            temp_file.write(content)
            temp_file_path = temp_file.name

        try:
            import time
            start_time = time.time()

            # Open and process image
            image = Image.open(temp_file_path)

            # Perform OCR
            ocr_data = pytesseract.image_to_data(image, output_type=pytesseract.Output.DICT)

            # Extract text and confidence
            text_parts = []
            total_confidence = 0
            confidence_count = 0

            for i in range(len(ocr_data['text'])):
                if int(ocr_data['conf'][i]) > 0:
                    text_parts.append(ocr_data['text'][i])
                    total_confidence += float(ocr_data['conf'][i])
                    confidence_count += 1

            text = ' '.join(text_parts).strip()
            avg_confidence = total_confidence / confidence_count if confidence_count > 0 else None
            processing_time = time.time() - start_time

            # Clean up temporary file
            background_tasks.add_task(os.unlink, temp_file_path)

            return OCRResponse(
                text=text,
                pages=1,
                processing_time=round(processing_time, 2),
                confidence=round(avg_confidence, 2) if avg_confidence else None
            )

        except Exception as e:
            # Clean up temporary file on error
            if os.path.exists(temp_file_path):
                os.unlink(temp_file_path)
            raise e

    except Exception as e:
        logger.error(f"Error processing image: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Error processing image: {str(e)}")

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