"""
Swarm: entrypoint for users. Orchestrates planner → scheduler → executor → results.

User code:
    swarm = Swarm(worker_count=4)
    result = swarm.run("Analyze diffusion model research")
    # Or with config file:
    swarm = Swarm(config="hivemind.toml")
    result = swarm.run("analyze diffusion models")
"""

import json
import os
import threading
from pathlib import Path
from datetime import datetime, timezone

from hivemind.types.task import Task
from hivemind.types.event import Event, events
from hivemind.utils.event_logger import EventLog
from hivemind.utils.models import resolve_model

from hivemind.swarm.planner import Planner
from hivemind.swarm.scheduler import Scheduler
from hivemind.swarm.executor import Executor
from hivemind.agents.agent import Agent


def _persist_dag(scheduler: Scheduler, event_log: EventLog) -> None:
    """Write task DAG to events dir as {run_id}_dag.json for graph export."""
    run_id = getattr(event_log, "run_id", None)
    if not run_id:
        return
    log_path = getattr(event_log, "log_path", None)
    if not log_path:
        return
    events_dir = os.path.dirname(log_path)
    nodes = [
        {"id": t.id, "description": (t.description or "")[:200]}
        for t in scheduler._tasks.values()
    ]
    edges = list(scheduler._graph.edges())
    path = os.path.join(events_dir, f"{run_id}_dag.json")
    with open(path, "w", encoding="utf-8") as f:
        json.dump({"nodes": nodes, "edges": edges}, f, indent=0)


