"""
DSPy + GEPA skill optimizer.

Wraps SKILL.md body as an optimizable DSPy module and uses GEPA
(or MIPROv2 fallback) to improve the text based on eval harness
feedback.  Only enqueues a pending proposal when the candidate
beats the baseline on held-out validation tasks.

dspy-ai is auto-installed into ~/.mugi/venv on first use when
text evolution is enabled.
"""
from __future__ import annotations

import json
import os
import random
import time
from datetime import datetime, timezone
from typing import Any, Callable

_has_dspy = False
try:
    import dspy
    _has_dspy = True
except ImportError:
    pass

_dspy_install_attempted = False


def _auto_install_dspy() -> bool:
    """Try to pip-install dspy-ai into the backend venv. Returns True on success."""
    global _has_dspy, _dspy_install_attempted
    if _dspy_install_attempted:
        return _has_dspy
    _dspy_install_attempted = True
    import subprocess
    venv_pip = os.path.join(os.path.expanduser("~"), ".mugi", "venv", "bin", "pip")
    if not os.path.isfile(venv_pip):
        _log("Cannot auto-install dspy: backend venv pip not found", "WARNING")
        return False
    req_file = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
                            "requirements-text-evolution.txt")
    if not os.path.isfile(req_file):
        _log("Cannot auto-install dspy: requirements-text-evolution.txt not found", "WARNING")
        return False
    _log("Auto-installing dspy-ai into backend venv...", "INFO")
    try:
        result = subprocess.run(
            [venv_pip, "install", "--prefer-binary", "-r", req_file],
            capture_output=True, text=True, timeout=300,
        )
        if result.returncode == 0:
            import importlib
            global dspy
            dspy = importlib.import_module("dspy")
            _has_dspy = True
            _log("dspy-ai installed successfully", "INFO")
            return True
        _log(f"dspy-ai install failed: {(result.stderr or result.stdout or '')[:300]}", "ERROR")
    except Exception as exc:
        _log(f"dspy-ai auto-install error: {exc}", "ERROR")
    return False


def _log(msg: str, level: str = "WARNING") -> None:
    try:
        from tools._logging import _debug_log
        _debug_log(f"[evolution/skill_gepa] {msg}", level)
    except Exception:
        pass


def has_dspy_library() -> bool:
    if _has_dspy:
        return True
    return _auto_install_dspy()


# ---------------------------------------------------------------------------
# Status tracking
# ---------------------------------------------------------------------------

_EVOLUTION_DIR = os.path.join(os.path.expanduser("~"), ".mugi", "evolution")
_STATUS_PATH = os.path.join(_EVOLUTION_DIR, "text_evolution_status.json")
_META_PATH = os.path.join(_EVOLUTION_DIR, "text_evolution_meta.json")


def _read_json(path: str, defaults: dict | None = None) -> dict:
    try:
        if os.path.isfile(path):
            with open(path, "r", encoding="utf-8") as f:
                data = json.load(f)
            if isinstance(data, dict):
                return data
    except Exception:
        pass
    return dict(defaults or {})


def _write_json(path: str, data: dict) -> None:
    import tempfile
    os.makedirs(os.path.dirname(path), mode=0o700, exist_ok=True)
    fd, tmp = tempfile.mkstemp(dir=os.path.dirname(path), suffix=".tmp")
    try:
        with os.fdopen(fd, "w") as f:
            json.dump(data, f, indent=2, default=str)
            f.flush()
            os.fsync(f.fileno())
        os.replace(tmp, path)
    except Exception:
        try:
            os.unlink(tmp)
        except OSError:
            pass
        raise


def get_text_evolution_status() -> dict:
    defaults = {
        "running": False,
        "last_run_at": None,
        "last_error": None,
        "skipped_reason": None,
        "current_skill": None,
    }
    return _read_json(_STATUS_PATH, defaults)


def _set_status(**fields: Any) -> None:
    status = get_text_evolution_status()
    status.update(fields)
    _write_json(_STATUS_PATH, status)


def get_text_evolution_meta() -> dict:
    defaults = {
        "last_text_evolution_at": None,
        "metric_calls_today": 0,
        "metric_calls_date": None,
    }
    return _read_json(_META_PATH, defaults)


def _increment_metric_calls(n: int = 1) -> int:
    """Increment daily metric call counter. Returns new total."""
    meta = get_text_evolution_meta()
    today = datetime.now(timezone.utc).strftime("%Y-%m-%d")
    if meta.get("metric_calls_date") != today:
        meta["metric_calls_today"] = 0
        meta["metric_calls_date"] = today
    meta["metric_calls_today"] = meta.get("metric_calls_today", 0) + n
    _write_json(_META_PATH, meta)
    return meta["metric_calls_today"]


