from __future__ import annotations

import os
import asyncio
from omegaconf import OmegaConf
from service_forge.workflow.workflow_factory import create_workflows
from service_forge.api.http_api import start_fastapi_server
from service_forge.api.kafka_api import start_kafka_server
from service_forge.db.database import DatabaseManager
from loguru import logger
from typing import Callable, AsyncIterator, Awaitable, Any, TYPE_CHECKING
from service_forge.api.http_api_doc import generate_service_http_api_doc
from service_forge.sft.config.sf_metadata import SfMetadata

if TYPE_CHECKING:
    from service_forge.workflow.workflow_group import WorkflowGroup

class Service:
    def __init__(
        self,
        metadata: SfMetadata,
        config_path: str,
        workflow_config_paths: list[str],
        _handle_stream_output: Callable[[str, AsyncIterator[str]], Awaitable[None]] = None,
        _handle_query_user: Callable[[str, str], Awaitable[str]] = None,
        enable_http: bool = True,
        http_host: str = "0.0.0.0",
        http_port: int = 8000,
        enable_kafka: bool = True,
        kafka_host: str = "localhost",
        kafka_port: int = 9092,
        service_env: dict[str, Any] = None,
        database_manager: DatabaseManager = None,
    ) -> None:
        self.metadata = metadata
        self.config_path = config_path
        self.workflow_config_paths = workflow_config_paths
        self._handle_stream_output = _handle_stream_output
        self._handle_query_user = _handle_query_user
        self.enable_http = enable_http
        self.http_host = http_host
        self.http_port = http_port
        self.enable_kafka = enable_kafka
        self.kafka_host = kafka_host
        self.kafka_port = kafka_port
        self.service_env = {} if service_env is None else service_env
        self.database_manager = database_manager
        self.workflow_groups: list[WorkflowGroup] = []

    @property
    def name(self) -> str:
        return self.metadata.name
    
    @property
    def version(self) -> str:
        return self.metadata.version
    
    @property
    def description(self) -> str:
        return self.metadata.description

    async def start(self):
        if self.enable_http:
            fastapi_task = asyncio.create_task(start_fastapi_server(self.http_host, self.http_port))
            doc_task = asyncio.create_task(generate_service_http_api_doc(self))
        else:
            fastapi_task = None
            doc_task = None
        if self.enable_kafka:
            kafka_task = asyncio.create_task(start_kafka_server(f"{self.kafka_host}:{self.kafka_port}"))
        else:
            kafka_task = None

        workflow_tasks: list[asyncio.Task] = []

        for workflow_config_path in self.workflow_config_paths:
            workflow_group = create_workflows(
                self.parse_workflow_path(workflow_config_path),
                service_env=self.service_env,
                _handle_stream_output=self._handle_stream_output, 
                _handle_query_user=self._handle_query_user,
                database_manager=self.database_manager,
            )
            self.workflow_groups.append(workflow_group)
            workflow_tasks.append(asyncio.create_task(workflow_group.run()))

        try:
            tasks = []
            if fastapi_task:
                tasks.append(fastapi_task)
            if doc_task:
                tasks.append(doc_task)
            if kafka_task:
                tasks.append(kafka_task)
            tasks.extend(workflow_tasks)
            await asyncio.gather(*tasks)
        except Exception as e:
            logger.error(f"Error in service {self.name}: {e}")
            if fastapi_task:
                fastapi_task.cancel()
            if kafka_task:
                kafka_task.cancel()
            for workflow_task in workflow_tasks:
                workflow_task.cancel()
            raise

    def parse_workflow_path(self, workflow_config_path: str) -> str:
        if os.path.isabs(workflow_config_path):
            return workflow_config_path
        else:
            return os.path.join(os.path.dirname(self.config_path), workflow_config_path)
    
    @staticmethod
    def from_config(metadata, service_env: dict[str, Any] = None) -> Service:
        config = OmegaConf.to_object(OmegaConf.load(metadata.service_config))
        database_manager = DatabaseManager.from_config(config=config)
        return Service(
            metadata=metadata,
            config_path=metadata.service_config,
            workflow_config_paths=config.get('workflows', []),
            _handle_stream_output=None,
            _handle_query_user=None,
            enable_http=config.get('enable_http', True),
            http_host=config.get('http_host', '0.0.0.0'),
            http_port=config.get('http_port', 8000),
            enable_kafka=config.get('enable_kafka', True),
            kafka_host=config.get('kafka_host', 'localhost'),
            kafka_port=config.get('kafka_port', 9092),
            service_env=service_env,
            database_manager=database_manager,
        )

def create_service(config_path: str, name: str, version: str, service_env: dict[str, Any] = None) -> Service:
    return Service.from_config(config_path, name, version, service_env)
