import asyncio, time, json
from concurrent.futures import ThreadPoolExecutor

from unstructured_client import UnstructuredClient
from unstructured_client.models import shared
from unstructured_client.models.errors import SDKError

from animuz_core.utils import init_logger

LOGGER = init_logger(__name__)

class MyUnstructuredClient():
    def __init__(self, host: str, port: int) -> None:
        self.client = UnstructuredClient(server_url=f"http://{host}:{port}", 
                                         api_key_auth=None,
                                        )
        
    def do_partition(self, req: shared.PartitionParameters):
        return self.client.general.partition(req)
    async def a_partition(self, filename: str):
        with open(filename, "rb") as f:
            # Note that this currently only supports a single file
            files=shared.Files(
                content=f.read(),
                file_name=filename,
            )

        req = shared.PartitionParameters(
            files=files,
            chunking_strategy="by_title",
            combine_under_n_chars=1500,
            max_characters=2000,
            new_after_n_chars=1500,
            # Other partition params
            strategy='auto',
            languages=["eng","tha"],
        )

        try:
            start = time.time()
            loop = asyncio.get_running_loop()
            with ThreadPoolExecutor() as pool:
                # Run the partition_document function asynchronously
                resp = await loop.run_in_executor(pool, self.do_partition, req)
                for e in resp.elements:
                    e["source"] = "unstructured"
                # LOGGER.debug(f"{time.strftime('[%Y-%b-%d %H:%M:%S %z]')} - Unstructured latency: {time.time() - start: .4f}s")
                LOGGER.debug(json.dumps({
                    "timestamp": time.strftime('[%Y-%b-%d %H:%M:%S %z]'),
                    "message": "Unstructured parsing latency",
                    "latency": str(round(time.time() - start, 4)),
                    "exception": False
                }))
                return resp.elements
        except SDKError as e:
            LOGGER.critical(json.dumps({
                "timestamp": time.strftime('[%Y-%b-%d %H:%M:%S %z]'),
                "message": "ERROR Unstructured parsing",
                "latency": None,
                "exception": True
            }), exc_info=True)
            # LOGGER.critical(f"{time.strftime('[%Y-%b-%d %H:%M:%S %z]')} - Unstructured error", exc_info=True)
            return None
        
class PaidUnstructuredClient(MyUnstructuredClient):
    def __init__(self, url: str, api_key: int) -> None:
        self.client = UnstructuredClient(server_url=url, 
                                         api_key_auth=api_key,
                                        )
        
if __name__ == "__main__":
    import time
    client = MyUnstructuredClient(host="192.168.1.60", port=8000)

    start = time.time()
    res = asyncio.run(client.a_partition("test_docs/test1.docx"))
    print(res)
    print(time.time() - start)