class Swarm:
    """Orchestrates planner, scheduler, executor, and agent. Single entrypoint for running a task."""

    def __init__(
        self,
        worker_count: int | None = None,
        worker_model: str | None = None,
        planner_model: str | None = None,
        event_log: EventLog | None = None,
        adaptive: bool | None = None,
        memory_router=None,
        store_swarm_memory: bool = True,
        use_tools: bool | None = None,
        config: str | Path | object | None = None,
    ) -> None:
        # Load from config file or config object if provided
        cfg = None
        if config is not None:
            if isinstance(config, (str, Path)):
                from hivemind.config import get_config

                cfg = get_config(config_path=str(config))
            else:
                cfg = config
        if cfg is not None:
            self.worker_count = (
                worker_count
                if worker_count is not None
                else getattr(cfg.swarm, "workers", 4)
            )
            worker_raw = worker_model if worker_model is not None else cfg.models.worker
            planner_raw = (
                planner_model if planner_model is not None else cfg.models.planner
            )
            self.worker_model = resolve_model(worker_raw, "analysis")
            self.planner_model = resolve_model(planner_raw, "planning")
            self.adaptive = (
                adaptive
                if adaptive is not None
                else getattr(cfg.swarm, "adaptive_planning", False)
                or getattr(cfg.swarm, "adaptive_execution", False)
            )
            self.use_tools = use_tools if use_tools is not None else True
            self.speculative_execution = getattr(
                cfg.swarm, "speculative_execution", False
            )
            self.cache_enabled = getattr(cfg.swarm, "cache_enabled", False)
            self.parallel_tools = getattr(cfg.swarm, "parallel_tools", True)
            self._config = cfg
        else:
            self.worker_count = worker_count if worker_count is not None else 4
            worker_raw = worker_model if worker_model is not None else "mock"
            planner_raw = planner_model if planner_model is not None else "mock"
            self.worker_model = resolve_model(worker_raw, "analysis")
            self.planner_model = resolve_model(planner_raw, "planning")
            self.adaptive = adaptive if adaptive is not None else False
            self.use_tools = use_tools if use_tools is not None else False
            self.speculative_execution = False
            self.cache_enabled = False
            self.parallel_tools = True
            self._config = None
        self.event_log = event_log or EventLog()
        self.memory_router = memory_router
        self.store_swarm_memory = store_swarm_memory
        self._last_scheduler: Scheduler | None = None
        self._last_reasoning_store = None
        self._pause_event = threading.Event()
        self._pause_event.set()

    def pause(self) -> None:
        """Pause the executor: currently-running tasks finish, no new tasks start."""
        self._pause_event.clear()

    def resume(self) -> None:
        """Resume the executor so new tasks can be picked."""
        self._pause_event.set()

    def run(self, user_task: str) -> dict[str, str]:
        """
        Create root task → plan subtasks → add to scheduler → run executor → return task_id → result.
        """
        self._emit(events.SWARM_STARTED, {"user_task": user_task[:200]})

        root = Task(id="root", description=user_task, dependencies=[])
        from hivemind.intelligence.strategy_selector import StrategySelector
        from hivemind.intelligence.strategies import get_strategy_for

        selector = StrategySelector()
        selected = selector.select(root)
        strategy_instance = get_strategy_for(selected)
        prompt_suffix = selector.suggest_planner_prompt_suffix(selected)
        planner = Planner(
            model_name=self.planner_model,
            event_log=self.event_log,
            strategy=strategy_instance,
            prompt_suffix=prompt_suffix,
        )
        subtasks = planner.plan(root)

        scheduler = Scheduler()
        scheduler.add_tasks(subtasks)
        _persist_dag(scheduler, self.event_log)

        from hivemind.reasoning.store import ReasoningStore

        reasoning_store = ReasoningStore()
        agent = Agent(
            model_name=self.worker_model,
            event_log=self.event_log,
            memory_router=self.memory_router,
            store_result_to_memory=False,
            use_tools=self.use_tools,
            reasoning_store=reasoning_store,
            user_task=user_task,
            parallel_tools=getattr(self, "parallel_tools", True),
        )
        task_cache = None
        semantic_cache = None
        if getattr(self, "cache_enabled", False):
            from hivemind.cache import TaskCache

            task_cache = TaskCache()
            cfg = getattr(self, "_config", None)
            if cfg and getattr(getattr(cfg, "cache", None), "semantic", False):
                from hivemind.cache.task_cache import SemanticTaskCache

                cache_cfg = cfg.cache
                semantic_cache = SemanticTaskCache(
                    similarity_threshold=getattr(cache_cfg, "similarity_threshold", 0.92),
                    max_age_hours=getattr(cache_cfg, "max_age_hours", 168.0),
                )
        complexity_router = None
        models_config = None
        if self._config is not None:
            from hivemind.providers.complexity_router import TaskComplexityRouter

            complexity_router = TaskComplexityRouter()
            models_config = self._config.models
        executor = Executor(
            scheduler=scheduler,
            agent=agent,
            worker_count=self.worker_count,
            event_log=self.event_log,
            planner=planner if self.adaptive else None,
            adaptive=self.adaptive,
            speculative_execution=getattr(self, "speculative_execution", False),
            task_cache=task_cache,
            pause_event=self._pause_event,
            semantic_cache=semantic_cache,
            complexity_router=complexity_router,
            models_config=models_config,
            streaming_dag=True,
        )
        executor.run_sync()

        self._last_scheduler = scheduler
        self._last_reasoning_store = reasoning_store
        results = scheduler.get_results()
        if self.store_swarm_memory and self.memory_router and results:
            self._store_swarm_memory(user_task, scheduler)
        self._emit(events.SWARM_FINISHED, {"task_count": len(results)})
        try:
            from hivemind.intelligence.analysis.run_report import build_report_from_events
            from hivemind.runtime.run_history import RunHistory
            log_path = getattr(self.event_log, "log_path", None)
            if log_path:
                events_dir = os.path.dirname(log_path)
                run_id = getattr(self.event_log, "run_id", None)
                if run_id:
                    report = build_report_from_events(run_id, events_dir)
                    RunHistory().record_run(report)
        except Exception:
            pass
        return results

    @property
    def last_completed_tasks(self) -> list[Task]:
        """After run(), return completed tasks (id, description, result) for report building."""
        if self._last_scheduler is None:
            return []
        return self._last_scheduler.get_completed_tasks()

    def map_reduce(
        self, dataset: list, map_fn, reduce_fn, worker_count: int | None = None
    ):
        """
        First-class map-reduce: run map_fn on each item in parallel, then reduce_fn on results.
        Uses the same worker pool pattern as the executor.
        """
        from hivemind.swarm.map_reduce import map_reduce as _map_reduce

        workers = worker_count if worker_count is not None else self.worker_count
        return _map_reduce(dataset, map_fn, reduce_fn, worker_count=workers)

    def _store_swarm_memory(self, user_task: str, scheduler: Scheduler) -> None:
        """Store important outputs (research findings, summaries, results) into memory after run."""
        from hivemind.memory.memory_store import (
            MemoryStore,
            get_default_store,
            generate_memory_id,
        )
        from hivemind.memory.memory_types import MemoryRecord, MemoryType
        from hivemind.memory.memory_index import MemoryIndex

        store = getattr(self.memory_router, "store", None)
        if not isinstance(store, MemoryStore):
            store = get_default_store()
        index = getattr(self.memory_router, "index", None) or MemoryIndex(store)
        for task in scheduler.get_completed_tasks():
            content = (task.result or "").strip()
            if not content or len(content) < 10:
                continue
            desc = (task.description or "").lower()
            if "research" in desc or "paper" in desc or "literature" in desc:
                mt = MemoryType.RESEARCH
            elif "code" in desc or "codebase" in desc or "analyze" in desc:
                mt = MemoryType.ARTIFACT
            elif "data" in desc or "dataset" in desc or "experiment" in desc:
                mt = MemoryType.SEMANTIC
            else:
                mt = MemoryType.EPISODIC
            record = MemoryRecord(
                id=generate_memory_id(),
                memory_type=mt,
                source_task=task.id,
                content=content[:15000],
                tags=["swarm", "task", task.id, user_task[:100]],
            )
            record = index.ensure_embedding(record)
            store.store(record)

    def _emit(self, event_type: events, payload: dict) -> None:
        self.event_log.append_event(
            Event(
                timestamp=datetime.now(timezone.utc), type=event_type, payload=payload
            )
        )