def clear_stale_running() -> None:
    """Clear running flag on startup (crash recovery)."""
    status = get_text_evolution_status()
    if status.get("running"):
        _log("Clearing stale text_evolution running flag (crash recovery)", "INFO")
        status["running"] = False
        status["last_error"] = "cleared stale running flag on startup"
        _write_json(_STATUS_PATH, status)


# ---------------------------------------------------------------------------
# DSPy module wrapper
# ---------------------------------------------------------------------------

class _SkillModule:
    """Thin wrapper presenting SKILL.md body as an optimizable text field.

    When dspy is not installed this class is a no-op placeholder used only
    for type hints.
    """

    def __init__(self, baseline_text: str):
        self.skill_body = baseline_text

    def get_text(self) -> str:
        return self.skill_body

    def set_text(self, text: str) -> None:
        self.skill_body = text


# ---------------------------------------------------------------------------
# GEPA optimization loop
# ---------------------------------------------------------------------------

def run_gepa_optimization(
    skill_name: str,
    baseline_text: str,
    eval_tasks: list[dict],
    config: dict | None = None,
    stop_check: Callable[[], bool] | None = None,
) -> dict:
    """Run GEPA optimization for a single skill.

    Returns {"improved": bool, "candidate_text": str, "baseline_score": float,
             "candidate_score": float, "proposal_id": str|None, "error": str|None}.
    """
    if config is None:
        from core.config import get_text_evolution_config
        config = get_text_evolution_config()

    if not _has_dspy:
        return {"improved": False, "error": "dspy not installed"}

    min_tasks = config.get("min_eval_tasks", 3)
    if len(eval_tasks) < min_tasks:
        return {"improved": False, "error": f"Need at least {min_tasks} eval tasks, have {len(eval_tasks)}"}

    max_calls_per_day = config.get("max_metric_calls_per_day", 50)
    meta = get_text_evolution_meta()
    today = datetime.now(timezone.utc).strftime("%Y-%m-%d")
    calls_today = meta.get("metric_calls_today", 0) if meta.get("metric_calls_date") == today else 0
    if calls_today >= max_calls_per_day:
        return {"improved": False, "error": f"Daily metric call limit reached ({max_calls_per_day})"}

    _set_status(running=True, current_skill=skill_name, last_error=None, skipped_reason=None)

    try:
        return _run_gepa_inner(skill_name, baseline_text, eval_tasks, config, stop_check)
    except Exception as exc:
        _log(f"GEPA optimization failed: {exc}", "ERROR")
        _set_status(running=False, last_error=str(exc)[:500])
        return {"improved": False, "error": str(exc)}
    finally:
        _set_status(
            running=False,
            current_skill=None,
            last_run_at=datetime.now(timezone.utc).isoformat(),
        )
        meta = get_text_evolution_meta()
        meta["last_text_evolution_at"] = datetime.now(timezone.utc).isoformat()
        _write_json(_META_PATH, meta)


