from __future__ import annotations
import boto3
from botocore.config import Config
import json
import time
import os
import io
from typing import Optional, List, Dict
from enum import Enum
import pymupdf
import concurrent.futures
from dataclasses import dataclass
from my_aws_helpers.logging import select_powertools_logger


logger = select_powertools_logger('bedrock-boy')

logger.info("Got logger")

class PromptType(str, Enum):
    json = "json_system_prompt.txt"
    markdown = "markdown_system_prompt.txt"


@dataclass
class TokenUsage:
    input_tokens: int
    output_tokens: int
    total_tokens: int

    @classmethod
    def from_dict(cls, data: Dict[str, int]) -> TokenUsage:
        return cls(
            input_tokens=data.get('inputTokens', 0),
            output_tokens=data.get('outputTokens', 0),
            total_tokens=data.get('totalTokens', 0),
        )


@dataclass
class OCRResult:
    content: List[Dict[str, str]]
    token_usage: TokenUsage
    page_number: int

    @classmethod
    def from_dict(cls, data: Dict) -> OCRResult:
        return cls(
            content = data.get("content", []),
            token_usage = TokenUsage.from_dict(data.get("token_usage", {})),
            page_number = data.get("page_number", 0)
        )


class Bedrock:
    def __init__(
        self,
        model_id: str = "apac.anthropic.claude-3-5-sonnet-20241022-v2:0",
    ):

        self.session = Bedrock._set_session_params()
        region_name = "ap-southeast-2"
        if self.session is None:
            self.session = boto3.Session(region_name = region_name)

        
        custom_config = Config(
            retries={
                'max_attempts': 2,  # Total attempts = 1 initial + 1 retry
                'mode': 'standard',  # or 'adaptive'
            }
        )      
        self.client = self.session.client("bedrock-runtime", region_name=region_name, config = custom_config)
        self.model_id = model_id

    @staticmethod
    def _set_session_params():
        try:
            aws_access_key_id = os.environ["AWS_ACCESS_KEY_ID"]
            aws_secret_access_key = os.environ["AWS_SECRET_ACCESS_KEY"]
            aws_session_token = os.environ["AWS_SESSION_TOKEN"]
            region_name = os.environ["AWS_DEFAULT_REGION"]
            return boto3.Session(
                aws_access_key_id=aws_access_key_id,
                aws_secret_access_key=aws_secret_access_key,
                aws_session_token=aws_session_token,
                region_name=region_name
            )
        except Exception as e:
            logger.exception(e)
            return None

    def _get_prompt(self, prompt_type: str) -> Optional[str]:
        if prompt_type not in list(PromptType):
            raise Exception(f"Error: Invalid prompt type")
        
        path = os.path.join(os.path.dirname(__file__), "prompts", prompt_type)
        try:
            with open(path, "r") as f:
                prompt = f.read()
                return prompt
        except Exception as e:
            logger.exception(f"Failed to get {prompt_type} prompt due to {e}")
        return None

    def _ocr(
        self, 
        prompt: str, 
        image_bytes: bytes,
        page_number: Optional[int] = 0
    ) -> Optional[OCRResult]:
        system_prompt = [{"text": prompt}]
        message = [
            {
                "role": "user",
                "content": [
                    {
                        "image": {
                            "format": "png",
                            "source": {
                                "bytes": image_bytes,
                            }
                        }
                    }
                ]
            }
        ]
        retries = 3
        for i in range(retries):
            logger.info(f"Attempt number {i} for {self.model_id} converse")
            try:
                response = self.client.converse(modelId = self.model_id, messages = message, system = system_prompt)
                if response['ResponseMetadata']['HTTPStatusCode'] == 200:
                    break
            except Exception as e:
                logger.exception(f"Error during conversation due to {e}")
                if i >= len(retries) - 1: raise Exception(e)
                continue
        
        result = {}
        result["content"] = json.loads(response["output"]["message"]["content"][0]["text"])
        result["token_usage"] = response["usage"]
        result["page_number"] = page_number
        return OCRResult.from_dict(data = result)
    
    def _parallel_ocr(
            self, 
            image_bytes_list: List[bytes], 
            prompt_type: str,
            max_workers: int = 10,
    ):
        results = list()
        prompt = self._get_prompt(prompt_type=prompt_type)
        with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
            logger.info("Some log")
            execution_futures = {
                executor.submit(
                    self._ocr, 
                    prompt = prompt,
                    image_bytes = img,
                ): img for img in image_bytes_list
            }
            for future in concurrent.futures.as_completed(execution_futures):
                result = future.result()
                if result:
                    results.append(result)
        return results
    
    def get_ocr_result(
        self,
        pdf_bytes: io.BytesIO, 
        prompt_type: str,
        zoom: int = 7,
    ) -> List[OCRResult]:
        logger.info("Getting OCR Results")
        try:
            document = pymupdf.open(stream=pdf_bytes, filetype="pdf")
        except Exception as e:
            logger.exception(f"Failed to open pdf due to {e}")
            return []

        pages: List[pymupdf.Page] = [p for p in document]

        image_bytes_list: List[bytes] = list()
        for i, p in enumerate(pages):
            try:
                image_bytes: bytes = p.get_pixmap(matrix = pymupdf.Matrix(zoom, zoom)).tobytes("png")
                image_bytes_list.append(image_bytes)
            except Exception as e:
                logger.error(f"Could not get pix map for page {i}")
                continue
        prompt = self._get_prompt(prompt_type=prompt_type)
        logger.info("Got Prompt")
        results = list()
        for i, image_bytes in enumerate(image_bytes_list):
            logger.info(f"Starting OCR for page: {i}")
            results.append(self._ocr(image_bytes=image_bytes, prompt=prompt))
        return results
            