from __future__ import annotations

import asyncio
import json
import time
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any

import httpx

from nexus.core.events.bus import EventBus
from nexus.core.graph.dynamic_graph import DynamicGraph
from nexus.core.graph.loop_detective import LoopDetective
from nexus.core.graph.node import NodeConfig
from nexus.core.state.fluid_mesh import FluidMesh
from nexus.economy.accountant import TokenAccountant
from nexus.economy.budget_allocator import BudgetAllocator
from nexus.economy.deduplicator import Deduplicator
from nexus.observability.audit import AuditLog
from nexus.observability.tracer import NexusTracer
from nexus.persona.role_assignor import AgentProfile, RoleAssignor
from nexus.protocols.providers.base import LLMRequest, Message
from nexus.protocols.providers.local_provider import LocalProvider
from nexus.protocols.tools.registry import ToolRegistry
from nexus.resilience.checkpoint import CheckpointManager
from nexus.resilience.circuit_breaker import CircuitBreaker
from nexus.resilience.retry import SmartRetry


@dataclass
class SandboxEvent:
    stage: str
    node_id: str
    agent_id: str
    status: str
    timestamp: datetime
    message: str
    partial_answer: str = ""
    agent_output: str = ""
    input_tokens: int = 0
    output_tokens: int = 0
    cumulative_tokens: int = 0
    hot_tokens: int = 0
    warm_entries: int = 0
    cold_count: int = 0
    tool_calls: int = 0
    checkpoint_count: int = 0
    graph_nodes: int = 0
    completed_nodes: int = 0
    ready_nodes: int = 0
    mutation_count: int = 0
    circuit_state: str = "closed"
    planner_used: int = 0
    researcher_used: int = 0
    synthesizer_used: int = 0
    duration_ms: float = 0.0
    meta: dict[str, Any] = field(default_factory=dict)


