import os
import asyncio
from typing import Optional, Dict, Any, List
from temporalio.client import Client, Schedule, ScheduleActionStartWorkflow, ScheduleSpec, ScheduleCalendarSpec, ScheduleIntervalSpec, ScheduleRange
from temporalio.common import SearchAttributePair, SearchAttributeKey,TypedSearchAttributes
from temporalio.worker import Worker
from dataclasses import dataclass

from .playground.workflow import playgroundRun
from .observability import log_with_context
from .security import DataConverter

@dataclass
class CloudConnectionOptions:
    engine_id: str
    api_key: str
    address: Optional[str] = "localhost:7233"
    temporal_namespace: Optional[str] = "default"
    data_converter: Optional[DataConverter] = None

@dataclass
class ServiceOptions:
    rate_limit: Optional[int] = 100000
    max_concurrent_workflow_runs: Optional[int] = 3000
    max_concurrent_function_runs: Optional[int] = 1000

class Restack:
    api_key: Optional[str] = None
    client: Optional[Client] = None
    options: Optional[CloudConnectionOptions] = None
    functions: Optional[List[Any]] = None

    def __init__(self, options: Optional[CloudConnectionOptions] = None):
        super().__init__()
        self.client = None
        self.options = options
        log_with_context("DEBUG", "Restack instance initialized", options=str(options))

    def get_connection_options(self) -> Dict[str, Any]:
        target_host = self.options.address if self.options and self.options.address is not None else 'localhost:7233'
        engine_id = self.options.engine_id if self.options and self.options.engine_id is not None else 'local'
        options = {
            "target_host": target_host,
            "metadata": {
                "restack-engineId": engine_id,
            }
        }
        if self.options and self.options.api_key is not None:
            options["tls"] = True
            options["api_key"] = self.options.api_key
        return options
    
    async def connect(self, connection_options: Optional[CloudConnectionOptions] = None):
        if self.client:
            log_with_context("INFO", "Reusing existing client connection")
            return
        try:
            self.client = await self.create_client(connection_options)
            log_with_context("INFO", "Connected to Restack Engine")
        except Exception as e:
            log_with_context("ERROR", "Failed to connect to Restack Engine", error=str(e))
            raise e

    async def create_client(self, connection_options: Optional[CloudConnectionOptions] = None) -> Client:
        connect_options = connection_options or self.get_connection_options()

        target_host = connect_options['target_host']
        namespace = self.options.temporal_namespace if self.options and self.options.temporal_namespace is not None else 'default'
        api_key = self.options.api_key if self.options and self.options.api_key is not None else None
        tls = True if self.options and self.options.api_key is not None else False
        metadata = connect_options['metadata']
        data_converter = self.options.data_converter if self.options and self.options.data_converter is not None else None
        try:
            client= await Client.connect(
              target_host,
              namespace=namespace,
              api_key=api_key,
              tls=tls,
              rpc_metadata=metadata,
              data_converter=data_converter
            )
        except Exception as e:
            log_with_context("ERROR", "Failed to connect to Restack Engine", error=str(e))
            raise e
        return client

    async def create_service(self, workflows: Optional[List[Any]] = None, functions: Optional[List[Any]] = None, task_queue: Optional[str] = None,options: Optional[ServiceOptions] = None) -> Worker:
        try:
            log_with_context("INFO", "Starting service...")
            client = await self.create_client()
            engine_id = self.get_connection_options()['metadata']['restack-engineId']

            service = Worker(
                identity=f"{engine_id}-{os.getpid()}-{task_queue or 'restack'}",
                client=client,
                task_queue=task_queue or "restack",
                workflows=workflows or [],
                activities=functions or [],
                max_activities_per_second=options.rate_limit,
                max_concurrent_activities=options.max_concurrent_function_runs,
                max_concurrent_workflow_tasks=options.max_concurrent_workflow_runs,
            )
            log_with_context("INFO", "Service created successfully")
            return service
        except Exception as e:
            if e.__cause__ and "http.client.incompleteread.__mro_entries__".lower() in str(e.__cause__).lower():
                log_with_context("ERROR", "Failed to start service: Functions in workflow steps need to be imported with import_functions(). See docs at https://docs.restack.io/libraries/python/reference/workflows")
            else:
                log_with_context("ERROR", "Failed to start service", error=str(e))
            raise e

    async def run_service(self, service: Worker):
        try:
            if service.task_queue == "restack":
                log_with_context("INFO", "Creating playground service...")

                playground_service = Worker(
                    identity=f"{os.getpid()}-playground",
                    client=service.client,
                    task_queue="playground",
                    workflows=[playgroundRun],
                )

                log_with_context("INFO", "Services ready to receive workflows and events")
                await asyncio.gather(service.run(), playground_service.run())
            else:
                log_with_context("INFO", "Service ready to receive workflows and events")
                await service.run()
        except Exception as e:
            log_with_context("ERROR", "Failed to run service", error=str(e))
            raise e

    async def start_service(
        self,
        workflows: Optional[List[Any]] = None,
        functions: Optional[List[Any]] = None,
        task_queue: Optional[str] = None,
        options: Optional[ServiceOptions] = None
    ):
        """
        Start a service with the specified configurations.

        Parameters:
        - workflows (Optional[List[Any]]): A list of workflows to be used.
        - functions (Optional[List[Any]]): A list of functions to be used. 
        - task_queue (Optional[str]): The task queue name. 
        - options (Optional[ServiceOptions]): Service options for rate limiting and concurrency.
        """
        service = await self.create_service(
            task_queue=task_queue,
            workflows=workflows,
            functions=functions or [],
            options=options or ServiceOptions()
        )
        await self.run_service(service)

    async def schedule_workflow(self, workflow_name: str, workflow_id: str, input: Optional[Dict[str, Any]] = None, schedule: Optional[ScheduleSpec] = None, task_queue: Optional[str] = 'restack') -> str:
        await self.connect()
        if self.client:
            try:
                connection_options = self.get_connection_options()
                engine_id = connection_options['metadata']['restack-engineId']

                search_attribute_pairs = [
                    SearchAttributePair(key=SearchAttributeKey.for_text('engineId'), value=engine_id)
                ]
                search_attributes = TypedSearchAttributes(search_attributes=search_attribute_pairs)

                if not schedule:
                    handle = await self.client.start_workflow(
                        workflow_name,
                        args=[input] if input else [],
                        id=f"{engine_id}-{workflow_id}",
                        memo={'engineId': engine_id},
                        search_attributes=search_attributes,
                        task_queue=task_queue,
                    )
                    return handle.first_execution_run_id
                else:
                    schedule_action = ScheduleActionStartWorkflow(
                        workflow=workflow_name,
                        args=[input] if input else [],
                        id=f"{engine_id}-{workflow_id}",
                        task_queue=task_queue,
                        memo={'engineId': engine_id},
                        typed_search_attributes=search_attributes
                    )
                    schedule_obj = Schedule(
                        action=schedule_action,
                        spec=schedule,
                    )
                    scheduled = await self.client.create_schedule(
                        id=f"{engine_id}-{workflow_id}",
                        schedule=schedule_obj,
                        memo={'engineId': engine_id},
                        search_attributes=search_attributes
                    )
                    return scheduled.id
                    
            except Exception as e:
                log_with_context("ERROR", "Failed to start or schedule workflow", error=str(e))
                raise e
        else:
            raise Exception("Workflow result not retrieved due to failed connection.")

    async def get_workflow_handle(self, workflow_id: str, run_id: str):
        await self.connect()
        if self.client:
            try:
                connection_options = self.get_connection_options()
                engine_id = connection_options['metadata']['restack-engineId']
                return self.client.get_workflow_handle(f"{engine_id}-{workflow_id}", run_id=run_id)
            except Exception as e:
                log_with_context("ERROR", "Failed to get workflow result", error=str(e))
                raise e
        else:
            raise Exception("Workflow result not retrieved due to failed connection.")

    async def get_workflow_result(self, workflow_id: str, run_id: str) -> Any:
        handle = await self.get_workflow_handle(workflow_id, run_id)
        try:
            return await handle.result()
        except Exception as e:
            log_with_context("ERROR", "Failed to get workflow result", error=str(e))
            raise e

    async def get_workflow_memory(self, workflow_id: str, run_id: str, memory_name: str) -> Any:
        handle = await self.get_workflow_handle(workflow_id, run_id)
        try:
            return await handle.query(memory_name)
        except Exception as e:
            log_with_context("ERROR", "Failed to get workflow memory", error=str(e))
            raise e
    
    async def send_workflow_event(self, event_name: str, workflow_id: str, run_id: Optional[str] = None, event_input: Optional[Dict[str, Any]] = None):
        handle = await self.get_workflow_handle(workflow_id, run_id)
        try:
            return await handle.execute_update(event_name, event_input)
        except Exception as e:
            log_with_context("ERROR", "Failed to send workflow event", error=str(e))
            raise e
