# guardianhub_sdk/agents/specialist_base.py
import inspect
import uuid
from typing import Dict, Any, List, Optional, TypeVar
from fastapi import APIRouter, HTTPException

from guardianhub.clients.llm_client import LLMClient

from guardianhub.clients.consul_client import ConsulClient
from guardianhub.clients.tool_registry_client import ToolRegistryClient
from guardianhub.config.settings import settings
from guardianhub.agents.workflows.constants import get_all_activities
from guardianhub.agents.workflows.agent_contract import ActivityRoles
from guardianhub.models.agent_models import AgentSubMission
from guardianhub.models.template.agent_plan import MacroPlan
from guardianhub.agents.services.memory_manager import MemoryManager
from guardianhub.agents.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.tool_registry_client import ToolRegistryClient
from guardianhub.clients.classification_client import ClassificationClient
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
from temporalio import activity
from httpx import AsyncClient
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_client: Optional[LLMClient] = None,
            vector_client: Optional[VectorClient] = None,
            graph_client: Optional[GraphDBClient] = None,
            tool_registry_client: Optional[ToolRegistryClient] = None,
            consul_client: Optional[ConsulClient] = None,
            temporal_client: Any = None,  # Keep as Any since it's from Temporal SDK
            # Additional specialized clients
            classification_client: Optional[ClassificationClient] = 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_client or LLMClient()
        self.vector_client = vector_client or VectorClient()
        self.graph_client = graph_client or GraphDBClient()
        self.tool_registry_client = tool_registry_client or ToolRegistryClient()
        self.consul_client = consul_client or ConsulClient()
        self.temporal_client = temporal_client

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

        # 4. Initialize core services
        try:
            self.memory = MemoryManager(
                vector_client=self.vector_client,
                graph_client=self.graph_client,
                tool_registry=self.tool_registry_client
            )
            self.episodes = EpisodicManager(
                vector_client=self.vector_client,
                graph_client=self.graph_client,
                tool_registry=self.tool_registry_client
            )
            self.a2a = A2AClient(
                sender_name=self.name,
                consul_service=self.consul_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.tool_registry_client,self.consul_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_activities(self) -> list:
        from temporalio import activity
        registry = {}

        # 1. HARVEST EXPLICIT OVERRIDES (Muscle)
        if hasattr(self.activities_instance, "get_muscle_registry"):
            muscle_map = self.activities_instance.get_muscle_registry()
            for act_name, method in muscle_map.items():
                # Extract the mutable function from the read-only method proxy
                target = method.__func__ if hasattr(method, "__func__") else method

                if not hasattr(target, "__temporal_activity_definition"):
                    activity.defn(target, name=act_name)

                defn = getattr(target, "__temporal_activity_definition")
                target._defn = defn
                registry[defn.name] = target
                logger.info(f"✅ [FUSED] Muscle capability: {act_name}")

        # 2. HARVEST KERNEL FALLBACKS (SDK Base)
        import inspect
        # We use inspect to find all async methods on 'self' (the SDK base)
        for _, method in inspect.getmembers(self, predicate=inspect.iscoroutinefunction):
            # Extract the mutable function from the SDK's own bound methods
            target = method.__func__ if hasattr(method, "__func__") else method

            defn = getattr(target, "__temporal_activity_definition", None)
            if defn:
                if defn.name not in registry:
                    # 🟢 THE SYMMETRICAL FIX: Pin to 'target' (function), not 'method'
                    target._defn = defn
                    registry[defn.name] = target
                    logger.info(f"🛡️ [KERNEL] Fallback active: {defn.name}")

        return list(registry.values())

    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 _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"
                )

            # 1. Gather Context (Hindsight + Ground Truth)
            # This pulls past AARs (successes/failures) and Neo4j schema
            context = await self.memory.get_reasoning_context(
                query=mission.sub_objective,
                template_id=mission.metadata.get("template_id", "TPL-GENERIC"),
                tools=self.spec.capabilities  # These are the FUSED activities we just verified!
            )

            # 2. Invoke Structured Planner
            # We use the specific 'MacroPlan' model to ensure the output
            # fits our Temporal Workflow's 'steps' loop.
            proposal = await self.llm.invoke_structured_model(
                user_input=f"Mission: {mission.sub_objective}\nConstraints: {mission.metadata}",
                system_prompt_template=f"""
                        You are the Sovereign Specialist: {self.name}.
                        Your task is to generate a durable execution plan.

                        AVAILABLE TOOLS: {self.spec.capabilities}
                        ENVIRONMENT CONTEXT: {context}

                        RULES:
                        1. Every step must map to an available tool.
                        2. If the risk is high, include a 'verify_objective_attainment' step.
                        3. The plan must be formatted as a MacroPlan.
                        """,
                response_model_name="MacroPlan"
            )

            return {
                "plan": proposal.steps,
                "rationale": proposal.rationale,
                "confidence_score": proposal.confidence_score,
                "session_id": mission.session_id
            }

        @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}

    @activity.defn(name=ActivityRoles.TACTICAL)
    async def analyze_tactical_context(self, bundle: dict) -> dict:
        """
        SDK Default: Real-time Risk Assessment.
        Acts as a circuit breaker between RECON and INTERVENTION.
        """
        TACTICAL_RISK_PROMPT = """
        Role: You are the Sovereign Tactical Safety Officer (TSO).
        Objective: Perform a high-fidelity risk assessment of a proposed action plan based on real-time reconnaissance data.

        INPUT DATA:
        1. PROPOSED STEPS: {steps}
        2. RECONNAISSANCE FACTS: {facts}

        YOUR TASK:
        Analyze the delta between the 'Plan' and the 'Reality'. 
        Specifically look for:
        - RESOURCE COLLISION: Does the plan attempt to modify a CI that is currently 'unstable' or 'unverified'?
        - BLAST RADIUS: Will the steps impact downstream dependencies not mentioned in the plan?
        - POLICY VIOLATION: Are any steps violating 'safe-window' or 'least-privilege' protocols?

        OUTPUT FORMAT:
        Return a structured JSON with: risk_score (0.0-1.0), reasoning (punchy justification), and recommendation (PROCEED/HALT).
        """

        
        steps = bundle.get("steps", [])
        intelligence = bundle.get("intelligence", [])

        # 1. Forensic Check
        if not steps or not intelligence:
            return {
                "status": "REJECTED",
                "risk_score": 1.0,
                "justification": "Tactical bundle incomplete. Cannot attest without Recon data."
            }

        # 2. Invoke Structured Attestation
        # We pass the variable prompt here
        analysis = await self.llm.invoke_structured_model(
            user_input=f"STEPS: {steps}\nFACTS: {intelligence}",
            system_prompt_template=TACTICAL_RISK_PROMPT,
            response_model_name="RiskAnalysisReport"
        )

        # 3. Log the Decision
        status = "APPROVED" if analysis.risk_score < 0.8 else "REJECTED"
        logger.info(f"🏛️ [TACTICAL_AUDIT] Result: {status} | Risk: {analysis.risk_score}")

        return {
            "status": status,
            "risk_score": analysis.risk_score,
            "justification": analysis.reasoning
        }

        # 3. Log the Attestation
        status = "APPROVED" if analysis.risk_score < 0.8 else "REJECTED"
        logger.info(f"🏛️ [TACTICAL_AUDIT] Result: {status} | Score: {analysis.risk_score}")

        if analysis.risk_score > 0.8:
            logger.warning(f"🚨 ALERT: {analysis.reasoning}")

        return {
            "status": status,
            "risk_score": analysis.risk_score,
            "justification": analysis.reasoning
        }


    @activity.defn(name=ActivityRoles.HISTORY)
    async def retrieve_intelligence_history(self, plan: Dict[str, Any]) -> List[Dict[str, Any]]:
        """Uses your VectorClient to pull past AARs (Hindsight)."""
        return await self.vector_client.get_recent_episodes(
            query=plan.get("sub_objective")
        )

    @activity.defn(name=ActivityRoles.AAR)
    async def commit_mission_after_action_report(self, results: List[Dict[str, Any]]) -> bool:
        """
        SDK Default: Converts mission results into a permanent 'Long-Term Memory'
        for future reconnaissance.
        """
        # Create a structured episode
        episode = {
            "agent": self.name,
            "timestamp": settings.current_time,
            "outcomes": results,
            "success_rate": sum(1 for r in results if r.get("verified")) / len(results) if results else 0
        }

        # Save to Vector DB so NEXT time 'conduct_reconnaissance' finds this
        await self.vector_client.upsert_document_from_text(
            document_content=str(episode),
            doc_id=f"aar-{uuid.uuid4().hex[:6]}",
            collection="episodes",
            metadata={"type": "mission_aar", "agent": self.name}
        )
        return True

    @activity.defn(name=ActivityRoles.COMPLETION)
    async def transmit_mission_completion(self, debrief: Dict[str, Any]) -> bool:
        """
        SDK Default: Closes the mission loop with the central orchestrator.
        """
        payload = {
            "agent": self.name,
            "mission_id": debrief.get("mission_id"),
            "final_report": debrief,
            "status": "COMPLETED"
        }

        async with AsyncClient() as client:
            # We use the Consul-discovered URL for Sutram
            sutram_url = await self.consul_client.get_service_url("sutram-orchestrator")
            response = await client.post(f"{sutram_url}/v1/callback", json=payload)
            return response.is_success