from datetime import timedelta
from typing import Dict, Any
from temporalio import workflow

# Contract Imports
from guardianhub.models.agent.contracts.discovery import DiscoveryRequest , DiscoveryResponse
from guardianhub.models.agent.contracts.history import HistoryRequest,HistoryResponse
from guardianhub.models.agent.contracts.tactical import TacticalBundle,TacticalAuditReport
from guardianhub.models.agent.contracts.action import ActionOutcome
from guardianhub.models.agent.contracts.attainment import AttainmentCheck,AttainmentReport
from guardianhub.models.agent.contracts.debrief import SummaryBrief,IntelligenceReport
from guardianhub.models.agent.contracts.learning import AfterActionReport,LearningAnchor
from guardianhub.models.agent.contracts.completion import MissionManifest,CallbackAck

from .agent_contract import ActivityRoles
from .constants import get_activity_options
from guardianhub.config.settings import settings
from ...models.agent_models import AgentSubMission
from ...models.template.agent_plan import MacroPlan, MissionRequest

logger = workflow.logger


@workflow.defn(name="SpecialistMissionWorkflow")
class SpecialistMissionWorkflow:
    @workflow.run
    async def run(self, input_payload: Any) -> Dict[str, Any]:
        # 1. 🛡️ CONTRACT INGESTION
        # Standardizes the input immediately. Handles raw dict or [dict].
        try:
            raw_data = input_payload[0] if isinstance(input_payload, list) else input_payload
            request = MissionRequest.model_validate(raw_data)
        except Exception as e:
            workflow.logger.error(f"❌ [DENIED] Input Schema Violation: {str(e)}")
            raise

        # 🟢 0. RESOURCE PREP
        mapping = settings.workflow_settings.activity_mapping
        debug_mode = request.metadata.get("debug_mode", True)
        execution_results = []

        # 2. 📡 RECONNAISSANCE: The Eyes (Kurma)
        intelligence = await workflow.execute_activity(
            mapping[ActivityRoles.RECON],
            arg=DiscoveryRequest(
                sub_objective=request.sub_objective,
                environment=request.metadata.get("environment", "production"),
                **request.technical_context  # 🚀 Standardized injection
            ),
            **get_activity_options(ActivityRoles.RECON)
        )

        # 3. 📚 HISTORY: The Leela (Varaha)
        past_lessons = await workflow.execute_activity(
            mapping[ActivityRoles.HISTORY],
            arg=HistoryRequest(
                search_query=request.sub_objective,
                limit=request.metadata.get("history_limit", 5),
                min_success_score=0.7,
                **request.technical_context
            ),
            **get_activity_options(ActivityRoles.HISTORY)
        )

        # 4. 🛡️ TACTICAL AUDIT: The Safety Gate (Narasimha)
        # 🎯 FIX: Intelligence comes from the RECON activity results (Kurma)
        audit = await workflow.execute_activity(
            mapping[ActivityRoles.TACTICAL],
            arg=TacticalBundle(
                proposed_plan=request.macro_plan,
                recon_intelligence=intelligence,  # 👈 Passed from Kurma
                history_intelligence=past_lessons,
                risk_threshold=request.metadata.get("risk_threshold", 0.7),
                **request.technical_context
            ),
            **get_activity_options(ActivityRoles.TACTICAL)
        )

        workflow.logger.info(f" audit output from LLM : {audit}")

        audit_score = getattr(audit, 'risk_score', audit.get('risk_score', 0.0))
        audit_decision = getattr(audit, 'decision', audit.get('decision', 'UNKNOWN'))


        if audit_decision == "HALT":
            workflow.logger.warning(
                f"🛑 Mission Halted: {getattr(audit, 'justification', audit.get('justification', 'Unknown'))}")
            return {
                "status": "HALTED",
                "mission_id": request.session_id,
                "reason": "Safety Audit Violation"
            }

        # 5. 🚀 ITERATIVE INTERVENTION: The Hands (Vamana/Parashurama)
        # SpecialistMissionWorkflow.run
        # 5. 🚀 MISSION EXECUTION: The Hands (Universal Intervention)
        # Instead of a loop, we pass the ENTIRE plan to the Agent once.

        mission_outcome = await workflow.execute_activity(
            mapping[ActivityRoles.INTERVENTION],
            arg=request,  # 🎯 Pass the whole MissionRequest!
            **get_activity_options(ActivityRoles.INTERVENTION)
        )

        # Execution results are now synthesized inside the agent's response
        # 🚀 RESILIENT ACCESS
        if isinstance(mission_outcome, dict):
            execution_results = mission_outcome.get("step_reports", [])
        else:
            execution_results = getattr(mission_outcome, "step_reports", [])

        # 6. DEBRIEF: Narrative Synthesis (Avatar: Rama)
            # 🎯 6. DEBRIEF: Narrative Synthesis
        total_steps = len(execution_results)

        # 🚀 FIX: Robustly check for 'attained' status regardless of type
        def is_attained(res):
            if hasattr(res, 'attained'):
                return res.attained
            if isinstance(res, dict):
                return res.get('attained', False)
            return False

        success_count = sum(1 for r in execution_results if is_attained(r))
        any_step_failed = success_count < total_steps

        # 6. DEBRIEF
        # 🚀 FIX: Resilient step_results packaging
        formatted_results = []
        for r in execution_results:
            if hasattr(r, "model_dump"):
                formatted_results.append(r.model_dump())
            elif isinstance(r, dict):
                formatted_results.append(r)
            else:
                # Fallback for basic types
                formatted_results.append({"result": str(r)})

        debrief_brief = SummaryBrief(
            session_id=request.session_id,
            trace_id=request.trace_id,
            template_id=request.template_id,
            agent_name=request.agent_name,
            step_results=formatted_results,  # 🎯 Clean list of dicts
            original_objective=request.sub_objective,
            is_partial_success=any_step_failed,
            total_steps=total_steps,
            success_count=success_count
        )
        
        # 🎯 DEBUG: Log the debrief_brief object
        logger.info(f"🎯 [DEBRIEF] Created SummaryBrief: - {debrief_brief.model_dump()}")

        final_report: IntelligenceReport = await workflow.execute_activity(
            mapping[ActivityRoles.DEBRIEF],
            args=[debrief_brief.model_dump()],
            **get_activity_options(ActivityRoles.DEBRIEF)
        )

        # 7. AAR: Learning Commitment (Avatar: Krishna)
        # 🎯 THE RESILIENT ACCESS FIX
        # Check if the report is an object with the attribute, otherwise use get()
        # In mission_skeleton.py around line 183
        # 🎯 THE FIX: Handle both Object and Dict formats
        def get_val(obj, key, default=None):
            if hasattr(obj, key):
                return getattr(obj, key)
            if isinstance(obj, dict):
                return obj.get(key, default)
            return default

        # Use the helper to extract values safely
        summary_text = get_val(final_report, "narrative_summary", "No summary available.")
        takeaways_list = get_val(final_report, "key_takeaways", [])

        flattened_aha = " | ".join(takeaways_list) if takeaways_list else "Mission completed with standard results."
        # Now use them in your final call
        execution_data_bundle = []
        for r in execution_results:
            if hasattr(r, "model_dump"):
                execution_data_bundle.append(r.model_dump())
            elif isinstance(r, dict):
                execution_data_bundle.append(r)
            else:
                execution_data_bundle.append({"raw_result": str(r)})

        # 🚀 THE FINAL FIX: Robust Success Score calculation
        def check_attained(res):
            if hasattr(res, "attained"):
                return res.attained
            if isinstance(res, dict):
                return res.get("attained", False)
            return False

        # Now calculate safely

        # 🎯 Robust Success Check
        all_attained = all(
            (r.attained if hasattr(r, 'attained') else r.get('attained', False))
            for r in execution_results
        )

        # 🎯 7. AAR: Learning Commitment
        aar_request = AfterActionReport(
            session_id=request.session_id,
            trace_id=request.trace_id,
            template_id=request.template_id,
            agent_name=request.agent_name,
            mission_id=request.session_id,
            summary_brief=summary_text,
            aha_moment=flattened_aha,  # 🎯 Now populated with synthesized wisdom
            execution_data=execution_data_bundle,
            success_score=1.0 if all_attained else 0.5  # 🎯 Use the safe check
        )

        # RESILIENT PACKAGING: Handle both Pydantic and Dict formats

        # 🎯 FIXED: Returning the Sovereign LearningAnchor
        anchor: LearningAnchor = await workflow.execute_activity(
            mapping[ActivityRoles.AAR],
            args=[aar_request.model_dump()],
            **get_activity_options(ActivityRoles.AAR)
        )
        # 🎯 THE FULL SPECTRUM RESILIENCE FIX
        memory_id = "unknown"
        anchor_is_success = False

        if isinstance(anchor, bool):
            # Case A: Simple Status Return
            anchor_is_success = anchor
            memory_id = request.session_id if anchor else "failed_to_anchor"
            workflow.logger.info(f"🗃️ [AAR] Simple Boolean Anchor received. Status: {anchor_is_success}")

        elif anchor is not None:
            # Case B: Object or Dict Return
            # We use a helper-style approach to get the values safely
            if hasattr(anchor, 'success'):  # Pydantic Model
                anchor_is_success = anchor.success
                memory_id = anchor.memory_id
            elif isinstance(anchor, dict):  # Python Dict
                anchor_is_success = anchor.get('success', False)
                memory_id = anchor.get('memory_id', 'unknown')

            workflow.logger.info(f"🗃️ [AAR] Wisdom anchored at: {memory_id} | Status: {anchor_is_success}")

        # 3. Final linkage without 'Blind' property access
        anchored_wisdom_id = memory_id if anchor_is_success else None


        # 🎯 FUNCTIONAL USAGE:
        # We pass the memory_id back to Sutram.
        # This allows Sutram to 'pin' this specific lesson to the user's project graph.


        # 8. COMPLETION: The Callback (Avatar: Buddha)

        # 🎯 DERIVE FINAL STATUS:
        # Based on our attainment checks, we determine the high-level signal.
        # 🎯 Robust Final Status Logic
        if all_attained:  # Use the variable we just calculated safely
            final_status = "COMPLETED"
        elif any((r.attained if hasattr(r, 'attained') else r.get('attained', False)) for r in execution_results):
            final_status = "PARTIAL"
        else:
            final_status = "FAILED"

        # In mission_skeleton.py
        try:
            # 🎯 THE RESILIENT PACKAGING FIX
            # Handle both Pydantic object and serialized dictionary
            if hasattr(final_report, "model_dump"):
                report_data = final_report.model_dump()
            elif isinstance(final_report, dict):
                report_data = final_report
            else:
                # Fallback for unexpected types
                report_data = {"summary": str(final_report)}

            # Now inject the anchored wisdom safely
            report_data.update({
                "anchored_wisdom_id": anchored_wisdom_id,
                "success_score": getattr(aar_request, 'success_score', 0.95)  # High fidelity default
            })

            manifest = MissionManifest(
                session_id=request.session_id,
                trace_id=request.trace_id,
                template_id=request.template_id,
                agent_name=request.agent_name,
                mission_id=request.session_id,
                final_report=report_data,  # 🟢 Now a clean Dict
                mission_status=final_status,
                callback_url=request.callback_url
            )
        except Exception as e:
            workflow.logger.error(f"❌ [MANIFEST] Failed to package mission results: {str(e)}")
            raise


        # 8. COMPLETION: The Callback (Avatar: Buddha)
        ack: CallbackAck = await workflow.execute_activity(
            mapping[ActivityRoles.COMPLETION],
            arg=manifest,
            **get_activity_options(ActivityRoles.COMPLETION)
        )

        # 🎯 FUNCTIONAL USAGE: Recursive Continuity

        # 🎯 THE RESILIENT ACKNOWLEDGMENT FIX
        # Handle both Object and Dict formats for the Callback result
        is_received = ack.orchestrator_received if hasattr(ack, 'orchestrator_received') else ack.get(
            'orchestrator_received', False)

        if not is_received:
            error_msg = ack.error_message if hasattr(ack, 'error_message') else ack.get('error_message',
                                                                                        'Unknown Error')
            workflow.logger.error(f"❌ [COMPLETION] Sutram failed to acknowledge: {error_msg}")

        # 🎯 THE CONTINUITY FIX (Token extraction)
        is_success = ack.success if hasattr(ack, 'success') else ack.get('success', False)
        token = ack.next_mission_token if hasattr(ack, 'next_mission_token') else ack.get('next_mission_token')

        continuation_token = token if is_success else None
        report_text = final_report.narrative_summary if hasattr(final_report,
                                                                'narrative_summary') else final_report.get(
            'narrative_summary', 'Mission summary processed.')
        return {
            "status": final_status,
            "mission_id": request.session_id,
            "anchored_wisdom_id": anchored_wisdom_id,
            # 🎯 The Handover: Passing the baton to the next mission or the UI
            "next_steps_token": continuation_token,
            "debug_mode": debug_mode,
            "report": report_text
        }