def _run_gepa_inner(
    skill_name: str,
    baseline_text: str,
    eval_tasks: list[dict],
    config: dict,
    stop_check: Callable[[], bool] | None,
) -> dict:
    from evolution.skill_eval_harness import evaluate_skill
    from evolution.pending_store import enqueue_proposal, sha256_of, has_pending_for_skill

    if has_pending_for_skill(skill_name):
        _set_status(skipped_reason=f"pending proposal already exists for '{skill_name}'")
        return {"improved": False, "error": f"Pending proposal already exists for '{skill_name}'"}

    # Train/val split
    random.shuffle(eval_tasks)
    split = max(1, len(eval_tasks) // 4)  # ~25% validation
    val_tasks = eval_tasks[:split]
    train_tasks = eval_tasks[split:]

    if not train_tasks:
        return {"improved": False, "error": "Not enough tasks for train split"}

    # Baseline score on validation set
    baseline_result = evaluate_skill(baseline_text, val_tasks, config, stop_check)
    baseline_score = baseline_result["mean_score"]
    _increment_metric_calls(len(val_tasks))

    if stop_check and stop_check():
        return {"improved": False, "error": "stopped"}

    # Configure DSPy LM
    optimizer_model = (config.get("optimizer_model") or config.get("eval_model") or "").strip()
    base_url = (config.get("base_url") or "").strip()
    api_key = (config.get("api_key") or "").strip()

    if not base_url:
        from agent_loop import get_backend_base_url
        base_url = get_backend_base_url()
    if not api_key:
        api_key = "lm-studio"

    lm = dspy.LM(
        model=f"openai/{optimizer_model or 'local-model'}",
        api_base=base_url,
        api_key=api_key,
        temperature=0.7,
        max_tokens=4096,
    )
    dspy.configure(lm=lm)

    # Generate candidates via DSPy prompt optimization
    best_text = baseline_text
    best_score = baseline_score
    max_candidates = min(5, max(1, config.get("max_metric_calls_per_day", 50) - _get_calls_today()))

    for i in range(max_candidates):
        if stop_check and stop_check():
            break

        candidate_text = _generate_candidate(baseline_text, best_text, train_tasks, config)
        if not candidate_text or candidate_text.strip() == best_text.strip():
            continue

        # Evaluate candidate on train set first (cheap gate)
        train_result = evaluate_skill(candidate_text, train_tasks, config, stop_check)
        _increment_metric_calls(len(train_tasks))

        if train_result["mean_score"] <= best_score:
            continue

        # Validate on held-out set
        val_result = evaluate_skill(candidate_text, val_tasks, config, stop_check)
        _increment_metric_calls(len(val_tasks))

        if val_result["mean_score"] > best_score:
            best_text = candidate_text
            best_score = val_result["mean_score"]
            _log(f"Candidate {i} improved: {val_result['mean_score']:.3f} > {baseline_score:.3f}", "INFO")

    if best_text == baseline_text or best_score <= baseline_score:
        _set_status(skipped_reason="no improvement found")
        return {
            "improved": False,
            "baseline_score": baseline_score,
            "candidate_score": best_score,
            "error": None,
        }

    # Enqueue proposal
    try:
        proposal = enqueue_proposal(
            skill_name=skill_name,
            proposed_content=best_text,
            baseline_sha256=sha256_of(baseline_text),
            metrics={
                "baseline_score": round(baseline_score, 4),
                "candidate_score": round(best_score, 4),
                "n_train_tasks": len(train_tasks),
                "n_val_tasks": len(val_tasks),
                "optimization_model": optimizer_model or "local-model",
            },
        )

        # Emit background event
        try:
            from mugi_runner import _write_event
            _write_event({
                "type": "text_evolution_proposal",
                "skill_name": skill_name,
                "score_improvement": round(best_score - baseline_score, 4),
                "proposal_id": proposal["id"],
            })
        except Exception:
            pass

        return {
            "improved": True,
            "baseline_score": baseline_score,
            "candidate_score": best_score,
            "proposal_id": proposal["id"],
            "error": None,
        }
    except ValueError as exc:
        return {"improved": False, "error": str(exc)}


def _generate_candidate(
    baseline_text: str,
    current_best: str,
    train_tasks: list[dict],
    config: dict,
) -> str | None:
    """Use DSPy LM to generate an improved candidate skill text."""
    task_descriptions = "\n".join(
        f"- {t.get('user_message', '')[:200]}" for t in train_tasks[:5]
    )
    prompt = (
        "You are optimizing an AI workflow procedure (SKILL.md) to better handle these tasks:\n\n"
        f"{task_descriptions}\n\n"
        "Current procedure:\n"
        f"```\n{current_best}\n```\n\n"
        "Generate an improved version of this procedure that would produce better results "
        "on the tasks above. Maintain the same general structure but improve clarity, "
        "completeness, and specificity. Output ONLY the improved procedure text, nothing else."
    )
    try:
        lm = dspy.settings.lm
        response = lm(prompt)
        if isinstance(response, list) and response:
            text = response[0] if isinstance(response[0], str) else str(response[0])
        elif isinstance(response, str):
            text = response
        else:
            text = str(response)

        # Strip markdown fences if present
        text = text.strip()
        if text.startswith("```"):
            lines = text.split("\n")
            if lines[0].startswith("```"):
                lines = lines[1:]
            if lines and lines[-1].strip() == "```":
                lines = lines[:-1]
            text = "\n".join(lines)

        return text.strip() if text.strip() else None
    except Exception as exc:
        _log(f"Candidate generation failed: {exc}", "ERROR")
        return None


def _get_calls_today() -> int:
    meta = get_text_evolution_meta()
    today = datetime.now(timezone.utc).strftime("%Y-%m-%d")
    if meta.get("metric_calls_date") != today:
        return 0
    return meta.get("metric_calls_today", 0)
