"""
Week 5 Features API Endpoints
Autonomous workflow engine, agent communication, and enhanced reasoning capabilities.
"""

from __future__ import annotations

from typing import Any, Dict, List, Optional
from fastapi import APIRouter, Depends, HTTPException, Query, BackgroundTasks
from pydantic import BaseModel

from app.core.auth import get_current_tenant
from app.services.autonomous_workflow_engine import (
    autonomous_workflow_engine,
    WorkflowDefinition,
    WorkflowExecution,
    WorkflowStatus,
    TaskStatus
)
from app.services.agent_communication import (
    agent_communication_manager,
    AgentMessage,
    MessageType,
    MessagePriority
)
from app.services.enhanced_reasoning import (
    enhanced_reasoning_engine,
    ReasoningMethod,
    ReasoningResult
)
from app.services.agentic_rag_service import AgentType, AgentTask
from app.core.cache import cache_service
import uuid
from datetime import datetime


router = APIRouter()


# Request/Response Models
class WorkflowCreateRequest(BaseModel):
    name: str
    description: str
    tasks: List[Dict[str, Any]]
    dependencies: Optional[Dict[str, List[str]]] = None
    max_parallel_tasks: int = 5
    timeout_seconds: int = 300


class WorkflowExecuteRequest(BaseModel):
    workflow_id: str
    context: Optional[Dict[str, Any]] = None


class AgentMessageRequest(BaseModel):
    recipient: str
    message_type: str
    payload: Dict[str, Any]
    priority: str = "normal"
    correlation_id: Optional[str] = None


class ReasoningRequest(BaseModel):
    query: str
    method: str = "chain_of_thought"
    max_steps: int = 10
    context: Optional[Dict[str, Any]] = None


class WorkflowStatusResponse(BaseModel):
    execution_id: str
    status: str
    task_results: Dict[str, Any]
    task_status: Dict[str, str]
    start_time: Optional[str] = None
    end_time: Optional[str] = None
    error: Optional[str] = None


# Workflow Engine Endpoints
@router.post("/workflows", summary="Create a new workflow")
async def create_workflow(
    request: WorkflowCreateRequest,
    tenant_id: str = Depends(get_current_tenant),
) -> Dict[str, Any]:
    """Create a new workflow definition."""
    try:
        # Convert task dictionaries to AgentTask objects
        tasks = []
        for task_data in request.tasks:
            task = AgentTask(
                id=task_data.get("id", str(uuid.uuid4())),
                agent_type=AgentType(task_data["agent_type"]),
                description=task_data["description"],
                input_data=task_data.get("input_data", {}),
                dependencies=task_data.get("dependencies", []),
                priority=task_data.get("priority", 1),
                created_at=datetime.utcnow()
            )
            tasks.append(task)
        
        workflow = await autonomous_workflow_engine.create_workflow(
            name=request.name,
            description=request.description,
            tasks=tasks,
            dependencies=request.dependencies,
            max_parallel_tasks=request.max_parallel_tasks,
            timeout_seconds=request.timeout_seconds
        )
        
        return {
            "workflow_id": workflow.id,
            "name": workflow.name,
            "description": workflow.description,
            "task_count": len(workflow.tasks),
            "status": "created"
        }
        
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Failed to create workflow: {str(e)}")


@router.post("/workflows/{workflow_id}/execute", summary="Execute a workflow")
async def execute_workflow(
    workflow_id: str,
    request: WorkflowExecuteRequest,
    background_tasks: BackgroundTasks,
    tenant_id: str = Depends(get_current_tenant),
) -> Dict[str, Any]:
    """Execute a workflow."""
    try:
        # Get agents from agentic RAG service
        from app.services.agentic_rag_service import agentic_rag_service
        agents = agentic_rag_service.agents
        
        # Execute workflow
        execution = await autonomous_workflow_engine.execute_workflow(
            workflow_id=workflow_id,
            tenant_id=str(tenant_id),
            agents=agents,
            context=request.context
        )
        
        return {
            "execution_id": execution.id,
            "workflow_id": workflow_id,
            "status": execution.status.value,
            "start_time": execution.start_time.isoformat() if execution.start_time else None,
            "message": "Workflow execution started"
        }
        
    except ValueError as e:
        raise HTTPException(status_code=404, detail=str(e))
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Failed to execute workflow: {str(e)}")


