import asyncio
from openai import AsyncOpenAI
import concurrent.futures
from typing import List, Dict, Any
from openai import OpenAI
import base64
import bilbystats as bs

def encode_image(image_path):
    with open(image_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode('utf-8')


def extract_text(image_path, max_tokens = 1024):

    image_base64 = encode_image(image_path)

    client = OpenAI(
        api_key=bs.read_api_key('DEEPINFRA'),
        base_url="https://api.deepinfra.com/v1/openai"
        )

    response = client.chat.completions.create(
    model="allenai/olmOCR-2-7B-1025",
    messages=[
        {
            "role": "user",
            "content": [
                {"type": "text", "text": "Extract all text from this image."},
                {
                    "type": "image_url",
                    "image_url": {
                        "url": f"data:image/jpeg;base64,{image_base64}"
                    }
                }
            ]
        }
    ],
    max_tokens=max_tokens
    )

    out = response.choices[0].message.content
    return out

def extract_text_batch(image_paths: List[str], max_tokens: int = 1024, max_workers: int = 5) -> List[Dict[str, Any]]:
    """
    Extract text from multiple images concurrently using threading.
    
    Args:
        image_paths: List of paths to images
        max_tokens: Maximum tokens for each extraction
        max_workers: Maximum number of concurrent threads
    
    Returns:
        List of dicts with 'path', 'text' (or 'error' if failed)
    """
    results = []
    
    def process_image(path):
        try:
            text = bs.extract_text(path, max_tokens)
            return {"path": path, "text": text}
        except Exception as e:
            return {"path": path, "error": str(e)}
    
    with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
        future_to_path = {executor.submit(process_image, path): path for path in image_paths}
        
        for future in concurrent.futures.as_completed(future_to_path):
            results.append(future.result())
    
    # Sort results to match input order
    path_order = {path: i for i, path in enumerate(image_paths)}
    results.sort(key=lambda x: path_order[x["path"]])
    
    return results




async def extract_text_async(image_paths, max_tokens=1024, max_concurrent=5):
    if isinstance(image_paths, str):
        image_paths = [image_paths]
        single_input = True
    else:
        single_input = False

    client = AsyncOpenAI(
        api_key=bs.read_api_key('DEEPINFRA'),
        base_url="https://api.deepinfra.com/v1/openai"
    )
    
    semaphore = asyncio.Semaphore(max_concurrent)

    async def process_single_image(image_path):
        async with semaphore:
            image_base64 = bs.encode_image(image_path)
            
            response = await client.chat.completions.create(
                model="allenai/olmOCR-2-7B-1025",
                messages=[
                    {
                        "role": "user",
                        "content": [
                            {"type": "text", "text": "Extract all text from this image."},
                            {
                                "type": "image_url",
                                "image_url": {
                                    "url": f"data:image/jpeg;base64,{image_base64}"
                                }
                            }
                        ]
                    }
                ],
                max_tokens=max_tokens
            )
            return response.choices[0].message.content

    results = await asyncio.gather(*[process_single_image(path) for path in image_paths])

    return results[0] if single_input else results