"""
Core Mind class — persistent algorithmic mind.

Uses create_mind() factory as the single entry point.
Lazy-loads everything. Nothing happens until first use.
"""

from __future__ import annotations

import time
import json
import traceback
from pathlib import Path
from typing import Any
from contextlib import contextmanager

from aalgoi._data import normalize, detect_type
from aalgoi._result import SolveResult
from aalgoi._status import box, table


# ── Lazy singleton for module-level solve() ──────────────────────────

_global_mind: Mind | None = None
_default_path = Path.home() / ".aalgoi" / "mind"


def _get_global_mind() -> Mind:
    global _global_mind
    if _global_mind is None:
        _global_mind = Mind(_default_path)
    return _global_mind


# ── Module-level API ─────────────────────────────────────────────────

def solve(
    problem: str,
    data: Any = None,
    *,
    examples: list[dict] | None = None,
    max_iterations: int = 50,
    time_limit: float = 30.0,
    hint: str | None = None,
) -> SolveResult:
    """
    Solve any algorithmic problem.

    Parameters
    ----------
    problem : str
        Natural language description of the problem.
    data : any
        Input data — numpy, pandas, torch, files, URLs, anything.
    examples : list[dict], optional
        Input/output examples for verification.
    max_iterations : int
        Maximum cognitive steps (default 50).
    time_limit : float
        Approximate time budget in seconds (default 30).
        Note: enforced approximately; the mind may slightly exceed it.
    hint : str, optional
        Guide the mind (e.g. "use_dijkstra", "think_dp").

    Returns
    -------
    SolveResult
        Transparent result — use like the raw output,
        or access .algorithm, .code, .explain(), etc.

    Examples
    --------
    >>> import aalgoi
    >>> aalgoi.solve("sort the array", [3, 1, 4, 1, 5])
    [1, 1, 3, 4, 5]

    >>> r = aalgoi.solve("find two numbers that sum to target",
    ...                   {"nums": [2, 7, 11, 15], "target": 9})
    >>> r.algorithm
    'hash_complement'
    """
    m = _get_global_mind()
    return m.solve(
        problem, data,
        examples=examples,
        max_iterations=max_iterations,
        time_limit=time_limit,
        hint=hint,
    )


@contextmanager
def session(path: str | Path | None = None):
    """
    Context manager for a solve session.

    Parameters
    ----------
    path : str or Path, optional
        Mind persistence path. Default: ~/.aalgoi/mind

    Yields
    ------
    MindSession
        Session object with .solve(), .learn(), .status().

    Examples
    --------
    >>> with aalgoi.session() as m:
    ...     m.solve("sort", [3, 1, 2])
    ...     m.status()
    """
    from aalgoi._session import MindSession
    p = Path(path).expanduser() if path else _default_path
    sess = MindSession(p)
    try:
        yield sess
    finally:
        sess._summarize()


# ── Mind class ───────────────────────────────────────────────────────

