"""
Agno Planning Strategy

Implementation using Agno framework for task planning.
"""

from typing import Dict, Any, AsyncIterator
import structlog

from control_plane_api.app.services.planning_strategy import PlanningStrategy
from control_plane_api.app.models.task_planning import TaskPlanResponse
from control_plane_api.app.lib.task_planning import create_planning_agent

logger = structlog.get_logger(__name__)


class AgnoPlanningStrategy(PlanningStrategy):
    """Task planning using Agno framework"""

    @property
    def name(self) -> str:
        return "agno"

    async def plan_task(self, planning_prompt: str) -> TaskPlanResponse:
        """Generate task plan using Agno (non-streaming)"""
        logger.info("using_agno_strategy_nonstream")

        # Create planning agent
        planning_agent = create_planning_agent(
            organization_id=self.organization_id,
            db=self.db,
            api_token=self.api_token
        )

        # Run agent
        response = planning_agent.run(planning_prompt)

        # Extract plan
        if isinstance(response.content, TaskPlanResponse):
            plan = response.content
        elif isinstance(response.content, dict):
            plan = TaskPlanResponse(**response.content)
        else:
            raise ValueError(f"Unexpected response type from Agno: {type(response.content)}")

        logger.info("agno_plan_generated", title=plan.title)
        return plan

    async def plan_task_stream(self, planning_prompt: str) -> AsyncIterator[Dict[str, Any]]:
        """Generate task plan using Agno with REAL agent streaming (tool calls, thinking, etc.)"""
        logger.info("using_agno_strategy_stream_with_hooks")

        try:
            import asyncio
            import time

            # Yield initial progress
            yield {
                "event": "progress",
                "data": {
                    "stage": "initializing",
                    "message": "🤖 Initializing planning agent with streaming...",
                    "progress": 10
                }
            }

            # Create event queue for real-time agent streaming
            event_queue = asyncio.Queue()
            main_loop = asyncio.get_event_loop()

            # Event publisher function for streaming hook
            def publish_event(event: dict):
                """Publish event to queue (sync wrapper for hook)"""
                try:
                    logger.info("publish_event_called", event_type=event.get("event"), from_thread=True)
                    main_loop.call_soon_threadsafe(event_queue.put_nowait, event)
                    logger.info("event_queued_successfully", event_type=event.get("event"))
                except Exception as e:
                    logger.error("failed_to_queue_event", event_type=event.get("event"), error=str(e), exc_info=True)

            # Import and create streaming hook
            from control_plane_api.app.lib.task_planning.streaming_hook import PlanningStreamingHook
            streaming_hook = PlanningStreamingHook(publish_event)

            # Create planning agent WITH streaming hook
            planning_agent = create_planning_agent(
                organization_id=self.organization_id,
                db=self.db,
                api_token=self.api_token,
                tool_hook=streaming_hook.create_tool_hook()  # Enable real-time tool streaming!
            )

            yield {
                "event": "progress",
                "data": {
                    "stage": "generating",
                    "message": "🧠 Agent is analyzing your task...",
                    "progress": 30
                }
            }

            # Track state
            agent_completed = False
            agent_result = None
            agent_error = None
            reasoning_chunks = []
            final_response = None

            # Run agent in background thread (non-blocking)
            async def run_agent_in_background():
                nonlocal agent_completed, agent_result, agent_error, final_response
                try:
                    logger.info("starting_agent_run_with_streaming", prompt_length=len(planning_prompt))

                    # Run synchronous agent.run() in thread pool
                    loop = asyncio.get_event_loop()
                    chunks = await loop.run_in_executor(
                        None,
                        lambda: list(planning_agent.run(planning_prompt, stream=True))
                    )

                    # Process chunks to extract final response and reasoning
                    for chunk in chunks:
                        # Check for reasoning content
                        if hasattr(chunk, 'reasoning_content') and chunk.reasoning_content:
                            reasoning_text = str(chunk.reasoning_content)
                            reasoning_chunks.append(reasoning_text)

                            # Publish reasoning as thinking event
                            if len(reasoning_text.strip()) > 10:
                                publish_event({
                                    "event": "thinking",
                                    "data": {
                                        "content": reasoning_text,
                                        "timestamp": time.time()
                                    }
                                })

                        # Check for content chunks
                        elif hasattr(chunk, 'content') and chunk.content and not isinstance(chunk.content, (dict, TaskPlanResponse)):
                            reasoning_text = str(chunk.content)
                            reasoning_chunks.append(reasoning_text)

                            if len(reasoning_text.strip()) > 10:
                                publish_event({
                                    "event": "thinking",
                                    "data": {
                                        "content": reasoning_text,
                                        "timestamp": time.time()
                                    }
                                })

                        # Check for final structured response
                        if hasattr(chunk, 'content') and isinstance(chunk.content, (dict, TaskPlanResponse)):
                            final_response = chunk

                    agent_result = chunks
                    agent_completed = True
                    logger.info("agent_run_completed_with_streaming")

                except Exception as e:
                    agent_error = e
                    agent_completed = True
                    logger.error("agent_run_error_with_streaming", error=str(e))

            # Start agent task
            agent_task = asyncio.create_task(run_agent_in_background())

            # Stream events from queue as they arrive
            logger.info("starting_event_stream_loop")
            last_heartbeat = time.time()

            while not agent_completed:
                try:
                    # Try to get event from queue (with short timeout)
                    event = await asyncio.wait_for(event_queue.get(), timeout=2.0)

                    # Stream event to client
                    event_type = event.get("event")
                    event_data = event.get("data", {})

                    logger.info("streaming_event_to_client", event_type=event_type, has_data=bool(event_data))

                    if event_type in ["thinking", "tool_call", "tool_result"]:
                        yield {"event": event_type, "data": event_data}

                except asyncio.TimeoutError:
                    # No event in queue - send heartbeat if needed
                    if time.time() - last_heartbeat > 5:
                        yield {"event": "heartbeat", "data": {"timestamp": time.time()}}
                        last_heartbeat = time.time()

                    # Check if agent finished while we were waiting
                    if agent_task.done():
                        agent_completed = True
                        break

            # Check for agent error
            if agent_error:
                raise agent_error

            # Drain any remaining events in queue
            while not event_queue.empty():
                try:
                    event = event_queue.get_nowait()
                    event_type = event.get("event")
                    event_data = event.get("data", {})

                    if event_type in ["thinking", "tool_call", "tool_result"]:
                        yield {"event": event_type, "data": event_data}
                except asyncio.QueueEmpty:
                    break

            logger.info("event_stream_completed", reasoning_chunks=len(reasoning_chunks))

            # Extract final plan from agent result
            if not final_response and agent_result:
                for chunk in agent_result:
                    if hasattr(chunk, 'content') and isinstance(chunk.content, (dict, TaskPlanResponse)):
                        final_response = chunk
                        break

            # Extract plan
            if final_response:
                if isinstance(final_response.content, TaskPlanResponse):
                    plan = final_response.content
                elif isinstance(final_response.content, dict):
                    plan = TaskPlanResponse(**final_response.content)
                else:
                    raise ValueError(f"Unexpected response type from Agno: {type(final_response.content)}")
            else:
                raise ValueError("Agent did not return a final plan")

            logger.info("agno_plan_generated_stream", title=plan.title)

            # Yield complete event with the plan
            yield {
                "event": "complete",
                "data": {
                    "plan": plan.model_dump(),
                    "progress": 100,
                    "message": "✅ Plan ready!"
                }
            }

        except Exception as e:
            logger.error("agno_streaming_error", error=str(e))
            yield {
                "event": "error",
                "data": {
                    "message": f"Planning failed: {str(e)}",
                    "error": str(e)
                }
            }
