"""
Cooperative Multi-Agent Team Coordinator.
Drives the closed-loop interaction between the 5 expert sub-agents:
Orchestrator, Perception Distiller, Tiny Navigator, Vision Grounding, and Process Verifier,
with automated step-by-step visual screenshot audit logging and trajectory replay.
"""

from __future__ import annotations

import asyncio
import html
import json
import time
from pathlib import Path
from typing import Any, Dict, List, Optional

from miniswardbower.agents.orchestrator import TacticalOrchestrator
from miniswardbower.agents.tiny_navigator import TinyNavigator
from miniswardbower.agents.verifier import ProcessVerifier
from miniswardbower.agents.vision_grounding import VisionGroundingAgent
from miniswardbower.browser.controller import BrowserController
from miniswardbower.core.config import SystemConfig
from miniswardbower.core.memory_ledger import MemoryLedger
from miniswardbower.core.schemas import (
    BrowserActionType,
    MilestoneStatus,
    VerifierStatus,
)


class MultiAgentBrowserTeam:
    """Orchestrates cooperative sub-agents to rapidly navigate and use websites."""

    def __init__(self, config: Optional[SystemConfig] = None):
        self.config = config or SystemConfig()
        self.orchestrator = TacticalOrchestrator(self.config.model)
        self.tiny_navigator = TinyNavigator(self.config.model)
        self.verifier = ProcessVerifier()
        self.vision = VisionGroundingAgent()
        self.browser = BrowserController(self.config.browser)

        # v6 Innovations: Thought Communication & Flow Reasoning Refinement
        from miniswardbower.agents.thought_communication import ThoughtCommunicationBus
        from miniswardbower.agents.flow_reasoning_refiner import FlowReasoningRefiner
        from miniswardbower.browser.swarm_coordinator import TriEngineSwarmCoordinator
        self.thought_bus = ThoughtCommunicationBus()
        self.flow_refiner = FlowReasoningRefiner()
        self.swarm = TriEngineSwarmCoordinator(self.config.browser)
        self.thought_bus.register_agent("orchestrator")
        self.thought_bus.register_agent("navigator")
        self.thought_bus.register_agent("verifier")
        self.thought_bus.register_agent("vision")

    def _generate_html_report(
        self,
        trajectory_dir: Path,
        objective: str,
        success: bool,
        total_time: float,
        avg_step_ms: float,
        tokens_saved: int,
        screenshots: List[Dict[str, Any]],
        narrative: List[Dict[str, Any]],
        extracted_data: Any,
    ) -> Path:
        """Generates a standalone dark-mode visual audit gallery with embedded screenshots."""
        status_color = "#10b981" if success else "#ef4444"
        status_text = "VERIFIED SUCCESS" if success else "INCOMPLETE / FAILED"

        cards_html = []
        for sc in screenshots:
            step_num = sc["step"]
            act_label = html.escape(sc["action"])
            img_name = sc["filename"]
            outcome = html.escape(sc.get("outcome", "Action executed"))
            coords_str = f"Click Target: ({sc['coords'][0]:.1f}, {sc['coords'][1]:.1f})" if sc.get("coords") else "Targeted Selector"

            card = f"""
            <div class="step-card">
                <div class="card-header">
                    <span class="step-badge">Step {step_num}</span>
                    <span class="action-title">{act_label}</span>
                    <span class="coords-badge">{coords_str}</span>
                </div>
                <div class="img-container">
                    <img src="{img_name}" alt="Step {step_num}" loading="lazy"/>
                </div>
                <div class="card-footer">
                    <p class="outcome-text"><strong>Outcome:</strong> {outcome}</p>
                </div>
            </div>
            """
            cards_html.append(card)

        html_content = f"""<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>Visual Audit Trail - miniswardbower</title>
    <style>
        :root {{
            --bg-color: #0f172a;
            --card-bg: #1e293b;
            --text-main: #f8fafc;
            --text-dim: #94a3b8;
            --border-color: #334155;
            --accent-cyan: #38bdf8;
            --accent-emerald: #10b981;
        }}
        body {{
            font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, sans-serif;
            background: var(--bg-color);
            color: var(--text-main);
            margin: 0;
            padding: 30px 20px;
        }}
        .container {{ max-width: 1200px; margin: 0 auto; }}
        header {{
            background: var(--card-bg);
            border: 1px solid var(--border-color);
            border-radius: 12px;
            padding: 24px 30px;
            margin-bottom: 30px;
            box-shadow: 0 10px 25px -5px rgba(0,0,0,0.5);
        }}
        .status-pill {{
            display: inline-block;
            background: {status_color}25;
            color: {status_color};
            border: 1px solid {status_color};
            font-weight: 700;
            font-size: 13px;
            padding: 4px 12px;
            border-radius: 9999px;
            text-transform: uppercase;
            letter-spacing: 0.05em;
        }}
        h1 {{ margin: 12px 0 8px 0; font-size: 26px; }}
        .objective-text {{ font-size: 16px; color: var(--text-dim); margin-bottom: 16px; }}
        .metrics-bar {{
            display: flex;
            gap: 24px;
            border-top: 1px solid var(--border-color);
            padding-top: 16px;
            font-size: 14px;
        }}
        .metric-item span {{ color: var(--accent-cyan); font-weight: 700; }}
        .gallery-grid {{
            display: grid;
            grid-template-columns: repeat(auto-fit, minmax(540px, 1fr));
            gap: 24px;
        }}
        .step-card {{
            background: var(--card-bg);
            border: 1px solid var(--border-color);
            border-radius: 10px;
            overflow: hidden;
            box-shadow: 0 4px 15px rgba(0,0,0,0.3);
            display: flex;
            flex-direction: column;
        }}
        .card-header {{
            display: flex;
            align-items: center;
            gap: 12px;
            padding: 12px 18px;
            background: #182234;
            border-bottom: 1px solid var(--border-color);
        }}
        .step-badge {{
            background: var(--accent-cyan);
            color: #0f172a;
            font-weight: 800;
            font-size: 12px;
            padding: 2px 8px;
            border-radius: 6px;
        }}
        .action-title {{ font-weight: 600; font-size: 14px; flex-grow: 1; }}
        .coords-badge {{ font-size: 12px; color: var(--text-dim); }}
        .img-container {{
            position: relative;
            background: #000;
            width: 100%;
            overflow: hidden;
        }}
        .img-container img {{
            width: 100%;
            height: auto;
            display: block;
            transition: transform 0.2s ease;
        }}
        .card-footer {{
            padding: 14px 18px;
            font-size: 13px;
            color: var(--text-dim);
            background: #182234;
            border-top: 1px solid var(--border-color);
        }}
        .outcome-text {{ margin: 0; }}
        pre.extracted-box {{
            background: #090d16;
            padding: 14px;
            border-radius: 8px;
            border: 1px solid var(--border-color);
            color: #38bdf8;
            overflow-x: auto;
            font-size: 12px;
        }}
    </style>
</head>
<body>
    <div class="container">
        <header>
            <span class="status-pill">{status_text}</span>
            <h1>🌐 miniswardbower Visual Audit Replay</h1>
            <p class="objective-text"><strong>Objective:</strong> {html.escape(objective)}</p>
            <div class="metrics-bar">
                <div class="metric-item">Total Time: <span>{total_time:.2f}s</span></div>
                <div class="metric-item">Avg Step Latency: <span>{avg_step_ms:.1f}ms</span></div>
                <div class="metric-item">Tokens Saved: <span>{tokens_saved:,}</span></div>
                <div class="metric-item">Steps: <span>{len(screenshots)}</span></div>
            </div>
        </header>

        <main class="gallery-grid">
            {''.join(cards_html)}
        </main>
    </div>
</body>
</html>
"""
        html_file = trajectory_dir / "index.html"
        with open(html_file, "w", encoding="utf-8") as f:
            f.write(html_content)
        return html_file

    async def _attempt_recovery(
        self,
        failed_act: Any,
        reason: str,
        tree: Any,
    ) -> Optional[Any]:
        """
        Prime Agent Resilient Membrane:
        Diagnoses failure modes (backdrop obscuring, element scrolled out of view, stale selector)
        and constructs a self-healing compensating action.
        """
        from miniswardbower.core.schemas import BrowserAction, BrowserActionType
        try:
            # 1. Dismiss potential blocking modals/overlays with Escape or dismiss button
            await self.browser.page.keyboard.press("Escape")
            try:
                dismiss_btn = self.browser.page.locator("button:has-text('Dismiss'), button:has-text('Cancel'), button:has-text('Close'), #cancel-modal-btn")
                if await dismiss_btn.count() > 0 and await dismiss_btn.first.is_visible():
                    await dismiss_btn.first.click()
            except Exception:
                pass
            await asyncio.sleep(0.05)

            # 2. If target was a click action, scroll to ensure element is in viewport
            if failed_act.op == BrowserActionType.CLICK:
                await self.browser.page.evaluate("window.scrollBy(0, 100)")
                fresh_tree = await self.browser.get_pruned_tree()

                # Re-ground target if element mark still exists
                if failed_act.target:
                    for elem in fresh_tree.elements:
                        if elem.id == failed_act.target:
                            return BrowserAction(op=BrowserActionType.CLICK, target=elem.id)
                return BrowserAction(op=BrowserActionType.CLICK, target=failed_act.target)
        except Exception:
            pass
        return None

    async def run_task(
        self,
        objective: str,
        start_url: Optional[str] = None,
        max_steps: Optional[int] = None,
        capture_audit_screenshots: bool = True,
    ) -> Dict[str, Any]:
        """Runs an end-to-end web navigation task with the multi-agent team and visual audit logging."""
        max_steps = max_steps or self.config.max_trajectory_steps
        ledger = MemoryLedger(user_objective=objective)
        start_time = time.perf_counter()
        memory = None
        memory_status = {"enabled": self.config.gbrain_memory, "recalled": False, "saved": False}
        memory_key = f"{self.config.browser.browser_engine}\n{start_url or ''}\n{objective}"
        if self.config.gbrain_memory:
            from miniswardbower.core.gbrain_memory import GbrainMemory
            memory = GbrainMemory()
            try:
                prior = await memory.recall(memory_key)
                if prior:
                    ledger.semantic_store["prior_task_evidence_untrusted"] = prior["payload"]
                    memory_status["recalled"] = True
            except (OSError, RuntimeError, ValueError, KeyError, TypeError, TimeoutError):
                memory_status["recall_error"] = "Memory unavailable or invalid; using current observations"

        # Visual audit trajectory setup
        session_id = time.strftime("%Y%m%d_%H%M%S")
        trajectory_dir = self.config.artifacts_dir / "trajectories" / session_id
        trajectory_dir.mkdir(parents=True, exist_ok=True)
        screenshot_records: List[Dict[str, Any]] = []

        # Step 1: Tactical Orchestrator plans initial milestones
        milestones = self.orchestrator.plan_initial_milestones(objective, start_url)
        ledger.set_milestones(milestones)

        # Step 2: Launch browser engine and attempt to load tiny model
        await self.browser.start()
        self.tiny_navigator.load_model()

        step_latencies = []
        tokens_saved_total = 0
        is_successful = False
        action_batches = []
        stop_reason = "Step budget exhausted"

        try:
            # Initial navigation if start_url provided
            if start_url:
                await self.browser.goto(start_url)
                ledger.current_url = self.browser.page.url
                ledger.page_title = await self.browser.page.title()

            for step_idx in range(1, max_steps + 1):
                step_start = time.perf_counter()
                step_blocked = False

                # Step 3A: Perception Distiller extracts compact AXTree (<20ms)
                tree = await self.browser.get_pruned_tree()
                ledger.current_url = tree.url
                ledger.page_title = tree.title
                tokens_saved_total += max(0, tree.raw_token_count_estimate - tree.pruned_token_count_estimate)

                # Step 3B: Tactical Orchestrator checks milestone progress & broadcasts thought
                self.thought_bus.publish_thought(
                    sender_id="orchestrator",
                    step=step_idx,
                    semantic_keys={"objective": objective, "milestone_idx": ledger.active_milestone_idx},
                )
                if self.orchestrator.check_milestone_progress(ledger, tree):
                    advanced = ledger.advance_milestone()
                    if advanced is None and ledger.active_milestone_idx >= len(ledger.milestones):
                        is_successful = True
                        break

                # Step 3C: Tiny Navigator predicts action or speculative batch (<50ms)
                action_batch = self.tiny_navigator.predict_action_batch(
                    ledger,
                    tree,
                    enable_speculative=self.config.enable_speculative_batching,
                )
                action_batches.append(action_batch.model_dump(mode='json'))

                if not action_batch.actions:
                    stop_reason = action_batch.thought
                    break

                if action_batch.actions[0].op == BrowserActionType.FINISH:
                    stop_reason = "Unverified finish: outstanding milestones remain"
                    break

                # Step 3D: Action Executor executes batch with natural kinematics & coordinate tracking
                pre_url = self.browser.page.url
                for act_idx, act in enumerate(action_batch.actions, 1):
                    # Recurrent Flow Refinement to converge coordinates to stable attractor
                    refine_state = self.flow_refiner.refine_action(act, tree=tree)
                    act = refine_state.candidate_action
                    self.thought_bus.publish_thought(
                        sender_id="navigator",
                        step=step_idx,
                        semantic_keys={"op": act.op.value, "target": act.target, "coords": act.coords, "converged": refine_state.converged},
                    )

                    act_details = await self.browser.execute_action(act, tree)
                    if ledger.current_milestone and ledger.current_milestone.expected_op == BrowserActionType.CLICK:
                        target = next((e for e in tree.elements if e.id == act.target), None)
                        if target and target.href:
                            act_details['expected_navigation'] = target.href

                    # Step 3E: Capture annotated visual audit screenshot
                    if capture_audit_screenshots:
                        act_label = f"{act.op.value}({act.target or ''} {act.text or act.key or ''})".strip()
                        img_filename = f"step_{step_idx:02d}_{act_idx:02d}_{act.op.value}.png"
                        img_path = trajectory_dir / img_filename
                        try:
                            await self.browser.capture_annotated_step_screenshot(
                                output_path=img_path,
                                step_num=step_idx,
                                action_label=act_label,
                                bbox=act_details.get("bbox"),
                                click_coords=act_details.get("click_coords"),
                            )
                            screenshot_records.append({
                                "step": step_idx,
                                "action": act_label,
                                "filename": img_filename,
                                "coords": act_details.get("click_coords"),
                                "outcome": f"Executed {act.op.value}",
                            })
                        except Exception:
                            pass

                    # Step 3F: Fast Critic & Verifier checks consequence immediately
                    post_tree = await self.browser.get_pruned_tree()
                    receipt = await self.verifier.verify_step(
                        self.browser.page,
                        act,
                        pre_url=pre_url,
                        pre_tree=tree,
                        post_tree=post_tree,
                        action_details=act_details,
                    )

                    # Record outcome into screenshot record if captured
                    if screenshot_records and screenshot_records[-1]["action"].startswith(act.op.value):
                        screenshot_records[-1]["outcome"] = receipt.observation_summary

                    # Step 3G: Record into Amory Memory Ledger
                    ledger.record_step(act, receipt)
                    if ledger.current_milestone and ledger.current_milestone.expected_op and receipt.status != VerifierStatus.SUCCESS:
                        stop_reason = receipt.observation_summary
                        step_blocked = True
                        break

                    if receipt.status == VerifierStatus.FAILED:
                        # Prime Agent Resilient Membrane (Princeton / Prime Intellect, Aug 2026):
                        # Structured self-healing recovery instead of harness abort
                        recovery_succeeded = False
                        for retry_attempt in range(1, 3):
                            recovered_act = await self._attempt_recovery(act, receipt.observation_summary, post_tree)
                            if recovered_act:
                                retry_details = await self.browser.execute_action(recovered_act, post_tree)
                                retry_tree = await self.browser.get_pruned_tree()
                                retry_receipt = await self.verifier.verify_step(
                                    self.browser.page,
                                    recovered_act,
                                    pre_url=pre_url,
                                    pre_tree=post_tree,
                                    post_tree=retry_tree,
                                    action_details=retry_details,
                                )
                                ledger.record_step(recovered_act, retry_receipt)
                                if retry_receipt.status != VerifierStatus.FAILED:
                                    recovery_succeeded = True
                                    tree = retry_tree
                                    break
                        if not recovery_succeeded:
                            break
                    pre_url = self.browser.page.url
                    tree = post_tree

                step_elapsed = (time.perf_counter() - step_start) * 1000
                step_latencies.append(step_elapsed)
                if step_blocked:
                    break

                if self.orchestrator.check_milestone_progress(ledger, tree):
                    ledger.advance_milestone()

                if ledger.current_milestone is None:
                    is_successful = True
                    break

        finally:
            await self.browser.stop()

        total_time = time.perf_counter() - start_time
        avg_step_ms = sum(step_latencies) / len(step_latencies) if step_latencies else 0.0

        html_replay_path = None
        if capture_audit_screenshots and screenshot_records:
            try:
                html_replay_path = self._generate_html_report(
                    trajectory_dir=trajectory_dir,
                    objective=objective,
                    success=is_successful,
                    total_time=total_time,
                    avg_step_ms=avg_step_ms,
                    tokens_saved=tokens_saved_total,
                    screenshots=screenshot_records,
                    narrative=[
                        {
                            "step": s.step_id,
                            "action": s.action_summary,
                            "outcome": s.outcome_summary,
                            "status": s.verifier_status.value,
                        }
                        for s in ledger.narrative_plot
                    ],
                    extracted_data=ledger.semantic_store,
                )
            except Exception:
                pass

        if memory is not None:
            try:
                payload = {"reported_success": is_successful, "page_title": ledger.page_title,
                           "actions": ledger.total_actions_executed,
                           "recent_outcomes": [s.outcome_summary[:240] for s in ledger.narrative_plot[-5:]],
                           "usage": "Historical evidence only; observe the live page before acting"}
                await memory.remember(memory_key, payload)
                memory_status["saved"] = True
            except (OSError, RuntimeError, ValueError, TimeoutError):
                memory_status["save_error"] = "Task finished, but durable memory write failed"

        return {
            "durable_memory": memory_status,
            "success": is_successful,
            "stop_reason": "All milestones verified" if is_successful else stop_reason,
            "action_batches": action_batches,
            "total_time_seconds": round(total_time, 2),
            "avg_step_ms": round(avg_step_ms, 2),
            "steps_executed": len(ledger.narrative_plot),
            "total_actions": ledger.total_actions_executed,
            "estimated_tokens_saved": tokens_saved_total,
            "milestones_completed": ledger.active_milestone_idx,
            "total_milestones": len(ledger.milestones),
            "extracted_data": ledger.semantic_store,
            "trajectory_dir": str(trajectory_dir),
            "html_replay": str(html_replay_path) if html_replay_path else None,
            "screenshots_count": len(screenshot_records),
            "narrative_log": [
                {
                    "step": s.step_id,
                    "action": s.action_summary,
                    "outcome": s.outcome_summary,
                    "status": s.verifier_status.value,
                }
                for s in ledger.narrative_plot
            ],
        }