import asyncio
import time
import logging
import json
from typing import Optional, List, Dict, Any
from mesh_core import agent_pb2

logger = logging.getLogger(__name__)

class SubAgent:
    """
    A stateful watcher for a specific task on an agent node.
    Handles execution, result accumulation, and state monitoring.
    """
    def __init__(self, name: str, task_fn, args: dict, retries: int = 1, 
                 llm_provider=None, assistant=None, on_event=None):
        self.name = name
        self.task_fn = task_fn
        self.args = args
        self.retries = retries
        self.llm = llm_provider
        self.assistant = assistant
        self.on_event = on_event
        self.status = "PENDING"
        self.result = None
        self.start_time = None
        self.end_time = None
        self.error = None
        self.task_id = None

    async def run(self):
        self.start_time = time.time()
        self.status = "RUNNING"
        
        # Every Sub-Agent task is now strictly atomic. 
        # The AI 'Observe-Think-Act' loop happens in the Main RagPipeline.
        return await self._run_standard()

    async def _run_standard(self):
        """Atomic execution with monitoring and simple retry logic."""
        node_id = self.args.get("node_id") or "swarm"
        
        # Improved task description logic
        cmd = self.args.get("cmd") or self.args.get("command")
        action = self.args.get("action")
        path = self.args.get("path")
        
        if cmd:
            item = cmd
        elif action and path:
            item = f"{action} {path}"
        elif action:
            item = str(action)
        else:
            item = self.name # Fallback to skill name if no command/action found
        
        # Phase 1: Dispatch
        if self.on_event:
            # Descriptive task name logic to reduce "Starting atomic step: ifconfig" noise
            desc = f"Performing {self.name}"
            if "command" in self.args or "cmd" in self.args:
                desc = f"Running: {self.args.get('command') or self.args.get('cmd')}"
            elif "path" in self.args:
                desc = f"Inspecting: {self.args['path']}"
            elif "action" in self.args:
                desc = f"Executing {self.args['action']}"
                
            await self.on_event({"type": "subagent_thought", "node_id": node_id, "content": f"✅ {desc}"})

        for attempt in range(self.retries + 1):
            try:
                # Dispatch task via assistant (which handles gRPC signing/sending)
                import inspect
                if inspect.iscoroutinefunction(self.task_fn):
                    self.result = await self.task_fn(**self.args)
                else:
                    self.result = await asyncio.to_thread(self.task_fn, **self.args)
                
                # If it's a long-running task (Wait-Pending), monitor it with heartbeats
                if isinstance(self.result, dict) and self.result.get("status") == "TIMEOUT_PENDING":
                    self.result = await self._monitor_atomic_task(self.result, node_id)

                if isinstance(self.result, dict) and self.result.get("error"):
                    err_msg = str(self.result.get("error")).lower()
                    if any(x in err_msg for x in ["timeout", "offline", "disconnected", "capacity", "rejected"]):
                        if attempt < self.retries:
                            backoff = (attempt + 1) * 3
                            self.status = f"RETRYING ({attempt+1}/{self.retries})"
                            await asyncio.sleep(backoff)
                            continue
                self.status = "COMPLETED"
                break
            except Exception as e:
                logger.error(f"SubAgent {self.name} error: {e}")
                self.error = str(e)
                if attempt < self.retries:
                    self.status = f"ERROR_RETRYING ({attempt+1}/{self.retries})"
                    await asyncio.sleep(2)
                else:
                    self.status = "FAILED"
        
        self.end_time = time.time()
        
        # --- UI Intelligence Report ---
        # We still emit the summary/status to the UI for human readability
        if self.on_event:
            summary = None
            if self.llm and self.result:
                summary = await self._generate_summary()
            
            if summary:
                rep = f"🧠 {summary}"
            else:
                # Fallback: Check if there's any data we can show directly
                if isinstance(self.result, dict):
                    eval_res = self.result.get("eval_result", "")
                    # Guard against the string "None" which means the JS returned null/undefined
                    if eval_res and eval_res != "None":
                        rep = f"📊 Data Extracted: {str(eval_res)[:500]}..."
                    elif self.result.get("stdout"):
                        rep = f"📟 Output: {str(self.result['stdout'])[:500]}..."
                    elif self.result.get("content"):
                        rep = f"📄 Read {len(self.result['content'])} bytes."
                    elif self.result.get("results") and isinstance(self.result["results"], list):
                        count = len(self.result["results"])
                        success_count = sum(1 for r in self.result["results"] if r.get("success"))
                        rep = f"🔍 Fetched {success_count}/{count} URLs successfully."
                    else:
                        rep = "✅ Step Finished." if not self.error else f"❌ Step Failed: {self.error}"
                else:
                    rep = "✅ Step Finished." if not self.error else f"❌ Step Failed: {self.error}"

            await self.on_event({"type": "subagent_thought", "node_id": node_id, "content": rep})

        # --- High-Fidelity Handoff with Smart Summary ---
        if isinstance(self.result, dict):
             # 1. Generate summary for human UI and as a 'hint' for the Master AI
             summary = None
             if self.llm:
                  # We skip summary for simple directory listings/cat if they are clean 
                  # but ALWAYS generate for errors or large blobs
                  is_data_tool = any(x in self.name for x in ["ls", "cat", "explorer", "find"])
                  if not is_data_tool or self.error or self.status != "COMPLETED":
                       summary = await self._generate_summary()
             
             # 2. Enrich the result object
             self.result["sub_agent_status"] = self.status
             if summary:
                  self.result["sub_agent_summary"] = summary
             if self.error:
                  self.result["sub_agent_error"] = self.error
             
             return self.result
        
        # Fallback for non-dictionary results
        final_summary = await self._generate_summary() if self.llm else "Completed"
        return {
            "success": not self.error,
            "result": self.result,
            "summary": final_summary,
            "status": self.status,
            "error": self.error
        }

    async def _generate_summary(self) -> Optional[str]:
        """Uses LLM to summarize logs, file system changes, or errors into a single report sentence."""
        try:
            # Prepare a more generous context for the sub-agent reporter
            # We provide a larger chunk of the raw result
            raw_res = str(self.result)[:20000] 
            prompt = (
                f"You are a Sub-Agent reporter on node '{self.args.get('node_id')}'.\n"
                f"Task Performed: {self.name} {self.args}\n"
                f"Raw Outcome: {raw_res}\n"
                f"Status: {self.status}\n"
                f"Errors: {self.error if self.error else 'None'}\n\n"
                "Task: Summarize the result for your Master-Architect in ONE CONCISE SENTENCE. "
                "CRITICAL: If the task was to find/extract information (news, stocks, headlines), your summary MUST INCLUDE the actual data found. "
                "For example: 'Found headlines: [Title A], [Title B]...' or 'PANW price is $320.50'. "
                "Do not just say 'headlines extracted'. SHOW THEM. "
                "Include key findings. Do not use conversational filler. Be the eyes and ears on the ground."
            )
            
            response = await self.llm.acompletion(prompt=prompt, max_tokens=4096, temperature=0)
            summary = response.choices[0].message.content.strip()
            summary = summary.strip('"').strip("'")
            if ":" in summary[:15]:
                summary = summary.split(":", 1)[1].strip()
            return summary
        except Exception as e:
            logger.warning(f"Failed to generate sub-agent summary: {e}")
            return None

    async def _monitor_atomic_task(self, partial_res: dict, node_id: str):
        """Monitors a single atomic task until finished without making decisions."""
        task_id = partial_res.get("task_id")
        if not task_id: return partial_res
        
        task_map = {node_id: task_id}
        max_checks = 120 # 10 minutes max for an atomic task
        for _ in range(max_checks):
            # Non-blocking peek
            res = await asyncio.to_thread(self.assistant.wait_for_swarm, task_map, timeout=0, no_abort=True)
            node_res = res.get(node_id, {})
            
            if node_res.get("status") not in ["RUNNING", "TIMEOUT_PENDING"]:
                # Finished! Collect final and return
                return await asyncio.to_thread(self.assistant.wait_for_swarm, task_map, timeout=2)
            
            # Simple heartbeat for UI
            if self.on_event:
                 await self.on_event({"type": "subagent_thought", "node_id": node_id, "content": "*(Still executing...)*"})
            
            await asyncio.sleep(5)
            
        return {"error": "SubAgent timed out waiting for atomic task."}

    def get_elapsed(self) -> int:
        if not self.start_time:
            return 0
        end = self.end_time or time.time()
        return int(end - self.start_time)
