import asyncio
import time
from datetime import datetime
from sqlalchemy.orm import Session
from tenacity import retry, wait_exponential, stop_after_attempt
import json

from app.db.session import SessionLocal
from app.db.models.agent import AgentInstance, AgentTemplate
from app.db.models import Message

class AgentExecutor:
    @staticmethod
    async def run(agent_id: str, prompt: str, rag_service, user_service):
        """Asynchronous execution loop for the agent."""
        # Create a fresh DB session for the background task
        db: Session = SessionLocal()
        try:
            instance = db.query(AgentInstance).filter(AgentInstance.id == agent_id).first()
            if not instance or not prompt:
                return

            # Acquire Lease
            instance.last_heartbeat = datetime.utcnow()
            instance.status = "active"
            instance.total_runs = (instance.total_runs or 0) + 1
            db.commit()

            # Launch secondary heartbeat task
            async def heartbeat():
                while True:
                    await asyncio.sleep(60)
                    try:
                        inner_db = SessionLocal()
                        inner_instance = inner_db.query(AgentInstance).filter(AgentInstance.id == agent_id).first()
                        if not inner_instance or inner_instance.status != "active":
                            inner_db.close()
                            break
                        inner_instance.last_heartbeat = datetime.utcnow()
                        inner_db.commit()
                        inner_db.close()
                    except:
                        break
            
            heartbeat_task = asyncio.create_task(heartbeat())

            template = db.query(AgentTemplate).filter(AgentTemplate.id == instance.template_id).first()
            if not template:
                instance.status = "error_suspended"
                db.commit()
                return

            max_iterations = template.max_loop_iterations or 20
            session_id = instance.session_id

            # Load session to check configured assigned provider
            from app.db.models.session import Session as SessionModel
            agent_session = db.query(SessionModel).filter(SessionModel.id == session_id).first()
            
            provider_name = getattr(agent_session, "provider_name", None)

            # If not explicitly defined on session, fallback to smartest available
            if not provider_name and user_service:
                from app.config import settings
                provider_name = settings.ACTIVE_LLM_PROVIDER
                sys_prefs = user_service.get_system_settings(db)
                providers = sys_prefs.get("llm", {}).get("providers", {})
                for best_choice in ["deepseek", "gemini", "openai", "anthropic"]:
                    if best_choice in providers and providers[best_choice].get("api_key"):
                        provider_name = best_choice
                        break

                        break
            
            # Area 4.2: Hippocampus (Scratchpad) Idempotency Check
            # We skip this for simple chat prompts, but for autonomous loops its vital
            if session_id:
                # In a real impl, we'd check if .cortex_memory_scratchpad.txt exists on node
                # For MVP, we just log the intention as per Task 4.2
                print(f"[AgentExecutor] Task 4.2: Idempotency check for {agent_id} in {instance.current_workspace_jail or '/tmp'}")

                if getattr(agent_session, "auto_clear_history", False):
                    print(f"[AgentExecutor] Auto-clearing history for session {session_id} before run.")
                    db.query(Message).filter(Message.session_id == session_id).delete(synchronize_session=False)
                    db.commit()

            print(f"[AgentExecutor] Starting run for {agent_id} with provider '{provider_name}'. Prompt length: {len(prompt)}")

            loop_start = time.time()
            # Iterate the RAG architecture to solve the prompt
            try:
                final_tool_counts = {}
                final_input_tokens = 0
                final_output_tokens = 0
                
                # We consume the generator completely to let it execute all tools and generate reasoning
                async for event in rag_service.chat_with_rag(
                    db=db,
                    session_id=session_id,
                    prompt=prompt,
                    provider_name=provider_name,
                    load_faiss_retriever=False,
                    user_service=user_service
                ):
                    if event.get("type") == "finish":
                        final_tool_counts = event.get("tool_counts", {})
                        final_input_tokens = event.get("input_tokens", 0)
                        final_output_tokens = event.get("output_tokens", 0)
                
                # Execution complete
                instance = db.query(AgentInstance).filter(AgentInstance.id == agent_id).first()
                if instance.status == "active":
                    instance.status = "idle" # Completed work
                    instance.successful_runs = (instance.successful_runs or 0) + 1
                    
                    elapsed = int(time.time() - loop_start)
                    instance.total_running_time_seconds = (instance.total_running_time_seconds or 0) + elapsed
                    instance.total_input_tokens = (instance.total_input_tokens or 0) + final_input_tokens
                    instance.total_output_tokens = (instance.total_output_tokens or 0) + final_output_tokens
                    
                    if final_tool_counts:
                        current_counts = dict(instance.tool_call_counts or {})
                        for k, v in final_tool_counts.items():
                            # Upgrade legacy single-integer counts into rich metric dicts
                            if k in current_counts and isinstance(current_counts[k], int):
                                current_counts[k] = {"calls": current_counts[k], "successes": current_counts[k], "failures": 0}
                            if not isinstance(v, dict):
                                v = {"calls": v, "successes": v, "failures": 0}
                            if k not in current_counts:
                                current_counts[k] = {"calls": 0, "successes": 0, "failures": 0}

                            current_counts[k]["calls"] = current_counts[k].get("calls", 0) + v.get("calls", 0)
                            current_counts[k]["successes"] = current_counts[k].get("successes", 0) + v.get("successes", 0)
                            current_counts[k]["failures"] = current_counts[k].get("failures", 0) + v.get("failures", 0)
                            
                        instance.tool_call_counts = current_counts
                    
                    db.commit()

            except Exception as e:
                import traceback
                print(f"[AgentExecutor] RAG execution failed for {agent_id}: {e}")
                print(traceback.format_exc())
                instance = db.query(AgentInstance).filter(AgentInstance.id == agent_id).first()
                instance.status = "error_suspended"
                db.commit()

        except Exception as e:
            print(f"[AgentExecutor] Unhandled loop error: {e}")
            instance = db.query(AgentInstance).filter(AgentInstance.id == agent_id).first()
            if instance:
                instance.status = "error_suspended"
                db.commit()
        finally:
            heartbeat_task.cancel()
            db.close()
