import logging
import asyncio
import json
import time
from typing import Dict, Any, List, Optional, Tuple
import os
import re

from app.db.models.agent import AgentInstance, AgentTemplate
from app.db import models
from app.core.orchestration import Architect
from app.core._regex import ANSI_ESCAPE, FINAL_SCORE

logger = logging.getLogger(__name__)

class HarnessEvaluator:
    """
    Quality gate auditor that evaluates agent performance against dynamic rubrics.
    Refactored for structural clarity and modular evaluation flows.
    """
    def __init__(self, db, agent_id, mesh_node_id, sync_workspace_id, llm_provider, services):
        self.db = db
        self.agent_id = agent_id
        self.mesh_node_id = mesh_node_id
        self.sync_workspace_id = sync_workspace_id
        self.llm_provider = llm_provider
        self.services = services
        
        self.orchestrator = getattr(services, "orchestrator", None)
        self.assistant = self.orchestrator.assistant if self.orchestrator else None
        self.history_cache = None

    async def initialize_cortex(self):
        """Initializes the .cortex/ state directory on the mesh node."""
        if not self.assistant or not self.mesh_node_id: return
        # Parallelize initialization tasks for speed
        await asyncio.gather(
            self.assistant.adispatch_single(self.mesh_node_id, "mkdir -p .cortex", session_id=self.sync_workspace_id),
            self.assistant.awrite(self.mesh_node_id, ".cortex/history.log", "[]", session_id=self.sync_workspace_id),
            self.assistant.awrite(self.mesh_node_id, ".cortex/feedback.md", "# Session Started\n", session_id=self.sync_workspace_id)
        )

    async def log_event(self, name: str, details: str = "", duration: float = 0, event_type: str = "event", metadata: Dict = None):
        """Appends a lifecycle event to the history log."""
        history = await self._read_history()
        history.append({
            "type": event_type, "name": name, "details": details,
            "duration": round(duration, 2), "timestamp": time.time(), "metadata": metadata or {}
        })
        await self._write_history(history)

    async def _read_history(self) -> List[Dict]:
        """Safely reads history.log from the node."""
        if self.history_cache is not None:
            return self.history_cache
        if not self.assistant: return []
        try:
            res = await self.assistant.adispatch_single(self.mesh_node_id, "cat .cortex/history.log", session_id=self.sync_workspace_id, timeout=5)
            history = json.loads(res.get("stdout", "[]")) if res.get("status") == "SUCCESS" else []
            self.history_cache = history
            return history
        except: return []

    async def _write_history(self, history: List[Dict]):
        """Safely writes history.log back to the node."""
        self.history_cache = history
        if not self.assistant: return
        try:
            await self.assistant.awrite(self.mesh_node_id, ".cortex/history.log", json.dumps(history, indent=2), session_id=self.sync_workspace_id)
        except Exception as e:
            logger.error(f"[HarnessEvaluator] History write failed: {e}")

    async def ensure_coworker_ground_truth(self):
        """Bootstraps .coworker.md alignment doc if missing on node."""
        if not self.assistant or not self.mesh_node_id: return
        check = await self.assistant.adispatch_single(self.mesh_node_id, "ls .coworker.md", session_id=self.sync_workspace_id, timeout=5)
        if check.get("status") == "SUCCESS": return

        instr = self._get_agent_instructions()
        sys_p = "You are the Swarm Alchemist. Distill instructions into a high-density .coworker.md project edict file."
        try:
            prediction = await self.llm_provider.acompletion(messages=[{"role":"system","content":sys_p},{"role":"user","content":f"Instructions:\n{instr}"}], stream=False)
            await self.assistant.awrite(self.mesh_node_id, ".coworker.md", prediction.choices[0].message.content, session_id=self.sync_workspace_id)
        except Exception as e: logger.error(f"Ground truth failed: {e}")

    def _get_agent_instructions(self) -> str:
        """Resolves the primary system instructions for the agent instance."""
        try:
            instance = self.db.query(AgentInstance).filter(AgentInstance.id == self.agent_id).first()
            if instance and instance.session and instance.session.system_prompt_override:
                return instance.session.system_prompt_override
            return instance.template.system_prompt_content if instance and instance.template else "Follow standard protocols."
        except: return "No instructions found."

    async def generate_rubric(self, initial_prompt: str) -> Optional[str]:
        """Generates a task-specific evaluation rubric."""
        if not self.assistant: return None
        await self.ensure_coworker_ground_truth()
        
        ctx = await self._read_node_file(".coworker.md")
        sys_p = f"You are a Quality Architect. Context:\n{ctx}\nGenerate a '# Evaluation Rubric' (0-100) for this task."
        try:
            res = await self.llm_provider.acompletion(messages=[{"role":"system","content":sys_p},{"role":"user","content":f"Task: {initial_prompt}"}], stream=False)
            content = res.choices[0].message.content
            await self.assistant.awrite(self.mesh_node_id, ".cortex/rubric.md", content, session_id=self.sync_workspace_id)
            return content
        except Exception as e:
            logger.error(f"Rubric fail: {e}")
            return None

    async def _read_node_file(self, path: str) -> str:
        """Helper to cat a file from the node."""
        try:
            res = await self.assistant.adispatch_single(self.mesh_node_id, f"cat {path}", session_id=self.sync_workspace_id, timeout=5)
            return res.get("stdout", "") if res.get("status") == "SUCCESS" else ""
        except: return ""

    async def evaluate_blind(self, prompt: str, rubric: str, result: str, partner_context: Dict = None) -> Dict[str, Any]:
        """Performs an objective blind audit of the final response."""
        if os.getenv("MOCK_EVALUATION", "false").lower() == "true":
            logger.info("[HarnessEvaluator] MOCK_EVALUATION enabled, returning mock score.")
            from app.db.models.agent import AgentInstance
            instance = self.db.query(AgentInstance).filter(AgentInstance.id == self.agent_id).first()
            threshold = instance.template.rework_threshold if instance and instance.template else 80
            
            # If threshold is high (usually 100 or 101 in tests to force failure), return failing score
            if threshold >= 100:
                return {"score": 0, "justification": "Mocked failure to force rework or limit check. FINAL_SCORE: 0"}
            else:
                return {"score": 100, "justification": "Mocked success. FINAL_SCORE: 100"}
                
        sys_p = f"You are the Blind Auditor. Request: {prompt}\nResult: {result}\nRubric: {rubric}\nRate 0-100. End with: FINAL_SCORE: [number]"
        return await self._run_evaluator_agent(sys_p, "Audit the result state.")

    async def evaluate_delta(self, prompt: str, rubric: str, blind_just: str, history: List[Dict], transcript: str, partner_context: Dict = None) -> str:
        """Analyzes the delta between reasoning and final output to generate rework directives."""
        hist_text = "\n".join([f"- {h['round']}: {h['reason']}" for h in history if h['type'] == 'attempt'])
        sys_p = f"You are the Delta Architect. Request: {prompt}\nRubric: {rubric}\nAuditor: {blind_just}\nHistory: {hist_text}\nExecution: {transcript}\nGenerate '# Rework Instructions' as actionable directives."
        try:
            res = await self.llm_provider.acompletion(messages=[{"role":"system","content":sys_p},{"role":"user","content":"Generate directives."}], stream=False)
            return res.choices[0].message.content
        except Exception as e: return f"Delta Error: {e}"

    async def _run_evaluator_agent(self, system_prompt: str, user_request: str) -> Dict[str, Any]:
        """Runs a sub-architect to perform complex multi-tool audit logic."""
        architect = Architect()
        final_answer = ""
        
        # Gathering context chunks
        chunks = [{"id": "hub_manifesto", "content": self._read_local_manifesto(), "metadata": {"priority":"critical"}}]
        coworker = await self._read_node_file(".coworker.md")
        if coworker: chunks.append({"id": ".coworker.md", "content": coworker, "metadata": {"priority":"high"}})
        
        async for event in architect.run(
            question=user_request, context_chunks=chunks, history=[], llm_provider=self.llm_provider,
            db=self.db, sync_workspace_id=self.sync_workspace_id, feature_name="agent_harness", session_override=system_prompt
        ):
            if event["type"] == "content":
                final_answer += event["content"]
                
        if self.assistant:
            await self.assistant.awrite(self.mesh_node_id, ".cortex/feedback.md", f"# Finalizing Audit...\n\n{final_answer}", session_id=self.sync_workspace_id)
        
        score_match = FINAL_SCORE.search(final_answer)
        return {"score": int(score_match.group(1)) if score_match else 0, "justification": final_answer}

    def _read_local_manifesto(self) -> str:
        """Reads the regional hub manifesto for architectural alignment."""
        try:
            with open("/app/blueprints/manifesto.md", "r") as f: return f.read()
        except: return ""

    async def generate_compaction_summary(self, prompt: str, history: List[Dict]) -> str:
        """Compacts multiple failure reports into a single dense rework directive."""
        fail_path = "\n".join([f"Attempt {h.get('round')}: {h.get('reason')}" for h in history if h['type'] == 'attempt'])
        sys_p = "You are a Quality Compactor. Distill failures into a '# COMPACTED DIRECTIVE'."
        try:
            prediction = await self.llm_provider.acompletion(messages=[{"role":"system","content":sys_p},{"role":"user","content":f"Mission: {prompt}\nHistory:\n{fail_path}"}], stream=False)
            return prediction.choices[0].message.content
        except: return "# COMPACTED DIRECTIVE\n- Stabilize core logic and address all remaining gaps."

    async def log_round(self, round_num: int, score: int, reason: str, feedback: str = "", sub_events: List[Dict] = None, duration: float = 0):
        """Records a completed evaluation round in history.log."""
        history = await self._read_history()
        history.append({
            "type": "attempt", "round": round_num, "score": score, "reason": reason, "feedback": feedback,
            "duration": round(duration, 2), "timestamp": time.time(), "sub_events": sub_events or []
        })
        await self._write_history(history)
