from __future__ import annotations

import asyncio
import inspect
import os
import time
from collections.abc import Callable
from datetime import datetime, timezone
from typing import Any
from uuid import uuid4

from pydantic import BaseModel, Field

from nexus.api.autowire import AutoWire
from nexus.api.builder import WorkflowDefinition
from nexus.core.events.bus import EventBus
from nexus.core.events.types import SessionEndEvent, SessionStartEvent, TaskCompleteEvent, TaskFailedEvent, TaskStartEvent
from nexus.core.graph.branch_manager import BranchManager
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.observability.audit import AuditLog
from nexus.observability.tracer import NexusTracer
from nexus.persona.engine import PersonaEngine
from nexus.persona.evolver import PromptEvolver
from nexus.persona.reflector import PostTaskReflector
from nexus.persona.role_assignor import RoleAssignor
from nexus.persona.skill_registry import SkillRegistry
from nexus.protocols.providers.anthropic_provider import AnthropicProvider
from nexus.protocols.providers.base import AbstractProvider
from nexus.protocols.providers.gemini_provider import GeminiProvider
from nexus.protocols.providers.local_provider import LocalProvider
from nexus.protocols.providers.openai_provider import OpenAIProvider
from nexus.protocols.tools.registry import ToolRegistry
from nexus.resilience.checkpoint import CheckpointManager
from nexus.resilience.circuit_breaker import CircuitBreaker
from nexus.resilience.errors import AgentFailure, OrchestratorError
from nexus.resilience.retry import SmartRetry


class NexusConfig(BaseModel):
    provider: str = "anthropic"
    model: str = "claude-opus-4-6"
    api_key: str | None = None
    session_budget: int = 500_000
    checkpoint_interval: int = 300
    audit_log: bool = False
    checkpoint_dir: str = ".nexus/checkpoints"
    cold_store_dir: str = ".nexus/cold_store"
    resilience: str = "standard"
    memory: str = "session"
    local_base_url: str = "http://localhost:11434/v1"
    local_backend: str = "ollama"


class NexusResult(BaseModel):
    output: Any
    session_id: str
    total_tokens: int
    duration_ms: float
    checkpoints: list[str] = Field(default_factory=list)
    audit_log_path: str | None = None
    token_report: str