class Mind:
    """
    A persistent algorithmic mind.

    Parameters
    ----------
    path : str or Path
        Persistence directory. Default: ~/.aalgoi/mind
    mode : str
        "fast" | "universal" (default) | "creative"
    max_iterations : int
        Default max cognitive steps per solve (default 50).
    time_limit : float
        Default approximate time budget per solve in seconds (default 30.0).
    """

    def __init__(
        self,
        path: str | Path = "~/.aalgoi/mind",
        *,
        mode: str = "universal",
        max_iterations: int = 50,
        time_limit: float = 30.0,
    ) -> None:
        self._path = Path(path).expanduser()
        self._path.mkdir(parents=True, exist_ok=True)
        self._mode = mode
        self._max_iterations = max_iterations
        self._time_limit = time_limit

        # Lazy-loaded internals — set by _ensure_loaded
        self._mind = None       # object returned by create_mind()
        self._kg = None         # AlgorithmicKnowledgeGraph
        self._loaded = False
        self._load_error: str | None = None

        # Session tracking (local, not dependent on safety API)
        self._solve_count: int = 0
        self._success_count: int = 0
        self._total_time_ms: float = 0.0

    def _ensure_loaded(self) -> None:
        """Lazy initialization via create_mind() factory."""
        if self._loaded:
            return
        if self._load_error:
            raise RuntimeError(f"Mind failed to load: {self._load_error}")

        try:
            from core.mind import create_mind
            from core.mind.model_config import MindConfig
        except ImportError as e:
            self._load_error = str(e)
            raise RuntimeError(
                f"Mind dependencies not available: {e}\n"
                f"Install with: pip install torch networkx"
            ) from e

        try:
            config = MindConfig()

            self._mind = create_mind(
                config=config,
                solver=None,
                synthesizer=None,
                persist_dir=str(self._path),
            )

            self._kg = getattr(self._mind, "knowledge_graph", None)
            if self._kg is None:
                self._kg = getattr(self._mind, "kg", None)

            self._loaded = True

        except Exception as e:
            self._load_error = f"{type(e).__name__}: {e}"
            raise RuntimeError(f"Mind failed to load: {self._load_error}") from e

    # ── Solve ────────────────────────────────────────────────────────

    def solve(
        self,
        problem: str,
        data: Any = None,
        *,
        examples: list[dict] | None = None,
        max_iterations: int | None = None,
        time_limit: float | None = None,
        hint: str | None = None,
    ) -> SolveResult:
        """
        Solve an algorithmic problem.

        The returned SolveResult is transparent — use it like the
        raw output value, or access metadata via properties.
        """
        # Normalize data first — handles numpy, pandas, torch, files, etc.
        data = normalize(data)

        self._ensure_loaded()

        max_it = max_iterations or self._max_iterations
        t_limit = time_limit or self._time_limit

        t0 = time.time()
        try:
            solution = self._mind.solve(
                problem_text=problem,
                data=data,
                examples=examples or [],
                max_iterations=max_it,
            )
        except Exception as e:
            return SolveResult(
                output=None,
                code=None,
                algorithm=None,
                complexity=None,
                principle=None,
                time_ms=(time.time() - t0) * 1000,
                is_novel=False,
                confidence=0.0,
                iterations=0,
                error=str(e),
            )

        elapsed_ms = (time.time() - t0) * 1000

        # The loop returns output=None — execute code ourselves
        output = None
        code = getattr(solution, "code", None)
        if code:
            output = self._execute_code(code, data)

        # If somehow output was set, use it
        if output is None:
            output = getattr(solution, "output", None)

        # Algorithm name
        algo_name = getattr(solution, "algorithm_name", None)
        if algo_name is None and code and self._kg:
            algo_name = self._find_algorithm_name(code)

        # Confidence from CorrectnessProof
        confidence = 0.0
        proof = getattr(solution, "correctness_proof", None)
        if proof is not None:
            confidence = getattr(proof, "confidence", 0.0)

        # Extract other fields
        complexity = getattr(solution, "time_complexity", None)
        principle = getattr(solution, "principle_applied", None)
        is_novel = getattr(solution, "is_novel", False)
        iterations = getattr(solution, "iterations", 0)

        # Local tracking
        self._solve_count += 1
        self._total_time_ms += elapsed_ms
        if confidence > 0.8:
            self._success_count += 1

        return SolveResult(
            output=output,
            code=code,
            algorithm=algo_name,
            complexity=complexity,
            principle=principle,
            time_ms=elapsed_ms,
            is_novel=is_novel,
            confidence=confidence,
            iterations=iterations,
            error=None,
        )

    # ── Training ─────────────────────────────────────────────────────

    def train(self, epochs: int = 10, augmentations: int = 5) -> dict:
        """
        Bootstrap train from expert demonstrations.

        Returns training statistics.
        """
        self._ensure_loaded()

        from core.mind.training.bootstrap_trainer import BootstrapTrainer
        from core.mind.model_config import MindConfig

        config = MindConfig()
        model = getattr(self._mind, "mind", self._mind)
        trainer = BootstrapTrainer(model, self._kg, config)
        stats = trainer.train(
            n_epochs=epochs,
            n_augmentations=augmentations,
        )

        import torch
        model_path = self._path / "model.pt"
        torch.save(model.state_dict(), model_path)

        return stats

    # ── Checkpoint / Rollback ────────────────────────────────────────

    def checkpoint(self) -> str | None:
        """Save a checkpoint. Returns path or None."""
        self._ensure_loaded()

        safety = getattr(self._mind, "safety", None)
        if safety is None:
            return None

        try:
            count = getattr(safety, "solve_count", self._solve_count)
            path = safety.auto_checkpoint(self._mind, count)
            return str(path) if path else None
        except Exception:
            return None

    def rollback(self, target: str = "last_good") -> dict:
        """
        Rollback to a safe state.

        Targets: "last", "last_good", "base"
        """
        self._ensure_loaded()

        safety = getattr(self._mind, "safety", None)
        if safety is None:
            return {"success": False, "error": "Safety not available"}

        try:
            result = safety.rollback(target)
            if result is None:
                return {"success": True, "target": target}
            return result
        except Exception as e:
            return {"success": False, "error": str(e)}

    # ── Federation ───────────────────────────────────────────────────

    def share(self) -> int:
        """Export federated updates. Returns count."""
        self._ensure_loaded()

        federation = getattr(self._mind, "federation", None)
        if federation is None:
            return 0

        try:
            outbox = getattr(federation, "outbox_dir", None)
            if outbox is None:
                return 0
            outbox_path = Path(outbox)
            if not outbox_path.exists():
                return 0
            return len(list(outbox_path.glob("*.json")))
        except Exception:
            return 0

    def receive(self) -> dict:
        """Import federated updates. Returns stats."""
        self._ensure_loaded()

        federation = getattr(self._mind, "federation", None)
        if federation is None:
            return {"updates_processed": 0, "algorithms_imported": 0,
                    "stats_updated": 0, "errors": []}

        try:
            results = federation.sync()
            if results is None:
                return {"updates_processed": 0, "algorithms_imported": 0,
                        "stats_updated": 0, "errors": []}
            return {
                "updates_processed": len(results) if isinstance(results, list) else 1,
                "algorithms_imported": len(results) if isinstance(results, list) else 0,
                "stats_updated": 0,
                "errors": [],
            }
        except Exception as e:
            return {"updates_processed": 0, "algorithms_imported": 0,
                    "stats_updated": 0, "errors": [str(e)]}

    # ── Benchmark ────────────────────────────────────────────────────

    def benchmark(self, verbose: bool = False) -> BenchmarkReport:
        """Run 50-problem benchmark suite."""
        from core.mind.benchmark import run_benchmark

        def solver(text, data):
            result = self.solve(text, data, max_iterations=30, time_limit=10.0)
            return result.output

        results = run_benchmark(solver, verbose=verbose)
        return BenchmarkReport(results)

    # ── Status ───────────────────────────────────────────────────────

    def status(self) -> str:
        """Return a beautiful status string."""
        self._ensure_loaded()

        lines = []

        kg_stats = {}
        if self._kg and hasattr(self._kg, "stats"):
            try:
                kg_stats = self._kg.stats()
            except Exception:
                pass

        lines.append(f"Algorithms:  {kg_stats.get('algorithms', '?')}")
        lines.append(f"Principles:  {kg_stats.get('principles', '?')}")
        lines.append(f"Problems:    {kg_stats.get('problems', '?')}")
        lines.append(f"Discovered:  {kg_stats.get('discovered_algorithms', 0)}")
        lines.append(f"Mode:        {self._mode}")

        lines.append(f"Solved:      {self._solve_count}")
        if self._solve_count > 0:
            rate = self._success_count / self._solve_count
            lines.append(f"Success rate: {rate:.0%}")
            avg_time = self._total_time_ms / self._solve_count
            lines.append(f"Avg time:     {avg_time:.1f}ms")

        return box(lines, title="\U0001f9e0 Algorithmic Mind")

    # ── Properties ───────────────────────────────────────────────────

    @property
    def algorithms(self) -> dict[str, AlgorithmInfo]:
        """All known algorithms with metadata."""
        self._ensure_loaded()
        result = {}
        if not self._kg:
            return result

        try:
            for nid, attrs in self._kg.graph.nodes(data=True):
                if attrs.get("node_type") != "algorithm":
                    continue
                algo = attrs["data"]
                result[algo.name] = AlgorithmInfo(
                    name=algo.name,
                    code=algo.code,
                    time_complexity=algo.time_complexity,
                    space_complexity=algo.space_complexity,
                    principles=algo.principles,
                    best_for=algo.best_for,
                    discovered_by=algo.discovered_by,
                    times_used=algo.times_used,
                    times_succeeded=algo.times_succeeded,
                    performance=algo.performance_history[-1] if algo.performance_history else 0.0,
                )
        except Exception:
            pass

        return result

    @property
    def principles(self) -> list[str]:
        """All known mathematical principles."""
        self._ensure_loaded()
        if not self._kg:
            return []
        try:
            return sorted([
                attrs["data"].name
                for nid, attrs in self._kg.graph.nodes(data=True)
                if attrs.get("node_type") == "principle"
            ])
        except Exception:
            return []

    @property
    def problems(self) -> list[str]:
        """All known problem types."""
        self._ensure_loaded()
        if not self._kg:
            return []
        try:
            return sorted([
                attrs["data"].signature
                for nid, attrs in self._kg.graph.nodes(data=True)
                if attrs.get("node_type") == "problem"
            ])
        except Exception:
            return []

    @property
    def path(self) -> Path:
        """Mind persistence path."""
        return self._path

    # ── Internal helpers ─────────────────────────────────────────────

    def _execute_code(self, code: str, data: Any) -> Any:
        """Execute generated code and return output."""
        try:
            namespace = {}
            exec(code, namespace)
            solve_fn = namespace.get("solve")
            if solve_fn is None:
                return None
            if isinstance(data, dict):
                return solve_fn(**data)
            return solve_fn(data)
        except TypeError:
            try:
                return solve_fn(data)
            except Exception:
                return None
        except Exception:
            return None

    def _find_algorithm_name(self, code: str) -> str | None:
        """Find algorithm name by code match in KG."""
        if not self._kg:
            return None
        try:
            for nid, attrs in self._kg.graph.nodes(data=True):
                if attrs.get("node_type") != "algorithm":
                    continue
                if attrs["data"].code == code:
                    return attrs["data"].name
        except Exception:
            pass
        return None

    def _infer_domain(self, data: Any) -> str:
        """Infer problem domain from data shape."""
        if data is None:
            return "unknown"
        if isinstance(data, dict):
            if "edges" in data:
                return "graph"
            if "X" in data:
                return "feature_matrix"
            if "nums" in data:
                return "integers"
            if "s" in data or "text" in data:
                return "text"
        if isinstance(data, (list, tuple)):
            return "integers"
        return "unknown"


