# services/common/specialist_base.py

import uuid
from typing import Dict, Any, List
from fastapi import APIRouter, Depends, HTTPException
from pydantic import BaseModel

from .memory_manager import MemoryManager, HindsightSimulator
from .episodic_manager import EpisodicManager
from guardianhub.clients.a2a_client import  A2AClient
from .workflows import SpecialistMissionWorkflow
from ..models.template.agent_plan import AgentSubMission, MacroPlan


class SovereignSpecialistBase:
    """
    The DNA of a Colonel.
    Standardizes Handshake, Exploration, and Coming Home logic.
    """

    def __init__(self, agent_spec: Dict[str, Any], services: Dict[str, Any]):
        self.spec = agent_spec
        self.name = agent_spec['name']

        # 🧠 Core Intelligence Modules
        self.memory = MemoryManager(services['vector'], services['graph'], services['registry'])
        self.episodes = EpisodicManager(services['vector'], services['graph'], services['registry'])
        self.a2a = A2AClient(sender_name=self.name, consul=services['consul'])
        self.llm = services['llm']

        self.router = APIRouter(prefix="/v1/mission")
        self._setup_routes()

    def _setup_routes(self):
        @self.router.post("/propose")
        async def propose(mission: AgentSubMission):
            return await self.handle_proposal(mission)

        @self.router.post("/execute")
        async def execute(plan: MacroPlan):
            return await self.handle_execution(plan)

    # --- 1. THE BRAIN: Tactical Planning Handshake ---
    async def handle_proposal(self, mission: AgentSubMission) -> ProposalResponse:
        """
        PASS 2 Handshake: Generates a plan armed with Local Context & Hindsight.
        """
        # A. Sensory Discovery (The 'Fog of War' clearing)
        context = await self.memory.get_reasoning_context(
            query=mission.sub_objective,
            template_id=mission.metadata.get("template_id", "TPL-GENERIC"),
            tools=self.spec['tools']
        )

        # B. Hindsight Simulation (Risk Assessment)
        risk_profile = HindsightSimulator.calculate_step_risk(
            mission.sub_objective,
            context['episodes']
        )

        # C. Local Planning
        # Uses the Specialist Prompt armed with the 7k token schema
        plan = await self.llm.generate_specialist_plan(
            instruction=mission.sub_objective,
            context=context,
            spec=self.spec,
            risk_profile=risk_profile
        )

        return ProposalResponse(plan=plan, rationale=plan.reflection)

    # --- 2. THE MUSCLE: Temporal Dispatch ---
    async def handle_execution(self, plan: MacroPlan) -> ExecutionTriggerResponse:
        """
        Launches the durable mission.
        """
        workflow_id = f"mission-{plan.session_id}-{uuid.uuid4().hex[:6]}"

        # Start the durable ReAct loop in Temporal
        handle = await self.temporal_client.start_workflow(
            "SpecialistMissionWorkflow",
            args=[plan.model_dump()],
            id=workflow_id,
            task_queue=f"{self.name}-queue"
        )

        return ExecutionTriggerResponse(workflow_id=handle.id, status="running")

    # --- 3. THE MEMORY: Post-Mission Promotion ---
    async def promote_findings(self, session_id: str, results: Dict[str, Any]):
        """
        The 'Coming Home' Memory Update.
        Promotes local tool findings to the Global Blackboard.
        """
        # Promote Facts to Neo4j
        await self.episodes.record_resource_discovery(session_id, results)

        # Record the Session Narrative
        await self.episodes.record_session_narrative(session_id, results)

        logger.info(f"🛰️ Blackboard Updated for session {session_id}")