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 ActionStep,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.template.agent_plan import MacroPlan

logger = workflow.logger

@workflow.defn(name="SpecialistMissionWorkflow")
class SpecialistMissionWorkflow:
    @workflow.run
    async def run(self, payload: Dict[str, Any]) -> Dict[str, Any]:
        # 🟢 0. DEBUG INITIALIZATION
        debug_mode = payload.get("debug_mode", True)  # Default to True for your pressure test
        debug_ledger = {
            "initial_identity": {},
            "recon_findings": {},
            "historical_parallels": [],
            "tactical_audit_details": {},
            "step_by_step_traces": []
        }
        mapping = settings.workflow_settings.activity_mapping
        data = payload[0] if isinstance(payload, list) and payload else payload

        # 2. Extract the plan data
        raw_plan = data.get("macro_plan") or data.get("plan")

        if isinstance(raw_plan, list):
            logger.info("📦 Wrapping raw steps list into MacroPlan object")
            final_plan = {
                "steps": raw_plan,
                "reflection": data.get("rationale", "Synthesized from orchestrator"),
                "confidence_score": data.get("confidence_score", 1.0)
            }
        else:
            final_plan = raw_plan


        # 1. INITIALIZATION: Context Extraction
        trace_id = data.get("trace_id")
        session_id = data.get("session_id")
        agent_name = data.get("agent_name", "specialist")
        template_id = data.get("template_id", "TPL-GENERIC")
        sub_objective = data.get("sub_objective")
        is_dry_run = data.get("is_dry_run", False)
        depth_limit = data.get("depth_limit", 5)

        debug_ledger["initial_identity"] = {
            "agent": agent_name,
            "sub_objective": sub_objective,
            "template": template_id
        }

        # 2. RECONNAISSANCE: The Eyes (Avatar: Kurma)
        recon_request = DiscoveryRequest(
            session_id=session_id, trace_id=trace_id, template_id=template_id,
            agent_name=agent_name, sub_objective=sub_objective,
            environment=data.get("environment", "production"),
            depth_limit=depth_limit
        )
        intelligence: DiscoveryResponse = await workflow.execute_activity(
            mapping[ActivityRoles.RECON], arg=recon_request, **get_activity_options(ActivityRoles.RECON)
        )
        if debug_mode:
            debug_ledger["recon_findings"] = (
                intelligence.model_dump() if hasattr(intelligence, "model_dump")
                else intelligence
            )

        # 3. HISTORY: The Leela (Avatar: Varaha)
        # We pull these from the data so the Orchestrator can tune the "Wisdom Filter"
        history_limit = data.get("history_limit", 5)
        min_success = data.get("min_success_score", 0.7)  # 🎯 Defaulting to high-quality only

        history_request = HistoryRequest(
            session_id=session_id,
            trace_id=trace_id,
            template_id=template_id,
            agent_name=agent_name,
            search_query=sub_objective,
            limit=history_limit,
            min_success_score=min_success  # 🎯 FIXED: No more learning from failures
        )

        past_lessons: HistoryResponse = await workflow.execute_activity(
            mapping[ActivityRoles.HISTORY],
            arg=history_request,
            **get_activity_options(ActivityRoles.HISTORY)
        )

        if debug_mode:
            # This is where you'll find the "Colonel" memories
            episodes = (
                past_lessons.episodes if hasattr(past_lessons, "episodes")
                else past_lessons.get("episodes", [])
            )
            debug_ledger["historical_parallels"] = [
                {"content": getattr(e, 'content', e.get('content')),
                 "similarity": getattr(e, 'similarity', e.get('similarity')),
                 "id": getattr(e, 'id', e.get('id'))}
                for e in episodes
            ]

        # 4. TACTICAL AUDIT: The Safety Gate (Avatar: Narasimha)
        # Extract the risk appetite from the data
        risk_threshold = data.get("risk_threshold", 0.7)  # 🎯 Defaulting to a balanced stance
        validated_plan = MacroPlan.model_validate(final_plan)
        tactical_bundle = TacticalBundle(
            session_id=session_id,
            trace_id=trace_id,
            template_id=template_id,
            agent_name=agent_name,
            proposed_plan=validated_plan,
            recon_intelligence=intelligence,
            history_intelligence=past_lessons,
            risk_threshold=risk_threshold,  # 🎯 FIXED: Threshold is now mission-aware
            is_dry_run=is_dry_run
        )

        audit: TacticalAuditReport = await workflow.execute_activity(
            mapping[ActivityRoles.TACTICAL],
            arg=tactical_bundle,
            **get_activity_options(ActivityRoles.TACTICAL)
        )
        if debug_mode:
            debug_ledger["tactical_audit_details"] = {
                "decision": getattr(audit, 'decision', audit.get('decision', 'UNKNOWN')),
                "risk_score": getattr(audit, 'risk_score', audit.get('risk_score', 0.0)),
                "justification": getattr(audit, 'justification', audit.get('justification', 'N/A'))
            }


        audit_decision = audit.decision if hasattr(audit, "decision") else audit.get("decision")
        if audit_decision == "HALT":
            # Use get() pattern for the message as well
            reason = audit.justification if hasattr(audit, "justification") else audit.get("justification",
                                                                                           "No reason provided")
            workflow.logger.warning(f"🛑 Mission Halted: {reason}")
            return {"status": "HALTED", "reason": reason}

        # 5. ITERATIVE INTERVENTION: The Hands (Avatars: Vamana, Parashurama)
        execution_results = []
        for step in data.get("plan", {}).get("steps", []):
            action = ActionStep(
                session_id=session_id, trace_id=trace_id, template_id=template_id,
                agent_name=agent_name, action_name=step.get("tool_name"),
                action_input=step.get("tool_args"), step_id=step.get("step_id"),
                is_dry_run=is_dry_run
            )
            outcome: ActionOutcome = await workflow.execute_activity(
                mapping[ActivityRoles.INTERVENTION], arg=action, **get_activity_options(ActivityRoles.INTERVENTION)
            )

            check = AttainmentCheck(
                session_id=session_id, trace_id=trace_id, template_id=template_id,
                agent_name=agent_name, step_id=step.get("step_id"),
                verification_mode=step.get("verification_mode",None),
                action_name=step.get("tool_name"), action_outcome=outcome # 🎯 Passes full Outcome
            )
            report: AttainmentReport = await workflow.execute_activity(
                mapping[ActivityRoles.ATTAINMENT], arg=check, **get_activity_options(ActivityRoles.ATTAINMENT)
            )
            if debug_mode:
                debug_ledger["step_by_step_traces"].append({
                    "step_id": step.get("step_id"),
                    "action": step.get("tool_name"),
                    "input": step.get("tool_args"),
                    "outcome_observation": outcome.observation,
                    "attainment_status": report.attained
                })

            execution_results.append(report)

        # 🎯 These are now FUNCTIONAL parameters
        # 6. DEBRIEF: Narrative Synthesis (Avatar: Rama)
        total_steps = len(execution_results)
        success_count = sum(1 for r in execution_results if r.attained)
        any_step_failed = success_count < total_steps

        # 6. DEBRIEF
        debrief_brief = SummaryBrief(
            session_id=session_id,
            trace_id=trace_id,
            template_id=template_id,
            agent_name=agent_name,
            # 🎯 Ensure these are dicts, not objects
            step_results=[r.model_dump() for r in execution_results],
            original_objective=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: {type(debrief_brief)} - {debrief_brief}")

        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
        aar_request = AfterActionReport(
            session_id=session_id,
            trace_id=trace_id,
            template_id=template_id,
            agent_name=agent_name,
            mission_id=session_id,
            summary_brief=summary_text,
            aha_moment=flattened_aha,  # 🎯 Now populated with synthesized wisdom
            execution_data=[r.model_dump() for r in execution_results],
            success_score=1.0 if all(r.attained for r in execution_results) else 0.5
        )
        # 🎯 FIXED: Returning the Sovereign LearningAnchor
        anchor: LearningAnchor = await workflow.execute_activity(
            mapping[ActivityRoles.AAR],
            args=[aar_request.model_dump()],
            **get_activity_options(ActivityRoles.AAR)
        )
        # In mission_skeleton.py around line 205
        # 🎯 THE FULL SPECTRUM FIX
        memory_id = "unknown"
        anchor_is_success = False

        if anchor:
            # 1. Safely extract memory_id
            memory_id = anchor.memory_id if hasattr(anchor, 'memory_id') else anchor.get('memory_id', 'unknown')

            # 2. Safely extract success status (The 'Eyes' of the logic)
            anchor_is_success = anchor.success if hasattr(anchor, 'success') else anchor.get('success', False)

            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.
        if all(r.attained for r in execution_results):
            final_status = "COMPLETED"
        elif any(r.attained 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=session_id,
                trace_id=trace_id,
                template_id=template_id,
                agent_name=agent_name,
                mission_id=session_id,
                final_report=report_data,  # 🟢 Now a clean Dict
                mission_status=final_status,
                callback_url=data.get("callback_url", settings.endpoints.SUTRAM_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": 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
        }