class NexusSandboxWorkflow:
    """End-to-end query workflow using NEXUS components for live sandbox execution."""

    def __init__(self, base_url: str = "http://localhost:8080", model: str = "qwen2.5-3b-instruct-q4_k_m.gguf") -> None:
        normalized = base_url.rstrip("/")
        if not normalized.endswith("/v1"):
            normalized = f"{normalized}/v1"

        self.provider = LocalProvider(
            model=model,
            base_url=normalized,
            backend="llama-server",
            context_window_size=8192,
        )
        self.bus = EventBus()
        self.mesh = FluidMesh(self.bus, cold_persist_dir=".nexus/cold_store")
        self.accountant = TokenAccountant(self.bus)
        self.allocator = BudgetAllocator(session_budget=250_000)
        self.deduplicator = Deduplicator()
        self.tools = ToolRegistry()
        self.checkpoints = CheckpointManager(self.bus)
        self.tracer = NexusTracer(service_name="nexus-sandbox")
        self.audit = AuditLog(session_id="sandbox")
        self.retry = SmartRetry()
        self.loop_detective = LoopDetective()
        self.graph = DynamicGraph(self.bus, self.loop_detective)
        self.role_assignor = RoleAssignor()
        self._breakers: dict[str, CircuitBreaker] = {}
        self._register_agents()
        self._register_tools()

    def _register_agents(self) -> None:
        self.role_assignor.register_agent(
            AgentProfile(
                agent_id="planner",
                role="Task Planner",
                capabilities=["planning", "decomposition"],
                specialty_tags=["workflow", "graph"],
            )
        )
        self.role_assignor.register_agent(
            AgentProfile(
                agent_id="researcher",
                role="Web Research Analyst",
                capabilities=["research", "web_search"],
                specialty_tags=["web", "evidence"],
            )
        )
        self.role_assignor.register_agent(
            AgentProfile(
                agent_id="synthesizer",
                role="Final Response Writer",
                capabilities=["analysis", "synthesis", "writing"],
                specialty_tags=["summary", "final"],
            )
        )

    def _register_tools(self) -> None:
        self.tools.register(self.web_search, name="web_search", description="Search the web and return concise snippets.")
        self.tools.register(self.fetch_url_text, name="fetch_url_text", description="Fetch readable text from a URL.")

    async def web_search(self, query: str) -> dict[str, Any]:
        """Web search using DuckDuckGo instant answer plus related topics."""
        async with httpx.AsyncClient(timeout=30.0) as client:
            resp = await client.get(
                "https://api.duckduckgo.com/",
                params={"q": query, "format": "json", "no_html": 1, "skip_disambig": 1},
            )
            resp.raise_for_status()
            data = resp.json()

        snippets: list[str] = []
        if data.get("AbstractText"):
            snippets.append(data["AbstractText"])

        for topic in data.get("RelatedTopics", [])[:8]:
            if isinstance(topic, dict) and topic.get("Text"):
                snippets.append(topic["Text"])
            elif isinstance(topic, dict) and topic.get("Topics"):
                for sub in topic.get("Topics", [])[:3]:
                    if sub.get("Text"):
                        snippets.append(sub["Text"])

        return {
            "query": query,
            "heading": data.get("Heading", ""),
            "abstract_url": data.get("AbstractURL", ""),
            "snippets": snippets[:12],
        }

    async def fetch_url_text(self, url: str) -> str:
        """Fetch URL and return stripped text preview."""
        async with httpx.AsyncClient(timeout=30.0, follow_redirects=True) as client:
            resp = await client.get(url)
            resp.raise_for_status()
            text = resp.text
        compact = " ".join(text.split())
        return compact[:4000]

    async def _llm(self, prompt: str, system: str | None = None, max_tokens: int = 700) -> str:
        req = LLMRequest(
            messages=[Message(role="user", content=prompt)],
            system=system,
            max_tokens=max_tokens,
            temperature=0.3,
            tools=self.tools.to_nexus_format(["web_search", "fetch_url_text"]),
        )
        response = await self.provider.complete(req)
        return response.content

    async def run(self, query: str):
        """Async generator that yields incremental workflow updates."""
        await self.bus.start()
        session_id = f"sandbox-{int(time.time())}"
        self.audit = AuditLog(session_id=session_id)
        total_tokens = 0
        tool_calls = 0
        checkpoint_count = 0
        mutation_count = 0
        role_decisions: list[dict[str, Any]] = []
        agent_budgets = self.allocator.allocate(
            [
                {
                    "agent_id": "planner",
                    "description": f"Plan the workflow for user query: {query}",
                    "estimated_tool_calls": 1,
                    "complexity": 1.0,
                },
                {
                    "agent_id": "researcher",
                    "description": f"Research using web search for query: {query}",
                    "estimated_tool_calls": 3,
                    "complexity": 1.2,
                },
                {
                    "agent_id": "synthesizer",
                    "description": f"Produce final response and rationale for query: {query}",
                    "estimated_tool_calls": 1,
                    "complexity": 1.1,
                },
            ]
        )
        for aid, budget in agent_budgets.agent_allocations.items():
            self.accountant.register_agent(aid, budget)

        # Build DAG
        self.graph = DynamicGraph(self.bus, self.loop_detective)
        self.graph.add_node(NodeConfig(node_id="plan", agent_id="planner", task_fn=lambda **_: None, task_description="Plan steps"))
        self.graph.add_node(NodeConfig(node_id="research", agent_id="researcher", task_fn=lambda **_: None, task_description="Search and gather"))
        self.graph.add_node(NodeConfig(node_id="synthesize", agent_id="synthesizer", task_fn=lambda **_: None, task_description="Compose final answer", tags=["terminal"]))
        self.graph.add_edge("plan", "research")
        self.graph.add_edge("research", "synthesize")

        start_ts = datetime.now(timezone.utc)
        yield SandboxEvent(
            stage="init",
            node_id="init",
            agent_id="system",
            status="running",
            timestamp=start_ts,
            message="Initialized NEXUS sandbox workflow.",
            agent_output="",
            cumulative_tokens=0,
            hot_tokens=0,
            warm_entries=0,
            cold_count=self.mesh.cold.count(),
            tool_calls=0,
            checkpoint_count=checkpoint_count,
            graph_nodes=len(self.graph.all_nodes()),
            completed_nodes=0,
            ready_nodes=len(self.graph.get_ready_nodes(set())),
            mutation_count=mutation_count,
            circuit_state="closed",
            planner_used=0,
            researcher_used=0,
            synthesizer_used=0,
        )

        completed: set[str] = set()
        partial_answer = ""

        while not self.graph.is_complete():
            ready = self.graph.get_ready_nodes(completed)
            if not ready:
                break

            for node in ready:
                node.mark_active()
                stage_start = time.perf_counter()

                if node.config.node_id == "plan":
                    task_tags = ["workflow", "graph"]
                    req_caps = ["planning"]
                elif node.config.node_id == "research":
                    task_tags = ["web", "evidence"]
                    req_caps = ["research"]
                else:
                    task_tags = ["final", "summary"]
                    req_caps = ["synthesis"]

                assigned_agent = self.role_assignor.assign(task_tags=task_tags, required_capabilities=req_caps)
                role_decisions.append(
                    {
                        "node_id": node.config.node_id,
                        "requested": node.config.agent_id,
                        "assigned": assigned_agent,
                        "task_tags": task_tags,
                        "required_capabilities": req_caps,
                    }
                )

                if assigned_agent not in self._breakers:
                    self._breakers[assigned_agent] = CircuitBreaker(assigned_agent, self.bus)
                breaker = self._breakers[assigned_agent]

                current_circuit_state = breaker.state.value

                if node.config.node_id == "plan":
                    planner_prompt = (
                        "You are a planner. Return a compact JSON with keys: objective, steps (array), reasoning. "
                        f"User query: {query}"
                    )
                    plan_text = await self.retry.execute(lambda: breaker.call(self._llm, planner_prompt, "Planner agent", 500))
                    in_tok = self.provider.count_tokens(planner_prompt)
                    out_tok = self.provider.count_tokens(plan_text)
                    total_tokens += in_tok + out_tok
                    await self.accountant.record_call(assigned_agent, in_tok, out_tok)
                    await self.mesh.set(assigned_agent, "plan", plan_text, tags=["plan"])
                    partial_answer = "Plan completed. Starting research..."
                    message = "Planner produced execution strategy."
                    agent_output = plan_text
                    meta = {"plan": plan_text[:1200]}

                    # Runtime graph mutation signal: add verification node for complex plans.
                    if "steps" in plan_text.lower() and "verify" not in {n.config.node_id for n in self.graph.all_nodes()}:
                        verify_cfg = NodeConfig(
                            node_id="verify",
                            agent_id="researcher",
                            task_fn=lambda **_: None,
                            task_description="Cross-check and verify synthesis assumptions",
                            depends_on=["research"],
                            tags=["verification"],
                        )
                        self.graph.add_node(verify_cfg)
                        self.graph.add_edge("research", "verify")
                        synth = self.graph.node("synthesize")
                        if "verify" not in synth.config.depends_on:
                            synth.config.depends_on.append("verify")
                        mutation_count += 1
                        meta["graph_mutation"] = "Added verify node due to complex plan"

                elif node.config.node_id == "research":
                    sr = await self.tools.call("web_search", {"query": query})
                    tool_calls += 1
                    snippets = "\n".join(sr.get("snippets", [])[:8])
                    research_prompt = (
                        "Using these search snippets, produce structured research notes with: findings, unknowns, confidence.\n"
                        f"Query: {query}\nSnippets:\n{snippets}"
                    )
                    research_text = await self.retry.execute(
                        lambda: breaker.call(self._llm, research_prompt, "Research agent", 700)
                    )
                    in_tok = self.provider.count_tokens(research_prompt)
                    out_tok = self.provider.count_tokens(research_text)
                    total_tokens += in_tok + out_tok
                    await self.accountant.record_call(assigned_agent, in_tok, out_tok)
                    await self.mesh.set(assigned_agent, "search_results", sr, tags=["web", "search"])
                    await self.mesh.set(assigned_agent, "research_notes", research_text, tags=["analysis"])
                    partial_answer = "Research completed. Synthesizing final answer..."
                    message = "Researcher gathered external evidence and notes."
                    agent_output = research_text
                    meta = {"research": research_text[:1600], "search": sr}

                elif node.config.node_id == "verify":
                    notes = self.mesh.get("researcher", "research_notes") or self.mesh.get(assigned_agent, "research_notes") or ""
                    verify_prompt = (
                        "You are a verification agent. Validate claims from research notes and list possible weak points.\n"
                        f"User query: {query}\nResearch notes:\n{notes}\n"
                    )
                    verify_text = await self.retry.execute(
                        lambda: breaker.call(self._llm, verify_prompt, "Verification agent", 550)
                    )
                    in_tok = self.provider.count_tokens(verify_prompt)
                    out_tok = self.provider.count_tokens(verify_text)
                    total_tokens += in_tok + out_tok
                    await self.accountant.record_call(assigned_agent, in_tok, out_tok)
                    await self.mesh.set(assigned_agent, "verification_notes", verify_text, tags=["verification"])
                    partial_answer = "Verification complete. Proceeding to synthesis..."
                    message = "Verifier audited research quality before synthesis."
                    agent_output = verify_text
                    meta = {"verification": verify_text[:1200]}

                else:
                    plan = self.mesh.get("planner", "plan") or self.mesh.get(assigned_agent, "plan") or ""
                    notes = self.mesh.get("researcher", "research_notes") or ""
                    verify_notes = self.mesh.get("researcher", "verification_notes") or ""
                    suppress, sim = await self.deduplicator.should_suppress(str(notes), self.mesh.get_hot(assigned_agent))
                    if not suppress:
                        await self.mesh.set(assigned_agent, "research_inbox", notes, tags=["inbox"])
                    final_prompt = (
                        "Answer the user clearly with bullet points, short rationale, and a confidence score 0-1.\n"
                        f"User query: {query}\nPlan:\n{plan}\nResearch:\n{notes}\nVerification:\n{verify_notes}\n"
                    )
                    final_text = await self.retry.execute(
                        lambda: breaker.call(self._llm, final_prompt, "Synthesis agent", 900)
                    )
                    in_tok = self.provider.count_tokens(final_prompt)
                    out_tok = self.provider.count_tokens(final_text)
                    total_tokens += in_tok + out_tok
                    await self.accountant.record_call(assigned_agent, in_tok, out_tok)
                    await self.mesh.set(assigned_agent, "final_answer", final_text, tags=["final"])
                    partial_answer = final_text
                    message = "Synthesizer completed final response."
                    agent_output = final_text
                    meta = {"dedup_similarity": sim, "suppressed": suppress}

                # Eviction/migration cycle to show mesh evolution over time
                await self.mesh.run_eviction_cycle(assigned_agent)

                node.mark_complete(partial_answer)
                completed.add(node.config.node_id)

                if len(completed) % 2 == 0:
                    await self.checkpoints.write(
                        session_id=session_id,
                        node_id=node.config.node_id,
                        agent_states={
                            "planner": self.mesh.snapshot("planner"),
                            "researcher": self.mesh.snapshot("researcher"),
                            "synthesizer": self.mesh.snapshot("synthesizer"),
                        },
                        graph_state=self.graph.to_dict(),
                        token_ledgers={
                            k: v.model_dump() for k, v in self.accountant.get_all_ledgers().items()
                        },
                        metadata={"query": query},
                    )
                    checkpoint_count += 1

                await self.audit.write(
                    event_type="node_complete",
                    payload={
                        "node": node.config.node_id,
                        "agent": node.config.agent_id,
                        "tokens": total_tokens,
                        "tool_calls": tool_calls,
                    },
                )

                warm_entries = len(self.mesh.get_warm(assigned_agent).to_dict())
                duration_ms = (time.perf_counter() - stage_start) * 1000.0
                ledgers = self.accountant.get_all_ledgers()
                planner_used = ledgers.get("planner").total_used if "planner" in ledgers else 0
                researcher_used = ledgers.get("researcher").total_used if "researcher" in ledgers else 0
                synthesizer_used = ledgers.get("synthesizer").total_used if "synthesizer" in ledgers else 0
                ready_after = len(self.graph.get_ready_nodes(completed))

                meta["pillars"] = {
                    "fsm": {
                        "hot_context_preview": self.mesh.get_hot(assigned_agent).to_context_string()[:500],
                        "hot_tokens": self.mesh.get_hot(assigned_agent).total_tokens(),
                        "warm_entries": warm_entries,
                        "cold_count": self.mesh.cold.count(),
                    },
                    "deg": {
                        "topological_order": self.graph.topological_order(),
                        "completed_nodes": sorted(list(completed)),
                        "ready_nodes": ready_after,
                        "mutation_count": mutation_count,
                    },
                    "tee": {
                        "session_total_tokens": self.accountant.session_total(),
                        "ledgers": {k: v.model_dump() for k, v in ledgers.items()},
                    },
                    "rf": {
                        "circuit_state": current_circuit_state,
                        "circuit_failures": breaker.failure_count,
                        "checkpoint_count": checkpoint_count,
                    },
                    "ape": {
                        "role_decisions": role_decisions[-6:],
                        "assigned_agent": assigned_agent,
                    },
                    "or": {
                        "audit_verified": self.audit.verify()[0],
                    },
                }

                yield SandboxEvent(
                    stage=node.config.node_id,
                    node_id=node.config.node_id,
                    agent_id=assigned_agent,
                    status="complete",
                    timestamp=datetime.now(timezone.utc),
                    message=message,
                    partial_answer=partial_answer,
                    agent_output=agent_output,
                    cumulative_tokens=total_tokens,
                    hot_tokens=self.mesh.get_hot(assigned_agent).total_tokens(),
                    warm_entries=warm_entries,
                    cold_count=self.mesh.cold.count(),
                    tool_calls=tool_calls,
                    checkpoint_count=checkpoint_count,
                    graph_nodes=len(self.graph.all_nodes()),
                    completed_nodes=len(completed),
                    ready_nodes=ready_after,
                    mutation_count=mutation_count,
                    circuit_state=current_circuit_state,
                    planner_used=planner_used,
                    researcher_used=researcher_used,
                    synthesizer_used=synthesizer_used,
                    duration_ms=duration_ms,
                    meta=meta,
                )

        yield SandboxEvent(
            stage="done",
            node_id="done",
            agent_id="system",
            status="complete",
            timestamp=datetime.now(timezone.utc),
            message="Workflow complete.",
            partial_answer=self.mesh.get("synthesizer", "final_answer") or partial_answer,
            agent_output=self.mesh.get("synthesizer", "final_answer") or partial_answer,
            cumulative_tokens=total_tokens,
            hot_tokens=self.mesh.get_hot("synthesizer").total_tokens(),
            warm_entries=len(self.mesh.get_warm("synthesizer").to_dict()),
            cold_count=self.mesh.cold.count(),
            tool_calls=tool_calls,
            checkpoint_count=checkpoint_count,
            graph_nodes=len(self.graph.all_nodes()),
            completed_nodes=len(completed),
            ready_nodes=len(self.graph.get_ready_nodes(completed)),
            mutation_count=mutation_count,
            circuit_state="closed",
            planner_used=self.accountant.get_all_ledgers().get("planner").total_used if "planner" in self.accountant.get_all_ledgers() else 0,
            researcher_used=self.accountant.get_all_ledgers().get("researcher").total_used if "researcher" in self.accountant.get_all_ledgers() else 0,
            synthesizer_used=self.accountant.get_all_ledgers().get("synthesizer").total_used if "synthesizer" in self.accountant.get_all_ledgers() else 0,
            duration_ms=(datetime.now(timezone.utc) - start_ts).total_seconds() * 1000.0,
            meta={
                "token_report": self.accountant.to_report(),
                "pillars": {
                    "deg": {
                        "final_topological_order": self.graph.topological_order(),
                        "mutation_count": mutation_count,
                    },
                    "tee": {
                        "ledgers": {k: v.model_dump() for k, v in self.accountant.get_all_ledgers().items()},
                    },
                    "ape": {
                        "role_decisions": role_decisions,
                    },
                    "or": {
                        "audit_verify": self.audit.verify(),
                    },
                },
            },
        )

        await self.bus.stop()