@router.get("/workflows/{execution_id}/status", summary="Get workflow execution status")
async def get_workflow_status(
    execution_id: str,
    tenant_id: str = Depends(get_current_tenant),
) -> WorkflowStatusResponse:
    """Get the status of a workflow execution."""
    try:
        execution = await autonomous_workflow_engine.get_workflow_status(execution_id)
        
        if not execution:
            raise HTTPException(status_code=404, detail="Workflow execution not found")
        
        return WorkflowStatusResponse(
            execution_id=execution.id,
            status=execution.status.value,
            task_results=execution.task_results,
            task_status={k: v.value for k, v in execution.task_status.items()},
            start_time=execution.start_time.isoformat() if execution.start_time else None,
            end_time=execution.end_time.isoformat() if execution.end_time else None,
            error=execution.error
        )
        
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Failed to get workflow status: {str(e)}")


@router.post("/workflows/{execution_id}/cancel", summary="Cancel a workflow execution")
async def cancel_workflow(
    execution_id: str,
    tenant_id: str = Depends(get_current_tenant),
) -> Dict[str, Any]:
    """Cancel a running workflow execution."""
    try:
        success = await autonomous_workflow_engine.cancel_workflow(execution_id)
        
        if not success:
            raise HTTPException(status_code=404, detail="Workflow execution not found or already completed")
        
        return {
            "execution_id": execution_id,
            "status": "cancelled",
            "message": "Workflow execution cancelled successfully"
        }
        
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Failed to cancel workflow: {str(e)}")


@router.get("/workflows/metrics", summary="Get workflow engine metrics")
async def get_workflow_metrics(
    tenant_id: str = Depends(get_current_tenant),
) -> Dict[str, Any]:
    """Get workflow engine performance metrics."""
    try:
        metrics = await autonomous_workflow_engine.get_metrics()
        return metrics
        
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Failed to get metrics: {str(e)}")


# Agent Communication Endpoints
@router.post("/agents/register", summary="Register an agent")
async def register_agent(
    agent_id: str = Query(..., description="Agent ID"),
    agent_type: str = Query(..., description="Agent type"),
    capabilities: List[str] = Query(..., description="Agent capabilities"),
    tenant_id: str = Depends(get_current_tenant),
) -> Dict[str, Any]:
    """Register an agent with the communication system."""
    try:
        await agent_communication_manager.register_agent(
            agent_id=agent_id,
            agent_type=AgentType(agent_type),
            capabilities=capabilities
        )
        
        return {
            "agent_id": agent_id,
            "status": "registered",
            "message": "Agent registered successfully"
        }
        
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Failed to register agent: {str(e)}")


@router.post("/agents/unregister", summary="Unregister an agent")
async def unregister_agent(
    agent_id: str = Query(..., description="Agent ID"),
    tenant_id: str = Depends(get_current_tenant),
) -> Dict[str, Any]:
    """Unregister an agent from the communication system."""
    try:
        await agent_communication_manager.unregister_agent(agent_id)
        
        return {
            "agent_id": agent_id,
            "status": "unregistered",
            "message": "Agent unregistered successfully"
        }
        
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Failed to unregister agent: {str(e)}")


@router.post("/agents/message", summary="Send a message to an agent")
async def send_agent_message(
    request: AgentMessageRequest,
    tenant_id: str = Depends(get_current_tenant),
) -> Dict[str, Any]:
    """Send a message to a specific agent."""
    try:
        message = AgentMessage(
            id=str(uuid.uuid4()),
            sender="api",
            recipient=request.recipient,
            message_type=MessageType(request.message_type),
            payload=request.payload,
            priority=MessagePriority(request.priority),
            correlation_id=request.correlation_id
        )
        
        success = await agent_communication_manager.send_message(message)
        
        if not success:
            raise HTTPException(status_code=400, detail="Failed to send message")
        
        return {
            "message_id": message.id,
            "recipient": message.recipient,
            "status": "sent",
            "timestamp": message.timestamp.isoformat()
        }
        
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Failed to send message: {str(e)}")


@router.get("/agents/{agent_id}/messages", summary="Receive messages for an agent")
async def receive_agent_messages(
    agent_id: str,
    timeout: float = Query(1.0, description="Timeout in seconds"),
    tenant_id: str = Depends(get_current_tenant),
) -> Dict[str, Any]:
    """Receive messages for a specific agent."""
    try:
        message = await agent_communication_manager.receive_message(agent_id, timeout)
        
        if not message:
            return {
                "agent_id": agent_id,
                "message": None,
                "status": "no_messages"
            }
        
        return {
            "agent_id": agent_id,
            "message": {
                "id": message.id,
                "sender": message.sender,
                "message_type": message.message_type.value,
                "payload": message.payload,
                "priority": message.priority.value,
                "timestamp": message.timestamp.isoformat(),
                "correlation_id": message.correlation_id
            },
            "status": "received"
        }
        
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Failed to receive message: {str(e)}")