class Nexus:
    """Main NEXUS orchestration entrypoint."""

    def __init__(self, config: NexusConfig | None = None) -> None:
        self.config = config or NexusConfig()
        self.bus = EventBus()
        self.mesh = FluidMesh(self.bus, cold_persist_dir=self.config.cold_store_dir)
        self.accountant = TokenAccountant(self.bus)
        self.budget_allocator = BudgetAllocator(session_budget=self.config.session_budget)
        self.checkpoint_manager = CheckpointManager(self.bus, checkpoint_dir=self.config.checkpoint_dir)
        self.loop_detective = LoopDetective()
        self.graph = DynamicGraph(self.bus, self.loop_detective)
        self.branch_manager = BranchManager()
        self.tools = ToolRegistry()
        self.provider = self._build_provider(self.config)
        self.reflector = PostTaskReflector(self.provider)
        self.evolver = PromptEvolver(self.provider)
        self.skill_registry = SkillRegistry(self.mesh.cold)
        self.role_assignor = RoleAssignor()
        self.persona = PersonaEngine(self.bus, self.reflector, self.evolver, self.skill_registry, self.role_assignor)
        self.tracer = NexusTracer()
        self.audit = AuditLog(session_id="init") if self.config.audit_log else None
        self._circuit_breakers: dict[str, CircuitBreaker] = {}

    def run(self, tasks: list[Callable] | WorkflowDefinition, inputs: dict | None = None) -> NexusResult:
        """Synchronous run wrapper."""
        return asyncio.run(self._run_async(tasks, inputs or {}))

    async def _run_async(self, tasks: list[Callable] | WorkflowDefinition, inputs: dict) -> NexusResult:
        """Execute workflow graph with resilience, accounting, and memory integration."""
        session_id = str(uuid4())
        started = time.perf_counter()
        checkpoints: list[str] = []
        completed: set[str] = set()

        if self.config.audit_log:
            self.audit = AuditLog(session_id=session_id)

        await self.bus.start()
        await self.bus.publish_and_wait(SessionStartEvent(source_id="nexus", session_id=session_id, config=self.config.model_dump()))
        await self.persona.start()

        task_specs = self._build_task_specs(tasks)
        plan = self.budget_allocator.allocate(task_specs)
        for agent_id, budget in plan.agent_allocations.items():
            self.accountant.register_agent(agent_id, budget)

        for spec in task_specs:
            node_cfg = NodeConfig(
                node_id=spec["task_id"],
                agent_id=spec["agent_id"],
                task_fn=spec["task_fn"],
                task_description=spec["description"],
                depends_on=spec["depends_on"],
                loop_strategy=spec.get("loop_strategy", "max_iterations"),
                max_iterations=int(spec.get("max_iterations", 5)),
                enables_graph_mutation=bool(spec.get("enables_graph_mutation", False)),
                tags=list(spec.get("tags", [])),
            )
            self.graph.add_node(node_cfg)

        for spec in task_specs:
            for dep in spec["depends_on"]:
                self.graph.add_edge(dep, spec["task_id"])

        last_checkpoint = time.perf_counter()
        retry = SmartRetry()

        while not self.graph.is_complete():
            ready = self.graph.get_ready_nodes(completed)
            if not ready:
                raise OrchestratorError("No ready nodes but graph is incomplete")

            async def run_node(node) -> None:
                node.mark_active()
                await self.bus.publish_and_wait(
                    TaskStartEvent(source_id="nexus", node_id=node.config.node_id, agent_id=node.config.agent_id, task_input=inputs)
                )
                if self.audit:
                    await self.audit.write("task_start", {"node_id": node.config.node_id, "agent_id": node.config.agent_id})

                breaker = self._circuit_breakers.get(node.config.agent_id)
                if breaker is None:
                    breaker = CircuitBreaker(node.config.agent_id, self.bus)
                    self._circuit_breakers[node.config.agent_id] = breaker

                async def invoke_task():
                    fn = node.config.task_fn
                    if fn is None:
                        return {}
                    sig = inspect.signature(fn)
                    kwargs = {k: v for k, v in inputs.items() if k in sig.parameters}
                    hot_context = self.mesh.get_hot(node.config.agent_id).to_context_string()
                    kwargs.setdefault("context", hot_context)
                    evolved_prompt = await self.persona.get_evolved_prompt(node.config.agent_id, "")
                    kwargs.setdefault("system_prompt", evolved_prompt)
                    examples = await self.skill_registry.to_few_shot_examples(node.config.task_description, n=2)
                    kwargs.setdefault("few_shot_examples", examples)
                    if inspect.iscoroutinefunction(fn):
                        return await fn(**kwargs)
                    return await asyncio.to_thread(fn, **kwargs)

                try:
                    result = await retry.execute(lambda: breaker.call(invoke_task))
                    content = str(result)
                    in_tok = self.provider.count_tokens(node.config.task_description)
                    out_tok = self.provider.count_tokens(content)
                    await self.accountant.record_call(node.config.agent_id, in_tok, out_tok)
                    await self.mesh.set(node.config.agent_id, f"result:{node.config.node_id}", result)
                    node.mark_complete(result)
                    completed.add(node.config.node_id)
                    await self.bus.publish_and_wait(
                        TaskCompleteEvent(
                            source_id="nexus",
                            node_id=node.config.node_id,
                            agent_id=node.config.agent_id,
                            result=result,
                            tokens_used=in_tok + out_tok,
                            duration_ms=0.0,
                        )
                    )
                    if self.audit:
                        await self.audit.write("task_complete", {"node_id": node.config.node_id, "agent_id": node.config.agent_id})
                except Exception as exc:
                    err = AgentFailure(str(exc), attempt=1, agent_id=node.config.agent_id, node_id=node.config.node_id)
                    node.mark_failed(err)
                    await self.bus.publish_and_wait(
                        TaskFailedEvent(
                            source_id="nexus",
                            node_id=node.config.node_id,
                            agent_id=node.config.agent_id,
                            error=err,
                            attempt=1,
                        )
                    )
                    raise

            await asyncio.gather(*(run_node(node) for node in ready))

            now = time.perf_counter()
            if (now - last_checkpoint) >= self.config.checkpoint_interval:
                ckpt_id = await self.checkpoint_manager.write(
                    session_id=session_id,
                    node_id=ready[-1].config.node_id,
                    agent_states={agent: self.mesh.snapshot(agent) for agent in plan.agent_allocations},
                    graph_state=self.graph.to_dict(),
                    token_ledgers={aid: ledger.model_dump() for aid, ledger in self.accountant.get_all_ledgers().items()},
                    metadata={"auto": True},
                )
                checkpoints.append(ckpt_id)
                last_checkpoint = now

        terminal = self.graph.terminal_node()
        output = terminal.result if terminal else None
        duration_ms = (time.perf_counter() - started) * 1000.0

        await self.bus.publish_and_wait(
            SessionEndEvent(
                source_id="nexus",
                session_id=session_id,
                total_tokens=self.accountant.session_total(),
                duration_ms=duration_ms,
                success=True,
            )
        )
        await self.bus.stop()

        return NexusResult(
            output=output,
            session_id=session_id,
            total_tokens=self.accountant.session_total(),
            duration_ms=duration_ms,
            checkpoints=checkpoints,
            audit_log_path=str(self.audit.log_path) if self.audit else None,
            token_report=self.accountant.to_report(),
        )

    def _build_provider(self, config: NexusConfig) -> AbstractProvider:
        """Construct provider implementation from config."""
        p = config.provider.lower()
        if p == "anthropic":
            return AnthropicProvider(model=config.model, api_key=config.api_key)
        if p == "openai":
            return OpenAIProvider(model=config.model, api_key=config.api_key)
        if p == "gemini":
            return GeminiProvider(model=config.model, api_key=config.api_key)
        if p == "local":
            return LocalProvider(
                model=config.model,
                base_url=config.local_base_url,
                api_key=config.api_key or "not-required",
                backend=config.local_backend,
                context_window_size=8192,
            )
        raise OrchestratorError(f"Unsupported provider: {config.provider}")

    def _build_task_specs(self, tasks: list[Callable] | WorkflowDefinition) -> list[dict[str, Any]]:
        """Build normalized task specs from decorators list or workflow object."""
        if isinstance(tasks, WorkflowDefinition):
            out: list[dict[str, Any]] = []
            for task in tasks.tasks:
                out.append(
                    {
                        "task_id": task["task_id"],
                        "agent_id": task["agent"],
                        "description": task.get("description", task["task_id"]),
                        "depends_on": task.get("depends_on", []),
                        "estimated_tool_calls": int(task.get("estimated_tool_calls", 1)),
                        "complexity": float(task.get("complexity", 1.0)),
                        "task_fn": task.get("task_fn"),
                        "loop_strategy": task.get("loop_strategy", "max_iterations"),
                        "max_iterations": int(task.get("max_iterations", 5)),
                        "enables_graph_mutation": bool(task.get("enables_graph_mutation", False)),
                        "tags": list(task.get("tags", [])),
                    }
                )
            return out

        task_list = list(tasks)
        wiring = AutoWire()
        inferred = wiring.infer_edges(task_list)
        explicit = {
            fn.__name__: list(getattr(fn, "_nexus_task_config", {}).get("depends_on", []))
            for fn in task_list
        }
        merged_order = wiring.build_dependency_order(task_list, explicit)
        dep_map: dict[str, set[str]] = {name: set(explicit.get(name, [])) for name in merged_order}
        for a, b in inferred:
            dep_map.setdefault(b, set()).add(a)

        by_name = {fn.__name__: fn for fn in task_list}
        out = []
        for name in merged_order:
            fn = by_name[name]
            cfg = getattr(fn, "_nexus_task_config", {})
            agent_cfg = cfg.get("agent")
            agent_id = agent_cfg.__name__ if hasattr(agent_cfg, "__name__") else str(agent_cfg)
            out.append(
                {
                    "task_id": name,
                    "agent_id": agent_id,
                    "description": (fn.__doc__ or name).strip(),
                    "depends_on": sorted(dep_map.get(name, set())),
                    "estimated_tool_calls": 1,
                    "complexity": 1.0,
                    "task_fn": fn,
                    "loop_strategy": cfg.get("loop_strategy", "max_iterations"),
                    "max_iterations": int(cfg.get("max_iterations", 5)),
                    "enables_graph_mutation": bool(cfg.get("enables_graph_mutation", False)),
                    "tags": list(cfg.get("tags", [])),
                }
            )
        return out

    async def resume(self, checkpoint_id: str, tasks: list[Callable]) -> NexusResult:
        """Restore from checkpoint then continue workflow execution."""
        ckpt = self.checkpoint_manager.load(checkpoint_id)
        for agent_id, snapshot in ckpt.agent_states.items():
            self.mesh.restore_snapshot(agent_id, snapshot)
        return await self._run_async(tasks, inputs={})
