# guardianhub_sdk/agents/specialist_base.py

import uuid
from typing import Dict, Any, List, Optional, TypeVar
from fastapi import APIRouter, HTTPException

from guardianhub.clients.llm_service import LLMService
from guardianhub.config.settings import settings
from guardianhub.workflows.constants import get_all_activities
from guardianhub.models.agent_models import AgentSubMission
from guardianhub.models.template.agent_plan import MacroPlan
from guardianhub.services.memory_manager import MemoryManager
from guardianhub.services.episodic_manager import EpisodicManager
from guardianhub.clients.a2a_client import A2AClient
from guardianhub.clients.vector_client import VectorClient
from guardianhub.clients.graph_db_client import GraphDBClient
from guardianhub.clients.registry_client import RegistryClient
from guardianhub.clients.classification_client import ClassificationClient
from guardianhub.clients.metadata_extractor_client import MetadataExtractorClient
from guardianhub.clients.ocr_client import OCRClient
from guardianhub.clients.paperless_client import PaperlessClient
from guardianhub.clients.text_cleaner_client import TextCleanerClient
from guardianhub import get_logger
logger = get_logger(__name__)

# Type variable for client types
T = TypeVar('T')


class SovereignSpecialistBase:
    """
    The Foundation for all Specialist Agents.
    Encapsulates Planning, Execution, and Context Management.
    """

    def __init__(
            self,
            activities_instance: Any,
            # Core clients with defaults
            llm_service: Optional[LLMService] = None,
            vector_client: Optional[VectorClient] = None,
            graph_client: Optional[GraphDBClient] = None,
            registry_client: Optional[RegistryClient] = None,
            temporal_client: Any = None,  # Keep as Any since it's from Temporal SDK
            # Additional specialized clients
            classification_client: Optional[ClassificationClient] = None,
            metadata_extractor: Optional[MetadataExtractorClient] = None,
            ocr_client: Optional[OCRClient] = None,
            paperless_client: Optional[PaperlessClient] = None,
            text_cleaner: Optional[TextCleanerClient] = None,
            # For any other clients
            custom_clients: Optional[Dict[str, Any]] = None
    ):
        # 1. Identity & Config
        self.spec = settings.specialist_settings
        self.name = self.spec.agent_name

        # 2. Initialize core clients with defaults
        self.llm = llm_service or LLMService()
        self.vector_client = vector_client or VectorClient()
        self.graph_client = graph_client or GraphDBClient()
        self.registry_client = registry_client or RegistryClient()
        self.temporal_client = temporal_client

        # 3. Initialize specialized clients
        self.classifier = classification_client or ClassificationClient()
        self.metadata_extractor = metadata_extractor or MetadataExtractorClient()
        self.ocr = ocr_client or OCRClient()
        self.paperless = paperless_client or PaperlessClient()
        self.text_cleaner = text_cleaner or TextCleanerClient()

        # 4. Initialize core services
        try:
            self.memory = MemoryManager(
                vector_client=self.vector_client,
                graph_client=self.graph_client,
                tool_registry=self.registry_client
            )
            self.episodes = EpisodicManager(
                vector_client=self.vector_client,
                graph_client=self.graph_client,
                tool_registry=self.registry_client
            )
            self.a2a = A2AClient(
                sender_name=self.name,
                consul_service=self.registry_client
            )
        except Exception as e:
            logger.warning(f"Could not initialize all core services: {str(e)}")
            if not all([self.vector_client, self.graph_client, self.registry_client]):
                logger.warning("One or more required clients are missing. Some functionality may be limited.")

        # 5. Store custom clients
        self.custom_clients = custom_clients or {}

        # 6. Domain Activity Instance (The "Fuel")
        self.activities_instance = activities_instance

        # 7. API Gateway
        self.router = APIRouter(prefix="/v1/mission")
        self._setup_routes()

    def get_client(self, client_name: str, client_type: T = None) -> Optional[T]:
        """
        Get a client by name with optional type checking.

        Args:
            client_name: Name of the client to retrieve
            client_type: Optional type to validate the client against

        Returns:
            The client instance or None if not found
        """
        client = self.custom_clients.get(client_name)
        if client is not None and client_type is not None and not isinstance(client, client_type):
            raise TypeError(f"Client {client_name} is not of type {client_type.__name__}")
        return client

    def get_activities(self) -> list:
        """
        Discovery Hook: Aggregates domain tools for the Temporal Worker.
        Used by the MuscleFactory during ignition.
        """
        return get_all_activities(self.activities_instance)

    def _setup_routes(self):
        @self.router.post("/propose", summary="Tactical Planning Handshake")
        async def propose(mission: AgentSubMission):
            if not all([self.memory, self.llm]):
                raise HTTPException(
                    status_code=500,
                    detail="Memory manager and LLM service must be configured for proposal generation"
                )

            context = await self.memory.get_reasoning_context(
                query=mission.sub_objective,
                template_id=mission.metadata.get("template_id", "TPL-GENERIC"),
                tools=self.spec.capabilities
            )
            plan = await self.llm.generate_specialist_plan(mission, context)
            return {"plan": plan, "rationale": "Audit based on local domain schema."}

        @self.router.post("/execute", summary="Durable Mission Launch")
        async def execute(plan: MacroPlan):
            if not self.temporal_client:
                raise HTTPException(
                    status_code=500,
                    detail="Temporal client must be configured for mission execution"
                )

            workflow_id = f"mission-{plan.session_id}-{uuid.uuid4().hex[:6]}"
            await self.temporal_client.start_workflow(
                "SpecialistMissionWorkflow",  # Updated workflow name
                args=[plan.model_dump()],
                id=workflow_id,
                task_queue=settings.temporal_settings.task_queue
            )
            return {"status": "running", "workflow_id": workflow_id}