@router.get("/agents/status", summary="Get agent communication status")
async def get_agent_communication_status(
    tenant_id: str = Depends(get_current_tenant),
) -> Dict[str, Any]:
    """Get the status of the agent communication system."""
    try:
        status = await agent_communication_manager.get_status()
        return status
        
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Failed to get status: {str(e)}")


# Enhanced Reasoning Endpoints
@router.post("/reasoning", summary="Perform enhanced reasoning")
async def perform_reasoning(
    request: ReasoningRequest,
    tenant_id: str = Depends(get_current_tenant),
) -> Dict[str, Any]:
    """Perform enhanced reasoning using various methods."""
    try:
        # Convert method string to enum
        method_map = {
            "chain_of_thought": ReasoningMethod.CHAIN_OF_THOUGHT,
            "tree_of_thoughts": ReasoningMethod.TREE_OF_THOUGHTS,
            "multi_step": ReasoningMethod.MULTI_STEP,
            "parallel": ReasoningMethod.PARALLEL,
            "hybrid": ReasoningMethod.HYBRID
        }
        
        method = method_map.get(request.method, ReasoningMethod.CHAIN_OF_THOUGHT)
        
        # Prepare context
        context = request.context or {}
        context["tenant_id"] = str(tenant_id)
        context["query"] = request.query
        
        # Perform reasoning
        result = await enhanced_reasoning_engine.reason(
            query=request.query,
            context=context,
            method=method,
            max_steps=request.max_steps
        )
        
        return {
            "chain_id": result.chain_id,
            "method": result.method.value,
            "final_answer": result.final_answer,
            "confidence": result.confidence,
            "reasoning_steps": result.reasoning_steps,
            "validation_metrics": result.validation_metrics,
            "execution_time": result.execution_time,
            "metadata": result.metadata
        }
        
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Reasoning failed: {str(e)}")


@router.get("/reasoning/stats", summary="Get reasoning statistics")
async def get_reasoning_stats(
    tenant_id: str = Depends(get_current_tenant),
) -> Dict[str, Any]:
    """Get statistics about reasoning performance."""
    try:
        stats = await enhanced_reasoning_engine.get_reasoning_stats()
        return stats
        
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Failed to get reasoning stats: {str(e)}")


# Combined Week 5 Features Endpoint
@router.get("/week5/status", summary="Get Week 5 features status")
async def get_week5_status(
    tenant_id: str = Depends(get_current_tenant),
) -> Dict[str, Any]:
    """Get comprehensive status of all Week 5 features."""
    try:
        # Get status from all Week 5 components
        workflow_metrics = await autonomous_workflow_engine.get_metrics()
        agent_status = await agent_communication_manager.get_status()
        reasoning_stats = await enhanced_reasoning_engine.get_reasoning_stats()
        
        return {
            "workflow_engine": {
                "status": "active",
                "metrics": workflow_metrics
            },
            "agent_communication": {
                "status": "active" if agent_status["running"] else "inactive",
                "metrics": agent_status
            },
            "enhanced_reasoning": {
                "status": "active",
                "stats": reasoning_stats
            },
            "overall_status": "operational"
        }
        
    except Exception as e:
        return {
            "workflow_engine": {"status": "error", "error": str(e)},
            "agent_communication": {"status": "error", "error": str(e)},
            "enhanced_reasoning": {"status": "error", "error": str(e)},
            "overall_status": "error"
        }


# Health check endpoint for Week 5 features
@router.get("/week5/health", summary="Health check for Week 5 features")
async def week5_health_check(
    tenant_id: str = Depends(get_current_tenant),
) -> Dict[str, Any]:
    """Health check for Week 5 features."""
    try:
        # Basic health checks
        workflow_metrics = await autonomous_workflow_engine.get_metrics()
        agent_status = await agent_communication_manager.get_status()
        
        # Check if components are responding
        workflow_healthy = workflow_metrics.get("total_executions", 0) >= 0
        agent_healthy = agent_status.get("running", False)
        
        overall_healthy = workflow_healthy and agent_healthy
        
        return {
            "status": "healthy" if overall_healthy else "unhealthy",
            "components": {
                "workflow_engine": "healthy" if workflow_healthy else "unhealthy",
                "agent_communication": "healthy" if agent_healthy else "unhealthy",
                "enhanced_reasoning": "healthy"  # Always healthy if we can reach this point
            },
            "timestamp": datetime.utcnow().isoformat()
        }
        
    except Exception as e:
        return {
            "status": "unhealthy",
            "error": str(e),
            "timestamp": datetime.utcnow().isoformat()
        }