import asyncio
import time
from datetime import datetime
import logging
import json
import traceback
from typing import Dict, Any, List, Optional
from sqlalchemy.orm import Session
from sqlalchemy.orm.exc import ObjectDeletedError, StaleDataError

from app.db.session import SessionLocal, async_db_op
from app.db.models.agent import AgentInstance, AgentTemplate
from app.db.models.session import Message
from app.core._regex import TURN_THINKING_MARKER, STRATEGY_BOILERPLATE

logger = logging.getLogger(__name__)

class AgentExecutor:
    """
    Orchestrates the execution of an agent, including its rework loop,
    evaluation via a Co-Worker Auditor, and real-time metric tracking.
    """

    def __init__(self, agent_id: str, services, user_service):
        self.agent_id = agent_id
        self.services = services
        self.user_service = user_service
        self.db: Session = SessionLocal()
        self.instance: Optional[AgentInstance] = None
        self.template: Optional[AgentTemplate] = None
        self.evaluator: Optional[HarnessEvaluator] = None

    @staticmethod
    async def run(agent_id: str, prompt: str, services, user_service, skip_coworker: bool = False):
        """Entry point for the background execution task."""
        logger.info(f"[AgentExecutor] Starting background run for agent {agent_id}")
        executor = AgentExecutor(agent_id, services, user_service)
        try:
            result = await executor.execute(prompt, skip_coworker)
            logger.info(f"[AgentExecutor] Completed background run for agent {agent_id}")
            return result
        finally:
            executor.db.close()

    async def execute(self, prompt: str, skip_coworker: bool = False):
        """Main execution sequence."""
        print(f"[AgentExecutor] 🔥 Executing with prompt: {prompt}")
        if not await self._initialize_instance(prompt):
            return

        heartbeat_task = asyncio.create_task(self._heartbeat_loop())
        
        try:
            # Phase 1: Setup & Initialization
            rubric_task = await self._setup_evaluation(prompt, skip_coworker)
            
            # Phase 2: Execution & Rework Loop
            final_result = await self._run_rework_loop(prompt, rubric_task)
            
            # Phase 3: Finalization
            await self._finalize_execution()
            return final_result

        except Exception as e:
            logger.error(f"[AgentExecutor] Unhandled loop error: {e}\n{traceback.format_exc()}")
            return await self._handle_fatal_error(e)
        finally:
            heartbeat_task.cancel()
            self.db.close()

    async def _safe_commit(self) -> bool:
        """Commits current DB changes with robust error handling for concurrent deletions."""
        try:
            await async_db_op(self.db.commit)
            return True
        except (ObjectDeletedError, StaleDataError):
            logger.info(f"[AgentExecutor] Agent {self.agent_id} deleted or modified externally. Exiting.")
            await async_db_op(self.db.rollback)
            return False
        except Exception as e:
            logger.error(f"[AgentExecutor] Commit failed: {e}")
            await async_db_op(self.db.rollback)
            raise

    async def _initialize_instance(self, prompt: str) -> bool:
        """Loads and prepares the agent instance for execution."""
        self.instance = self.db.query(AgentInstance).filter(AgentInstance.id == self.agent_id).first()
        if not self.instance or not prompt:
            return False
            
        self.mesh_node_id = self.instance.mesh_node_id
        self.session_id = self.instance.session_id
        self.latest_quality_score = self.instance.latest_quality_score or 0

        self.template = self.db.query(AgentTemplate).filter(AgentTemplate.id == self.instance.template_id).first()
        if not self.template:
            self.instance.status = "error_suspended"
            self.instance.last_error = f"Template '{self.instance.template_id}' not found."
            await self._safe_commit()
            return False

        # Initialize base metrics and status
        self.instance.last_heartbeat = datetime.utcnow()
        self.instance.status = "active"
        self.instance.last_error = None
        self.instance.total_runs = (self.instance.total_runs or 0) + 1
        
        # Area 4.2: Optional history clearing
        session = self.instance.session
        if session and getattr(session, "auto_clear_history", False):
            self.db.query(Message).filter(Message.session_id == session.id).delete(synchronize_session=False)

        return await self._safe_commit()

    async def _heartbeat_loop(self):
        """Maintains the 'active' lease in the background."""
        while True:
            await asyncio.sleep(60)
            try:
                # We use a nested session to avoid interfering with the main thread's flushing
                async with asyncio.Lock(): # Simple guard
                    inner_db = SessionLocal()
                    try:
                        obj = inner_db.query(AgentInstance).filter(AgentInstance.id == self.agent_id).first()
                        if not obj or obj.status not in ["active", "starting"]:
                            break
                        obj.last_heartbeat = datetime.utcnow()
                        inner_db.commit()
                    finally:
                        inner_db.close()
            except Exception:
                break

    async def _setup_evaluation(self, prompt: str, skip_coworker: bool) -> Optional[asyncio.Task]:
        """Configures the Co-Worker Auditor and launches rubric generation if enabled."""
        if skip_coworker or not getattr(self.template, "co_worker_quality_gate", False):
            return None

        # Resolve provider for evaluation
        provider = self._resolve_eval_provider()
        workspace_id = self.instance.session.sync_workspace_id if self.instance.session else str(self.instance.session_id)
        
        from app.core.orchestration.harness_evaluator import HarnessEvaluator
        self.evaluator = HarnessEvaluator(self.db, self.agent_id, self.instance.mesh_node_id, workspace_id, provider, self.services)
        
        self.instance.status = "starting"
        self.instance.evaluation_status = "📋 Co-Worker: Initiating parallel rubric & mission setup..."
        self.instance.current_rework_attempt = 0
        await self._safe_commit()

        # Deterministic setup of the .cortex state directory
        await self.evaluator.initialize_cortex()

        # Parallel rubric generation task
        mesh_node_id = str(self.instance.mesh_node_id)
        return asyncio.create_task(self._rubric_generator_bg(prompt, provider, workspace_id, mesh_node_id))

    def _resolve_eval_provider(self):
        """Determines the LLM provider configuration for the Auditor using PreferenceService."""
        session = self.instance.session
        provider_name = getattr(session, "provider_name", None)
        model_name = getattr(session, "model_name", None)
        
        provider, _ = self.services.preference_service.resolve_llm_provider(
            self.db, 
            session.user if session else None,
            provider_name,
            model_name
        )
        return provider

    async def _rubric_generator_bg(self, prompt: str, provider, workspace_id, mesh_node_id: str) -> Optional[str]:
        """Background helper for non-blocking rubric generation."""
        bg_db = SessionLocal()
        from app.core.orchestration.harness_evaluator import HarnessEvaluator
        try:
            bg_evaluator = HarnessEvaluator(bg_db, self.agent_id, mesh_node_id, workspace_id, provider, self.services)
            return await bg_evaluator.generate_rubric(prompt)
        except Exception as e:
            logger.error(f"[AgentExecutor] Rubric generation failed: {e}")
            return None
        finally:
            bg_db.close()

    async def _run_rework_loop(self, prompt: str, rubric_task: Optional[asyncio.Task]) -> Dict[str, Any]:
        """Core loop handling primary execution and Auditor-driven reworks."""
        max_reworks = getattr(self.template, "max_rework_attempts", 3)
        rework_threshold = getattr(self.template, "rework_threshold", 80)
        
        current_prompt = prompt
        current_attempt = 0
        rubric_content = ""
        loop_start_time = time.time()
        
        while current_attempt <= max_reworks:
            round_metrics = {"start": time.time(), "sub_events": [], "input_tokens": 0, "output_tokens": 0, "tool_counts": {}}
            
            try:
                # 1. Main Agent Execution (Streaming)
                result = await self._execute_main_agent(current_prompt, current_attempt, round_metrics)
                if not result: break

                # 2. Evaluation Phase
                if self.evaluator and result.get("response"):
                    # Ensure rubric is ready
                    if rubric_task and not rubric_content:
                        rubric_content = await rubric_task or "# Evaluation Rubric\nTask completion required."
                        rubric_task = None
                    
                    passed, audit_summary = await self._perform_quality_audit(
                        current_prompt, rubric_content, result, current_attempt, rework_threshold, round_metrics
                    )
                    
                    if passed: break
                    
                    # Plan Rework
                    current_prompt = f"### CO-WORKER DIRECTIVE (ATTEMPT {current_attempt + 1})\n\n{audit_summary}\n\nProceed with rework."
                    current_attempt += 1
                else:
                    break # No auditor or no response to evaluate

            except Exception as e:
                logger.error(f"[AgentExecutor] Round {current_attempt} failed: {e}\n{traceback.format_exc()}")
                return await self._handle_round_error(round_metrics, e)

        # If we exited the loop without a break (pass), we hit the limit
        score = getattr(self, 'latest_quality_score', 0) or 0
        threshold = getattr(self.template, 'rework_threshold', 80) or 80
        if score < threshold and getattr(self.template, "co_worker_quality_gate", False):
            try:
                self.instance.evaluation_status = f"🚫 failed_limit ({score}%)"
                await self._safe_commit()
            except ObjectDeletedError:
                logger.info(f"[AgentExecutor] Agent {self.agent_id} deleted before updating limit status.")

        return result

    async def _execute_main_agent(self, prompt: str, attempt: int, metrics: Dict) -> Optional[Dict]:
        """Orchestrates the LLM stream for the primary task."""
        self.instance.last_reasoning = ""
        self.instance.status = "active"
        self.instance.evaluation_status = f"🤖 Main Agent (Rd {attempt + 1}): Executing..."
        if not await self._safe_commit(): return None

        final_answer, sync_buffer = "", ""
        last_sync, last_msg_id = time.time(), None
        registry = getattr(self.services.rag_service, "node_registry_service", None)

        async for event in self.services.rag_service.chat_with_rag(
            db=self.db,
            session_id=self.instance.session_id,
            prompt=prompt,
            provider_name=self.instance.session.provider_name if self.instance.session else None,
            load_faiss_retriever=False,
            user_service=self.user_service
        ):
            e_type = event.get("type")
            if e_type == "finish":
                last_msg_id = event.get("message_id")
                metrics["tool_counts"] = event.get("tool_counts", {})
                final_answer = event.get("full_answer", "")
            
            # Extract common event processing to helper
            final_answer, sync_buffer, last_sync = await self._process_stream_event(
                event, final_answer, sync_buffer, last_sync, metrics, registry
            )
            if self.instance.status == "error_suspended": # Break if safe_commit failed in helper
                return None

        # Final persistence refresh
        await self._flush_execution_metrics(sync_buffer, metrics)
        metrics["duration"] = time.time() - metrics["start"]
        metrics["sub_events"].append({"name": "Agent execution", "duration": round(metrics["duration"], 2), "timestamp": time.time()})
        metrics["last_msg_id"] = last_msg_id

        return {
            "response": final_answer,
            "reasoning": self._compress_reasoning(self.instance.last_reasoning or "")
        }

    async def _process_stream_event(self, event: Dict, final_answer: str, sync_buffer: str, last_sync: float, metrics: Dict, registry) -> (str, str, float):
        """Processes individual stream events and handles real-time persistence/broadcasts."""
        e_type = event.get("type")
        
        if e_type == "token_counted":
            usage = event.get("usage", {})
            metrics["input_tokens"] += usage.get("prompt_tokens", 0)
            metrics["output_tokens"] += usage.get("completion_tokens", 0)
            return final_answer, sync_buffer, last_sync

        if e_type in ("reasoning", "content"):
            content = event.get("content", "")
            if e_type == "content": final_answer += content
            sync_buffer += content
            
            # Periodic buffer flush for UI responsiveness
            if time.time() - last_sync > 2.0 or len(sync_buffer) > 200:
                self.instance.last_reasoning = (self.instance.last_reasoning or "") + sync_buffer
                if not await self._safe_commit(): 
                    self.instance.status = "error_suspended"
                sync_buffer, last_sync = "", time.time()

            if registry and self.instance.mesh_node_id:
                registry.emit(self.instance.mesh_node_id, "reasoning", {
                    "content": content, "agent_id": self.agent_id, "session_id": self.instance.session_id
                })

        return final_answer, sync_buffer, last_sync

    async def _flush_execution_metrics(self, sync_buffer: str, metrics: Dict):
        """Finalizes and persists agent execution statistics."""
        self.instance.last_reasoning = (self.instance.last_reasoning or "") + sync_buffer
        self.instance.total_input_tokens = (self.instance.total_input_tokens or 0) + metrics.get("input_tokens", 0)
        self.instance.total_output_tokens = (self.instance.total_output_tokens or 0) + metrics.get("output_tokens", 0)
        self._merge_tool_counts(metrics.get("tool_counts", {}))
        await self._safe_commit()

    async def _perform_quality_audit(self, prompt, rubric, result, attempt, threshold, metrics) -> (bool, str):
        """Runs the Auditor's blind rating and delta analysis."""
        self.instance.evaluation_status = f"🕵️ Co-Worker (Rd {attempt + 1}): Auditing result..."
        await self._safe_commit()
        
        # Auditor technical context
        available_tools = []
        if hasattr(self.services, "tool_service"):
            tools = self.services.tool_service.get_available_tools(self.db, self.instance.user_id, session_id=self.instance.session_id)
            available_tools = [t["function"]["name"] for t in tools]
        
        partner_ctx = {"system_prompt": self.template.system_prompt_content, "skills": available_tools}
        
        # 1. Blind Rating & History Fetching in parallel for efficiency
        blind_task = asyncio.create_task(self.evaluator.evaluate_blind(prompt, rubric, result["response"], partner_context=partner_ctx))
        hist_task = asyncio.create_task(self._fetch_tester_history())
        
        blind_eval = await blind_task
        score = blind_eval.get("score", 0)
        just_msg = blind_eval.get("justification", "")
        metrics["sub_events"].append({"name": "Co-Worker review", "duration": round(blind_eval.get("duration", 0), 2), "timestamp": time.time()})

        self.latest_quality_score = score
        try:
            self.instance.latest_quality_score = score
            await self._safe_commit()
        except ObjectDeletedError:
            logger.info(f"[AgentExecutor] Agent {self.agent_id} deleted before updating score.")

        if score >= threshold:
            await self._record_audit_passed(score, just_msg, rubric, metrics, attempt)
            return True, just_msg

        # 2. Delta Analysis (Directive for rework)
        self.instance.evaluation_status = f"🧠 Co-Worker (Rd {attempt + 1}): Analyzing delta..."
        await self._safe_commit()

        # Rubric and justification are ready, and history should be ready too
        hist_log = await hist_task
        
        delta_start = time.time()
        if attempt >= 2: # Compaction trigger
            directive = await self.evaluator.generate_compaction_summary(prompt, hist_log)
            metrics["sub_events"].append({"name": "Context compaction", "duration": round(time.time() - delta_start, 2), "timestamp": time.time()})
        else:
            directive = await self.evaluator.evaluate_delta(prompt, rubric, just_msg, hist_log, result["reasoning"], partner_context=partner_ctx)
            metrics["sub_events"].append({"name": "Delta analysis", "duration": round(time.time() - delta_start, 2), "timestamp": time.time()})

        await self._record_audit_failed(score, just_msg, directive, rubric, metrics, attempt, hist_log)
        return False, directive

    async def _fetch_tester_history(self) -> List:
        """Retrieves raw history from the Auditor's workspace logs."""
        try:
            if not self.evaluator or not self.evaluator.assistant: return []
            res = await self.evaluator.assistant.adispatch_single(self.instance.mesh_node_id, "cat .cortex/history.log", session_id=self.evaluator.sync_workspace_id)
            return json.loads(res.get("stdout", "[]"))
        except: return []

    async def _record_audit_passed(self, score, justification, rubric, metrics, attempt):
        """Records a successful quality gate pass in history and DB."""
        self.instance.evaluation_status = f"✅ PASSED (Score {score}%)"
        await self._safe_commit()
        
        feedback = f"# Evaluation Passed\n\n**Score**: {score}/100\n\n**Justification**:\n{justification}"
        if self.evaluator and self.evaluator.assistant:
            await self.evaluator.assistant.awrite(self.mesh_node_id, ".cortex/feedback.md", feedback, session_id=self.evaluator.sync_workspace_id)
        
        duration = sum(e.get("duration", 0) for e in metrics["sub_events"])
        summary = self._truncate_text(justification, 250)
        await self.evaluator.log_round(attempt + 1, score, summary, "Final answer passed quality gate.", sub_events=metrics["sub_events"], duration=duration)
        
        await self._update_message_metadata(metrics.get("last_msg_id"), rubric, feedback, score, passed=True)

    async def _record_audit_failed(self, score, justification, directive, rubric, metrics, attempt, history):
        """Records a quality gate failure and triggers rework protocol."""
        full_audit = f"# Co-Worker Review (Attempt {attempt + 1})\n\n**Justification**:\n{justification}\n\n---\n\n{directive}"
        if self.evaluator and self.evaluator.assistant:
            await self.evaluator.assistant.awrite(self.mesh_node_id, ".cortex/feedback.md", full_audit, session_id=self.evaluator.sync_workspace_id)
        
        duration = sum(e.get("duration", 0) for e in metrics["sub_events"])
        summary = self._truncate_text(justification, 250)
        await self.evaluator.log_round(attempt + 1, score, summary, directive, sub_events=metrics["sub_events"], duration=duration)
        
        self.db.add(Message(session_id=self.session_id, sender="system", content=f"⚠️ **Co-Worker**: Quality check FAILED ({score}/100). Requesting rework..."))
        await self._update_message_metadata(metrics.get("last_msg_id"), rubric, full_audit, score, passed=False, history=history)
        
        try:
            self.instance.evaluation_status = f"⚠️ Rework Triggered ({score}%)"
            await self._safe_commit()
        except ObjectDeletedError:
            logger.info(f"[AgentExecutor] Agent {self.agent_id} deleted before updating status.")

    async def _update_message_metadata(self, msg_id, rubric, feedback, score, passed, history=None):
        """Enriches the assistant message with deep evaluation metadata."""
        if not msg_id: return
        self.db.query(Message).filter(Message.id == msg_id).update({
            "message_metadata": {
                "evaluation": {
                    "rubric": rubric, "feedback": feedback, "history": history or [],
                    "score": score, "passed": passed
                }
            }
        })
        await self._safe_commit()

    def _merge_tool_counts(self, round_counts: Dict):
        """Merges round-level tool metrics into the global agent stats."""
        if not round_counts: return
        counts = (self.instance.tool_call_counts or {}).copy()
        for tool, results in round_counts.items():
            if tool not in counts: counts[tool] = {"calls": 0, "successes": 0, "failures": 0}
            
            # Standardize count structure if it was legacy int
            if isinstance(counts[tool], int): 
                counts[tool] = {"calls": counts[tool], "successes": counts[tool], "failures": 0}
            
            results_dict = results if isinstance(results, dict) else {"calls": results, "successes": results, "failures": 0}
            
            counts[tool]["calls"] += results_dict.get("calls", 0)
            counts[tool]["successes"] += results_dict.get("successes", 0)
            counts[tool]["failures"] += results_dict.get("failures", 0)
        
        self.instance.tool_call_counts = counts
        from sqlalchemy.orm.attributes import flag_modified
        flag_modified(self.instance, "tool_call_counts")

    async def _finalize_execution(self):
        """Performs final state cleanup, metric aggregation, and logging."""
        try:
            if self.instance.status == "active":
                self.instance.status = "idle"
                score = getattr(self, 'latest_quality_score', 0) or 0
                threshold = getattr(self.template, 'rework_threshold', 80) or 80
                if score >= threshold or not getattr(self.template, "co_worker_quality_gate", False):
                    self.instance.successful_runs = (self.instance.successful_runs or 0) + 1

            self.instance.last_reasoning = None
            await self._safe_commit()
            
            if self.evaluator:
                await self.evaluator.log_event("Process Completed", "Lifecycle finished successfully.")
        except Exception as e:
            logger.info(f"[AgentExecutor] Error during finalization (likely deleted): {e}")

    async def _handle_round_error(self, metrics, error):
        """Gracefully handles errors within a specific rework round."""
        self.instance.status = "error_suspended"
        self.instance.last_error = str(error)
        self.instance.total_input_tokens = (self.instance.total_input_tokens or 0) + metrics["input_tokens"]
        self.instance.total_output_tokens = (self.instance.total_output_tokens or 0) + metrics["output_tokens"]
        await self._safe_commit()
        return {"status": "error", "response": f"Execution failed: {str(error)}", "reasoning": ""}

    async def _handle_fatal_error(self, error):
        """Gracefully handles unhandled exceptions across the entire lifecycle."""
        self.instance = self.db.query(AgentInstance).filter(AgentInstance.id == self.agent_id).first()
        if self.instance:
            self.instance.status = "error_suspended"
            self.instance.last_error = f"Unhandled fatal error: {str(error)}"
            await self._safe_commit()
        return {"status": "error", "response": "Internal server error during execution."}

    @staticmethod
    def _truncate_text(text: str, length: int) -> str:
        """Splits text by double newline to get first paragraph or truncates to length."""
        summary = text.split('\n\n')[0] if '\n\n' in text else text
        return f"{summary[:length-3]}..." if len(summary) > length else summary

    @staticmethod
    def _compress_reasoning(text: str) -> str:
        """Deduplicates turn markers and collapses boilerplate using centralized patterns."""
        if not text: return ""

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

        for line in lines:
            l_strip = line.strip()
            if not l_strip:
                cleaned.append(line)
                continue
            
            # Use central patterns to detect boilerplate without inline regex
            if TURN_THINKING_MARKER.search(line) and TURN_THINKING_MARKER.search(last_line):
                continue
            if STRATEGY_BOILERPLATE.search(line) and STRATEGY_BOILERPLATE.search(last_line):
                continue
                
            cleaned.append(line)
            last_line = l_strip

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