# guardianhub_sdk/agents/specialist_base.py
import datetime
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():
                # Soul Extraction (for metadata pinning)
                target = method.__func__ if hasattr(method, "__func__") else method

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

                # Signature Promotion
                defn = getattr(target, "__temporal_activity_definition")
                target._defn = defn

                # 🟢 RETURN THE BODY: Keep it bound to the specialist instance
                registry[defn.name] = method
                logger.info(f"✅ [FUSED] Muscle capability: {act_name}")

        # 2. HARVEST KERNEL FALLBACKS (SDK Base)
        import inspect
        for _, method in inspect.getmembers(self, predicate=inspect.iscoroutinefunction):
            target = method.__func__ if hasattr(method, "__func__") else method
            defn = getattr(target, "__temporal_activity_definition", None)

            if defn and defn.name not in registry:
                # Signature Promotion
                target._defn = defn
                # 🟢 RETURN THE BODY: Keep it bound to 'self' (the SDK base)
                registry[defn.name] = method
                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="Sovereign Planning Handshake")
        async def propose(mission: AgentSubMission):
            # The SDK no longer fetches context here.
            # It trusts the Specialist's 'PROPOSAL' activity to be self-sufficient.

            activities = self.get_activities()
            try:
                plan_activity = next(a for a in activities if a._defn.name == ActivityRoles.PROPOSAL)

                # This call now handles Context + Logic + Planning
                proposal: MacroPlan = await plan_activity(mission)

                return {
                    "plan": proposal.plan,
                    "rationale": proposal.reflection,
                    "session_id": mission.session_id,
                    "metadata": proposal.metadata
                }
            except StopIteration:
                raise HTTPException(status_code=501, detail="Specialist has no PROPOSAL capability fused.")

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

    # ============================================================================
    # SOVEREIGN SPECIALIST ACTIVITY METHODS - Organized by ActivityRoles Sequence
    # ============================================================================
    @activity.defn(name=ActivityRoles.RECON)
    async def conduct_reconnaissance(self, mission: AgentSubMission) -> List[Dict[str, Any]]:
        """
        SDK Default: Uses the objective from the mission brief to pull facts.
        """
        logger.info(f"📡 [RECON] Starting reconnaissance for mission: {mission.sub_objective}")
        logger.info(f"📡 [RECON] Agent: {self.name} | Template: {mission.metadata.get('template_id', 'TPL-GENERIC')}")
        
        try:
            facts = await self.graph_client.get_infrastructure_facts(
                query=mission.sub_objective,
                metadata=mission.metadata,
                template_id=mission.metadata.get("template_id", "TPL-GENERIC")
            )
            logger.info(f"📡 [RECON] Successfully retrieved {len(facts)} infrastructure facts")
            return facts
        except Exception as e:
            logger.error(f"📡 [RECON] Failed to gather intelligence: {str(e)}")
            raise


    @activity.defn(name=ActivityRoles.HISTORY)
    async def retrieve_intelligence_history(self, mission: AgentSubMission) -> List[Dict[str, Any]]:
        """
        SDK Default: Uses mission objective for Vector Hindsight.
        """
        logger.info(f"📚 [HISTORY] Retrieving intelligence history for: {mission.sub_objective}")
        logger.info(f"📚 [HISTORY] Agent: {self.name} consulting episodic memory")
        
        try:
            episodes = await self.vector_client.get_recent_episodes(
                query=mission.sub_objective
            )
            logger.info(f"📚 [HISTORY] Found {len(episodes)} relevant historical episodes")
            return episodes
        except Exception as e:
            logger.error(f"📚 [HISTORY] Failed to retrieve intelligence history: {str(e)}")
            raise

    @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.
        """
        logger.info(f"🏛️ [TACTICAL] Starting tactical analysis for mission bundle")
        logger.info(f"🏛️ [TACTICAL] Agent: {self.name} performing risk assessment")
        
        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", [])

        logger.info(f"🏛️ [TACTICAL] Analyzing {len(steps)} proposed steps against {len(intelligence)} intelligence facts")

        # 1. Forensic Check
        if not steps or not intelligence:
            logger.warning(f"🏛️ [TACTICAL] Bundle incomplete - Steps: {len(steps)}, Intelligence: {len(intelligence)}")
            return {
                "status": "REJECTED",
                "risk_score": 1.0,
                "justification": "Tactical bundle incomplete. Cannot attest without Recon data."
            }

        # 2. Invoke Structured Attestation
        try:
            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:.2f}")
            logger.info(f"🏛️ [TACTICAL_AUDIT] Justification: {analysis.reasoning}")

            if analysis.risk_score > 0.8:
                logger.warning(f"🚨 [TACTICAL] HIGH RISK ALERT: {analysis.reasoning}")

            return {
                "status": status,
                "risk_score": analysis.risk_score,
                "justification": analysis.reasoning
            }
        except Exception as e:
            logger.error(f"🏛️ [TACTICAL] Risk assessment failed: {str(e)}")
            return {
                "status": "REJECTED",
                "risk_score": 1.0,
                "justification": f"Risk assessment system failure: {str(e)}"
            }

    @activity.defn(name=ActivityRoles.PROPOSAL)
    async def formulate_mission_proposal(self, mission: AgentSubMission) -> MacroPlan:
        """
        SDK Default: Generates the plan.
        Note: This is usually called via /propose, but can be a durable activity.
        """
        logger.info(f"📋 [PROPOSAL] Formulating mission proposal for: {mission.sub_objective}")
        logger.info(f"📋 [PROPOSAL] Agent: {self.name} creating structured execution plan")
        
        try:
            plan = await self.llm.generate_specialist_plan(mission.sub_objective, {})
            logger.info(f"📋 [PROPOSAL] Generated plan with {len(plan.steps)} steps")
            logger.info(f"📋 [PROPOSAL] Plan confidence: {plan.confidence_score:.2f}")
            return plan
        except Exception as e:
            logger.error(f"📋 [PROPOSAL] Failed to generate mission proposal: {str(e)}")
            raise

    @activity.defn(name=ActivityRoles.INTERVENTION)
    async def execute_direct_intervention(self, step: Dict[str, Any]) -> Dict[str, Any]:
        """
        SDK Default: A 'dry-run' placeholder.
        Specialists MUST override this to actually touch infrastructure.
        """
        action = step.get('action', 'unknown')
        logger.info(f"⚡ [INTERVENTION] Executing direct intervention: {action}")
        logger.info(f"⚡ [INTERVENTION] Agent: {self.name} performing action")
        
        if hasattr(self.activities_instance, 'execute_direct_intervention'):
            logger.info(f"⚡ [INTERVENTION] Delegating to specialist implementation")
            return await self.activities_instance.execute_direct_intervention(step)
        
        logger.warning(f"⚠️ [INTERVENTION] SDK_KERNEL: Dry-run intervention for step: {action}")
        logger.warning(f"⚠️ [INTERVENTION] Specialist should override this method for actual execution")
        return {"status": "dry_run_success", "action": action}

    @activity.defn(name=ActivityRoles.ATTAINMENT)
    async def verify_objective_attainment(self, step: Dict[str, Any]) -> bool:
        """
        SDK Default: Optimistic verification.
        Specialists should override this to perform a 'Post-Action Recon'.
        """
        action = step.get('action', 'unknown')
        logger.info(f"✅ [ATTAINMENT] Verifying objective attainment for: {action}")
        logger.info(f"✅ [ATTAINMENT] Agent: {self.name} performing post-action verification")
        
        if hasattr(self.activities_instance, 'verify_objective_attainment'):
            logger.info(f"✅ [ATTAINMENT] Delegating to specialist implementation")
            return await self.activities_instance.verify_objective_attainment(step)
        
        logger.info(f"✅ [ATTAINMENT] SDK_KERNEL: Optimistic verification for {action}")
        logger.info(f"✅ [ATTAINMENT] Specialist should override this method for actual verification")
        return True

    @activity.defn(name=ActivityRoles.DEBRIEF)
    async def summarize_intelligence_debrief(self, bundle: Dict[str, Any]) -> Dict[str, Any]:
        """
        SDK Default: Narrative synthesis of the mission outcomes.
        """
        mission_id = bundle.get("mission_id", "unknown")
        logger.info(f"📝 [DEBRIEF] Summarizing intelligence debrief for mission: {mission_id}")
        logger.info(f"📝 [DEBRIEF] Agent: {self.name} creating mission narrative")
        
        try:
            outcomes = bundle.get("outcomes", [])
            logger.info(f"📝 [DEBRIEF] Processing {len(outcomes)} mission outcomes")
            
            summary = await self.llm.summarize_results(outcomes)
            
            result = {
                "mission_id": mission_id,
                "summary": summary,
                "status": "SUCCESS"
            }
            
            logger.info(f"📝 [DEBRIEF] Successfully created debrief summary")
            return result
        except Exception as e:
            logger.error(f"📝 [DEBRIEF] Failed to create intelligence debrief: {str(e)}")
            raise

    @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.
        """
        logger.info(f"🗃️ [AAR] Committing After Action Report for mission results")
        logger.info(f"🗃️ [AAR] Agent: {self.name} storing {len(results)} mission outcomes")
        
        try:
            # Create a structured episode
            episode = {
                "agent": self.name,
                "timestamp": datetime.datetime.utcnow().isoformat(),  # 🟢 Fixed
                "outcomes": results,
                "success_rate": sum(1 for r in results if r.get("verified")) / len(results) if results else 0
            }

            success_rate = episode["success_rate"]
            logger.info(f"🗃️ [AAR] Mission success rate: {success_rate:.2%}")

            # Save to Vector DB so NEXT time 'conduct_reconnaissance' finds this
            doc_id = f"aar-{uuid.uuid4().hex[:6]}"
            await self.vector_client.upsert_document_from_text(
                document_content=str(episode),
                doc_id=doc_id,
                collection="episodes",
                metadata={"type": "mission_aar", "agent": self.name}
            )
            
            logger.info(f"🗃️ [AAR] Successfully stored AAR with document ID: {doc_id}")
            return True
        except Exception as e:
            logger.error(f"🗃️ [AAR] Failed to commit after action report: {str(e)}")
            raise

    @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.
        """
        mission_id = debrief.get("mission_id", "unknown")
        logger.info(f"📡 [COMPLETION] Signaling mission completion for: {mission_id}")

        try:
            payload = {
                "agent": self.name,
                "mission_id": mission_id,
                "final_report": debrief,
                "status": "COMPLETED"
            }

            # 🟢 DNS ALIGNMENT:
            # We use the confirmed local DNS entry for the orchestrator
            sutram_url = "http://sutram.guardianhub.com"

            async with AsyncClient() as client:
                logger.info(f"📡 [COMPLETION] Transmitting debrief to: {sutram_url}/v1/callback")
                response = await client.post(f"{sutram_url}/v1/callback", json=payload)

                if response.is_success:
                    logger.info("📡 [COMPLETION] Orchestrator acknowledged mission success.")
                return response.is_success

        except Exception as e:
            logger.error(f"📡 [COMPLETION] Transmission failure: {str(e)}")
            raise



    @activity.defn(name=ActivityRoles.SUPPORT)
    async def recruit_specialist_support(self, requirement: Dict[str, Any]) -> Dict[str, Any]:
        """
        SDK Default: Uses A2AClient to find a peer specialist.
        """
        specialty = requirement.get("specialty", "unknown")
        objective = requirement.get("objective", "unknown")
        logger.info(f"🤝 [SUPPORT] Recruiting specialist support for: {specialty}")
        logger.info(f"🤝 [SUPPORT] Agent: {self.name} seeking peer assistance for: {objective}")
        
        try:
            result = await self.a2a.recruit_ranger(
                specialty=specialty,
                objective=objective
            )
            
            if result.get("found"):
                logger.info(f"🤝 [SUPPORT] Successfully recruited specialist: {result.get('specialist', 'unknown')}")
            else:
                logger.warning(f"🤝 [SUPPORT] No suitable specialist found for: {specialty}")
            
            return result
        except Exception as e:
            logger.error(f"🤝 [SUPPORT] Failed to recruit specialist support: {str(e)}")
            raise