# ── Helper classes ───────────────────────────────────────────────────

class AlgorithmInfo:
    """Metadata about a known algorithm."""
    __slots__ = (
        'name', 'code', 'time_complexity', 'space_complexity',
        'principles', 'best_for', 'discovered_by',
        'times_used', 'times_succeeded', 'performance',
    )

    def __init__(self, **kwargs):
        for k, v in kwargs.items():
            setattr(self, k, v)

    def __repr__(self) -> str:
        return (
            f"AlgorithmInfo("
            f"name={self.name!r}, "
            f"complexity={self.time_complexity!r}, "
            f"performance={self.performance:.2f})"
        )

    def display(self) -> str:
        lines = [
            f"Name:        {self.name}",
            f"Time:        {self.time_complexity}",
            f"Space:       {self.space_complexity}",
            f"Principles:  {', '.join(self.principles)}",
            f"Best for:    {', '.join(self.best_for[:5])}",
            f"Discovered:  {self.discovered_by}",
            f"Used:        {self.times_used}x ({self.times_succeeded} succeeded)",
            f"Performance: {self.performance:.2f}",
        ]
        return box(lines, title=f"\U0001f4cb {self.name}")


class BenchmarkReport:
    """Beautiful benchmark report."""

    def __init__(self, data: dict) -> None:
        self._data = data

    @property
    def accuracy(self) -> float:
        return self._data.get("accuracy", 0.0)

    @property
    def correct(self) -> int:
        return self._data.get("correct", 0)

    @property
    def total(self) -> int:
        return self._data.get("total", 0)

    @property
    def by_domain(self) -> dict:
        return self._data.get("by_domain", {})

    def __repr__(self) -> str:
        lines = [
            f"Accuracy: {self.correct}/{self.total} ({self.accuracy:.0%})",
        ]
        if self.by_domain:
            lines.append("\u2500" * 36)
            for domain in sorted(self.by_domain):
                d = self.by_domain[domain]
                pct = d["correct"] / max(d["total"], 1)
                bar = "\u2588" * int(pct * 20) + "\u2591" * (20 - int(pct * 20))
                lines.append(f"{domain:<16} {d['correct']}/{d['total']} {bar} {pct:.0%}")
        return box(lines, title="\U0001f4ca Benchmark")

    def __str__(self) -> str:
        return repr(self)

    def details(self) -> list[dict]:
        """Per-problem results."""
        return self._data.get("problems", [])
