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, services, user_service, skip_coworker: bool = False):
        """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.last_error = None
            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"
                instance.last_error = f"Template '{instance.template_id}' not found."
                db.commit()
                return

            # Configuration for Rework Loop
            co_worker_enabled = getattr(template, "co_worker_quality_gate", False)
            rework_threshold = getattr(template, "rework_threshold", 80)
            max_rework_attempts = getattr(template, "max_rework_attempts", 3)
            
            # --- Phase 1: Pre-Execution Initialization (Harness Mirror) ---
            from app.core.orchestration.harness_evaluator import HarnessEvaluator
            evaluator = None
            rubric_content = ""
            
            if co_worker_enabled and not skip_coworker:
                from app.core.providers.factory import get_llm_provider
                # For Evaluation, we use the same provider/model as the main task for consistency
                # Load provider settings
                from app.db.models.session import Session as SessionModel
                agent_session = db.query(SessionModel).filter(SessionModel.id == instance.session_id).first()
                provider_name = getattr(agent_session, "provider_name", None)
                if not provider_name and user_service:
                    from app.config import settings
                    provider_name = settings.ACTIVE_LLM_PROVIDER
                
                # Resolve model/key from user preferences (mirrors rag.py logic)
                base_provider_key = provider_name.split("/")[0] if provider_name and "/" in provider_name else provider_name
                llm_prefs = {}
                if agent_session and agent_session.user and agent_session.user.preferences:
                    llm_prefs = agent_session.user.preferences.get("llm", {}).get("providers", {}).get(base_provider_key, {})
                if (not llm_prefs or not llm_prefs.get("api_key") or "*" in str(llm_prefs.get("api_key"))) and user_service:
                    system_prefs = user_service.get_system_settings(db)
                    system_provider_prefs = system_prefs.get("llm", {}).get("providers", {}).get(base_provider_key, {})
                    
                    if not system_provider_prefs or not system_provider_prefs.get("model"):
                       active_prov_key = system_prefs.get("llm", {}).get("active_provider")
                       if active_prov_key:
                           system_provider_prefs = system_prefs.get("llm", {}).get("providers", {}).get(active_prov_key, {})
                           if system_provider_prefs:
                               provider_name = active_prov_key
                    
                    if system_provider_prefs:
                        merged = system_provider_prefs.copy()
                        if llm_prefs: merged.update({k: v for k, v in llm_prefs.items() if v})
                        llm_prefs = merged
                eval_api_key = llm_prefs.get("api_key")
                eval_model = "" if "/" in (provider_name or "") else llm_prefs.get("model", "")
                eval_kwargs = {k: v for k, v in llm_prefs.items() if k not in ["api_key", "model"]}
                eval_provider = get_llm_provider(provider_name, model_name=eval_model, api_key_override=eval_api_key, **eval_kwargs)

                evaluator = HarnessEvaluator(db, agent_id, instance.mesh_node_id, instance.session.sync_workspace_id if instance.session else str(instance.session_id), eval_provider, services)
                await evaluator.initialize_cortex()
                # Round 0: Rubric Generation timing
                rubric_start = time.time()
                rubric_content = await evaluator.generate_rubric(prompt)
                rubric_duration = time.time() - rubric_start
                
                if not rubric_content:
                    rubric_content = "# Evaluation Rubric\nComplete the requested task with high technical accuracy."
                
                # Update status
                db.query(AgentInstance).filter(AgentInstance.id == agent_id).update({
                    "status": "starting",
                    "evaluation_status": "📋 Co-Worker: Generating request-specific rubric.md...",
                    "current_rework_attempt": 0
                })
                db.commit()
                
                # Emit status if registry exists
                registry = getattr(services.rag_service, "node_registry_service", None)
                if registry and instance.mesh_node_id:
                    registry.emit(instance.mesh_node_id, "status_update", {"evaluation_status": "📋 Co-Worker: Generating rubric..."})
                
                # Record initial timeline event for Rubric
                await evaluator.log_event("Rubric Generation", "Task-specific evaluation criteria established.", duration=rubric_duration)


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

            from app.db.models.session import Message
            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 system default
            if not provider_name and user_service:
                from app.config import settings
                provider_name = settings.ACTIVE_LLM_PROVIDER
            
            # Area 4.2: Hippocampus (Scratchpad) Idempotency Check
            if session_id:
                if getattr(agent_session, "auto_clear_history", False):
                    db.query(Message).filter(Message.session_id == session_id).delete(synchronize_session=False)
                    db.commit()

            current_prompt = prompt
            current_attempt = 0
            final_result = None

            # --- MAIN REWORK LOOP ---
            loop_start = time.time() # Handle scope for exception reporting
            while current_attempt <= max_rework_attempts:
                # Refresh instance for loop state
                instance = db.query(AgentInstance).filter(AgentInstance.id == agent_id).first()
                if not instance: break
                
                round_sub_events = []
                try:
                    registry = getattr(services.rag_service, "node_registry_service", None)
                    final_tool_counts = {}
                    final_input_tokens = 0
                    final_output_tokens = 0
                    final_answer = ""
                    
                    instance = db.query(AgentInstance).filter(AgentInstance.id == agent_id).first()
                    instance.last_reasoning = ""
                    instance.evaluation_status = f"🤖 Main Agent (Rd {current_attempt + 1}): Executing..."
                    db.commit()

                    execution_start = time.time()

                    if registry and instance.mesh_node_id:
                        registry.emit(instance.mesh_node_id, "status_update", {"evaluation_status": instance.evaluation_status})

                    # Buffers for real-time streaming to avoid O(N^2) regex and DB hammering
                    content_buffer = ""
                    last_db_sync_time = time.time()
                    sync_token_count = 0

                    async for event in services.rag_service.chat_with_rag(
                        db=db,
                        session_id=session_id,
                        prompt=current_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)
                            final_answer = event.get("full_answer", "")
                        elif event.get("type") == "token_counted":
                            usage = event.get("usage", {})
                            final_input_tokens += usage.get("prompt_tokens", 0)
                            final_output_tokens += usage.get("completion_tokens", 0)
                        elif event.get("type") in ("reasoning", "content"):
                            new_content = event.get("content", "")
                            if event.get("type") == "content":
                                final_answer += new_content
                            content_buffer += new_content
                            sync_token_count += 1
                            
                            now = time.time()
                            if now - last_db_sync_time > 2.0 or sync_token_count >= 50:
                                instance = db.query(AgentInstance).filter(AgentInstance.id == agent_id).first()
                                instance.last_reasoning = (instance.last_reasoning or "") + content_buffer
                                content_buffer = ""
                                last_db_sync_time = now
                                sync_token_count = 0
                                db.commit()

                            if registry and instance.mesh_node_id:
                                registry.emit(instance.mesh_node_id, "reasoning", {
                                    "content": new_content, 
                                    "agent_id": agent_id,
                                    "session_id": instance.session_id
                                })
                    
                    # Final flush
                    if content_buffer:
                        instance = db.query(AgentInstance).filter(AgentInstance.id == agent_id).first()
                        instance.last_reasoning = (instance.last_reasoning or "") + content_buffer
                        db.commit()
                        content_buffer = ""

                    exec_duration = time.time() - execution_start
                    round_sub_events.append({"name": "Agent execution", "duration": round(exec_duration, 2), "timestamp": time.time()})

                    # Execution complete
                    instance = db.query(AgentInstance).filter(AgentInstance.id == agent_id).first()
                    # 4.3: Post-processing to compress boilerplate from reasoning
                    final_reasoning = AgentExecutor._compress_reasoning(instance.last_reasoning or "")
                    
                    final_result = {
                        "response": final_answer,
                        "reasoning": final_reasoning
                    }

                    # --- EVALUATION PHASE (Co-Worker Loop) ---
                    if evaluator and final_answer:
                        instance.evaluation_status = "evaluating"
                        db.commit()
                        
                        # Stage 2A: Blind Rating
                        instance = db.query(AgentInstance).filter(AgentInstance.id == agent_id).first()
                        instance.evaluation_status = f"🕵️ Co-Worker (Rd {current_attempt + 1}): Auditing result against criteria..."
                        db.commit()
                        
                        if registry and instance.mesh_node_id:
                             registry.emit(instance.mesh_node_id, "status_update", {"evaluation_status": instance.evaluation_status})

                        blind_eval = await evaluator.evaluate_blind(prompt, rubric_content, final_answer)
                        score = blind_eval.get("score", 0)
                        justification = blind_eval.get("justification", "")
                        blind_duration = blind_eval.get("duration", 0)
                        round_sub_events.append({"name": "Co-Worker review", "duration": round(blind_duration, 2), "timestamp": time.time()})
                        
                        # Update instance with latest score
                        db.query(AgentInstance).filter(AgentInstance.id == agent_id).update({"latest_quality_score": score})
                        db.commit()

                        # Check Threshold
                        if score >= rework_threshold:
                            instance.evaluation_status = f"✅ PASSED (Score {score}%)"
                            db.commit()
                            
                            if registry and instance.mesh_node_id:
                                 registry.emit(instance.mesh_node_id, "status_update", {"evaluation_status": instance.evaluation_status})
                            
                            # Log final success feedback to workspace even if no rework was needed
                            success_feedback = f"# Evaluation Passed\n\n**Score**: {score}/100\n\n**Justification**:\n{justification}"
                            evaluator.assistant.write(evaluator.mesh_node_id, ".cortex/feedback.md", success_feedback, session_id=evaluator.sync_workspace_id)
                            
                            # M3: Aggregate total success duration and truncated summary for timeline
                            total_success_duration = sum(e.get("duration", 0) for e in round_sub_events)
                            summary_reason = justification.split('\n\n')[0] if '\n\n' in justification else justification
                            if len(summary_reason) > 250: summary_reason = summary_reason[:247] + "..."

                            await evaluator.log_round(current_attempt + 1, score, summary_reason, "Final answer passed quality gate.", sub_events=round_sub_events, duration=total_success_duration)
                            break # Success!
                        
                        # Check Rework Limits
                        if current_attempt >= max_rework_attempts:
                            instance.evaluation_status = "failed_limit"
                            instance.last_error = f"Co-Worker Gate: Quality fell below {rework_threshold}% after {max_rework_attempts} attempts."
                            db.commit()

                            # M3: Aggregate total failure duration and truncated summary for timeline
                            total_fail_duration = sum(e.get("duration", 0) for e in round_sub_events)
                            summary_reason = justification.split('\n\n')[0] if '\n\n' in justification else justification
                            if len(summary_reason) > 250: summary_reason = summary_reason[:247] + "..."

                            await evaluator.log_round(current_attempt + 1, score, summary_reason, "Failed quality gate after max attempts.", sub_events=round_sub_events, duration=total_fail_duration)
                            break # No more reworks
                        
                        # Stage Delta (Gap Analysis)
                        instance = db.query(AgentInstance).filter(AgentInstance.id == agent_id).first()
                        instance.evaluation_status = f"🧠 Co-Worker (Rd {current_attempt + 1}): Analyzing reasoning delta..."
                        db.commit()
                        
                        if registry and instance.mesh_node_id:
                             registry.emit(instance.mesh_node_id, "status_update", {"evaluation_status": instance.evaluation_status})
                        
                        # Fetch history for context
                        cmd_res = evaluator.assistant.dispatch_single(instance.mesh_node_id, "cat .cortex/history.log", session_id=evaluator.sync_workspace_id)
                        hist_log = []
                        try: hist_log = json.loads(cmd_res.get("stdout", "[]"))
                        except: pass

                        delta_start = time.time()
                        if current_attempt == 2:
                            directive_feedback = await evaluator.generate_compaction_summary(prompt, hist_log)
                            round_sub_events.append({"name": "Context compaction", "duration": round(time.time() - delta_start, 2), "timestamp": time.time()})
                        else:
                            directive_feedback = await evaluator.evaluate_delta(prompt, rubric_content, justification, hist_log, final_reasoning)
                            round_sub_events.append({"name": "Delta analysis", "duration": round(time.time() - delta_start, 2), "timestamp": time.time()})
                        
                        # M3: Categorization & Duration Metrics
                        full_audit_stream = f"# Co-Worker Review (Attempt {current_attempt + 1})\n\n**Justification**:\n{justification}\n\n---\n\n{directive_feedback}"
                        evaluator.assistant.write(evaluator.mesh_node_id, ".cortex/feedback.md", full_audit_stream, session_id=evaluator.sync_workspace_id)

                        # Extract high-density summary for timeline
                        summary_reason = justification.split('\n\n')[0] if '\n\n' in justification else justification
                        if len(summary_reason) > 250:
                            summary_reason = summary_reason[:247] + "..."
                        
                        # Calculate total round duration
                        total_round_duration = sum(e.get("duration", 0) for e in round_sub_events)
                        
                        # Log this round with summary and duration
                        await evaluator.log_round(current_attempt + 1, score, summary_reason, directive_feedback, sub_events=round_sub_events, duration=total_round_duration)

                        # Trigger next iteration
                        current_prompt = f"### CO-WORKER DIRECTIVE (ATTEMPT {current_attempt + 1})\n\n{directive_feedback}\n\nProceed with rework."
                        current_attempt += 1
                        
                        db.add(Message(session_id=session_id, sender="system", content=f"⚠️ **Co-Worker**: Quality check FAILED (Score: {score}/100). Requesting rework..."))
                        db.query(AgentInstance).filter(AgentInstance.id == agent_id).update({"evaluation_status": f"⚠️ Rework Triggered ({score}%)"})
                        db.commit()
                        continue # Start next loop iteration
                    else:
                        break # No co-worker or no answer

                except Exception as e:
                    import traceback
                    print(f"[AgentExecutor] RAG attempt failed for {agent_id}: {e}")
                    print(traceback.format_exc())
                    instance = db.query(AgentInstance).filter(AgentInstance.id == agent_id).first()
                    if instance:
                        instance.status = "error_suspended"
                        instance.last_error = str(e)
                        db.commit()
                    return {
                        "status": "error",
                        "response": f"Execution failed: {str(e)}",
                        "reasoning": instance.last_reasoning if instance else ""
                    }

            # Final loop cleanup & Stats
            instance = db.query(AgentInstance).filter(AgentInstance.id == agent_id).first()
            if instance and 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 if final_result else 0)
                instance.total_output_tokens = (instance.total_output_tokens or 0) + (final_output_tokens if final_result else 0)
                
                if final_tool_counts:
                    import copy
                    current_counts = copy.deepcopy(instance.tool_call_counts or {})
                    for k, v in final_tool_counts.items():
                        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"] += v.get("calls", 0)
                        current_counts[k]["successes"] += v.get("successes", 0)
                        current_counts[k]["failures"] += v.get("failures", 0)
                    instance.tool_call_counts = current_counts
                    from sqlalchemy.orm.attributes import flag_modified
                    flag_modified(instance, "tool_call_counts")
                
            # Clear reasoning as the task is now complete
            instance.last_reasoning = None
            db.commit()

            return final_result

        except Exception as e:
            import traceback
            print(f"[AgentExecutor] Unhandled loop error: {e}")
            print(traceback.format_exc())
            instance = db.query(AgentInstance).filter(AgentInstance.id == agent_id).first()
            if instance:
                instance.status = "error_suspended"
                instance.last_error = f"Unhandled loop error: {str(e)}"
                db.commit()
            return {
                "status": "error",
                "response": "Internal server error during execution.",
                "reasoning": instance.last_reasoning if instance else ""
            }
        finally:
            heartbeat_task.cancel()
            db.close()

    @staticmethod
    def _compress_reasoning(text: str) -> str:
        """Deduplicates turn markers and collapses boilerplate using high-perf string logic."""
        if not text: return ""

        lines = text.splitlines(keepends=True)
        cleaned = []
        last_important_line = ""

        # Collapse repeating Turn headers and Strategy boilerplate without regex
        for line in lines:
            l_strip = line.strip()
            if not l_strip:
                cleaned.append(line)
                continue
            
            l_lower = l_strip.lower()
            
            # Deduplicate the system thinking marker
            is_turn_marker = "🛰️" in line and "[turn" in l_lower and "thinking" in l_lower
            if is_turn_marker:
                if "🛰️" in last_important_line and "[turn" in last_important_line.lower():
                    continue
            
            # Collapse Strategy boilerplate
            is_strategy = "strategy:" in l_lower and "executing orchestrated tasks" in l_lower
            if is_strategy:
                if "strategy:" in last_important_line.lower() and "executing orchestrated tasks" in last_important_line.lower():
                    continue
            
            cleaned.append(line)
            last_important_line = l_strip

        return "".join(cleaned).strip()

