"""GEPA-compatible proposer that adapts optimizer.py's improve_skill() logic.

Wraps the existing 4-dimension improvement framework (tool usage, domain
modeling, task strategy, safety) into GEPA's ProposalFn interface.

LLM calls use litellm for provider-agnostic model routing.
"""

from __future__ import annotations

import json
import re
from pathlib import Path

import litellm

from clawbench._env import DEFAULT_OPTIMIZER_MODEL


class Proposer:
    """GEPA-compatible proposer wrapping optimizer.py's improvement logic.

    GEPA ProposalFn signature:
        (candidate: dict[str,str], dataset: list, components: list[str]) -> dict[str,str]
    """

    def __init__(
        self,
        model: str = DEFAULT_OPTIMIZER_MODEL,
        run_dir: Path | None = None,
    ):
        self.model = model
        self.run_dir = run_dir
        self._call_count = 0

    def __call__(
        self,
        candidate: dict[str, str],
        reflective_dataset: list,
        components_to_update: list[str],
    ) -> dict[str, str]:
        """Propose an improved candidate based on evaluation feedback.

        Args:
            candidate: Current candidate, e.g. {"skills": "...SKILL.md..."}
            reflective_dataset: List of side_info dicts from recent evaluations
            components_to_update: Which candidate keys to update, e.g. ["skills"]

        Returns:
            Dict with updated component values.
        """
        results = {}
        for component in components_to_update:
            current_value = candidate.get(component, "")
            # GEPA passes reflective_dataset as {component_name: [side_info_items...]}
            component_data = reflective_dataset.get(component, []) if isinstance(reflective_dataset, dict) else reflective_dataset
            prompt = self._build_prompt(current_value, component_data)
            new_value = self._call_llm(prompt)
            # Ensure frontmatter is preserved
            new_value = self._ensure_frontmatter(current_value, new_value)
            results[component] = new_value

        self._call_count += 1
        if self.run_dir:
            self._log(candidate, reflective_dataset, components_to_update, results)

        return results

    # Total char budget for all trajectory data in the prompt.
    # Keeps the prompt well within model context even with many tasks.
    TRAJECTORY_BUDGET = 12000

    def _build_prompt(self, current_skill: str, side_infos: list) -> str:
        """Build improvement prompt from current skill and evaluation feedback.

        Groups side_infos by task_name. For multi-task (each task_name unique),
        this is equivalent to the original per-entry behavior. For single-task
        with replicated workers, all N entries share a task_name and are
        collapsed into a single grouped entry with an LLM-synthesized summary.
        """
        # Extract structured entries from side_infos
        entries = []
        for si in side_infos:
            if isinstance(si, str):
                continue  # skip malformed entries
            feedback = si.get("Feedback", {})
            task_name = si.get("Input", {}).get("task_name", "unknown")
            status = feedback.get("Status", "unknown")
            rewards = feedback.get("Rewards", {})
            # GEPA renames "scores" to "Scores (Higher is Better)"
            if not rewards:
                rewards = si.get("Scores (Higher is Better)", {})
            score = rewards.get("score", rewards.get("reward", 0.0)) if rewards else 0.0

            gen_outputs = si.get("Generated Outputs", {})
            entries.append({
                "task_name": task_name,
                "status": status,
                "score": score,
                "rewards": rewards,
                "error": feedback.get("Error"),
                "agent_trace": gen_outputs.get("Agent Trace", ""),
                "trajectory_analysis": gen_outputs.get("Trajectory Analysis", ""),
                "duration": si.get("Input", {}).get("duration_sec"),
            })

        # Group by task_name
        from collections import defaultdict
        groups: dict[str, list[dict]] = defaultdict(list)
        for entry in entries:
            groups[entry["task_name"]].append(entry)

        # Collapse groups: if a task_name has N > 1 entries, synthesize them
        failed = []
        passed = []
        for task_name, group in groups.items():
            if len(group) == 1:
                # Single entry (multi-task case) — pass through unchanged
                entry = group[0]
                if entry["score"] > 0:
                    passed.append(entry)
                else:
                    failed.append(entry)
            else:
                # Replicated single-task: synthesize N entries into one grouped entry
                grouped_entry = self._synthesize_group(task_name, group)
                if grouped_entry["score"] > 0:
                    passed.append(grouped_entry)
                else:
                    failed.append(grouped_entry)

        prompt = (
            "You are optimizing an agent skill — instructions that guide an "
            "AI agent on how to complete tasks in a specific environment.\n\n"
            "The agent reads these instructions, then completes various tasks "
            "using CLI tools or APIs available in the environment. The current "
            "skill content already documents the tools and patterns to use — "
            "work with whatever is documented there.\n\n"
            f"Current results: {len(passed)} passed, {len(failed)} failed "
            f"out of {len(entries)} evaluations\n\n"
            "Here is the current skill content:\n"
            "<current_skill>\n"
            f"{current_skill}\n"
            "</current_skill>\n\n"
        )

        # Budget trajectory chars: failed tasks get ~80%, passed get ~20%
        budget = self.TRAJECTORY_BUDGET
        n_failed = len(failed)
        n_passed = len(passed)
        per_failed = int(budget * 0.8 / max(n_failed, 1))
        per_passed = int(budget * 0.2 / max(n_passed, 1))

        if failed:
            prompt += "FAILED TASKS:\n"
            for f in failed:
                prompt += f"\n  Task: {f['task_name']}\n"
                prompt += f"  Score: {f['score']:.2f}\n"
                if f.get("n_runs"):
                    prompt += f"  Runs: {f['n_runs']} ({f['n_passed']} passed, {f['n_failed']} failed)\n"
                if f["error"]:
                    prompt += f"  Error: {f['error']}\n"
                if f["rewards"]:
                    prompt += f"  Sub-scores: {json.dumps(f['rewards'])}\n"
                if f["trajectory_analysis"]:
                    prompt += (
                        f"  Trajectory analysis:\n"
                        f"  <analysis>\n{f['trajectory_analysis']}\n  </analysis>\n"
                    )
                if f["agent_trace"]:
                    trace = f["agent_trace"][:per_failed]
                    if len(f["agent_trace"]) > per_failed:
                        trace += "\n  ... (trajectory truncated for prompt budget)"
                    prompt += (
                        f"  Raw agent trajectory:\n"
                        f"  <agent_trajectory>\n{trace}\n  </agent_trajectory>\n"
                    )
            prompt += "\n"

        if passed:
            prompt += "PASSED TASKS:\n"
            for p in passed:
                prompt += f"\n  - {p['task_name']} (score: {p['score']:.2f})\n"
                if p.get("n_runs"):
                    prompt += f"    Runs: {p['n_runs']} ({p['n_passed']} passed, {p['n_failed']} failed)\n"
                if p["trajectory_analysis"]:
                    prompt += f"    Analysis: {p['trajectory_analysis'][:per_passed]}\n"
                elif p["agent_trace"]:
                    trace_lines = p["agent_trace"].split("\n")
                    key_lines = [l for l in trace_lines if any(kw in l for kw in ["[Bash", "[Summary", "curl"])]
                    summary = "; ".join(key_lines[:10])
                    if len(summary) > per_passed:
                        summary = summary[:per_passed] + "..."
                    if summary:
                        prompt += f"    Key actions: {summary}\n"
            prompt += "\n"

        prompt += (
            "Analyze the failures and improve the skill.\n\n"
            "HOW TO IMPROVE:\n"
            "- Start from the current skill content as a base\n"
            "- You MAY edit, reword, or fix existing lines if they are wrong, "
            "misleading, or incomplete\n"
            "- You MAY add new sections, examples, or strategies\n"
            "- You MUST keep the YAML frontmatter block (---...---) unchanged\n"
            "- You MUST keep the same tools and their invocation patterns\n"
            "- You MUST NOT discard the existing structure — the result should be "
            "recognizably the same document, improved\n"
            "- Keep under 500 lines\n\n"
            "IMPROVEMENT DIMENSIONS:\n"
            "- Tool usage: commands, queries, pagination, output parsing, batch ops\n"
            "- Domain modeling: entities, relationships, naming conventions\n"
            "- Task strategy: scope → discover → act in batches → verify\n"
            "- Safety: confirm scope before destructive actions\n\n"
            "ANTI-PATTERNS:\n"
            "- Don't memorize specific task names or entity IDs (overfitting)\n"
            "- Don't sacrifice passing tasks to fix failing ones\n"
            "- Don't add fluff — every instruction must be actionable\n\n"
            "Respond with the complete improved SKILL.md inside <new_skill> tags.\n\n"
            "<new_skill>\n"
            "... improved SKILL.md ...\n"
            "</new_skill>"
        )
        return prompt

    def _synthesize_group(self, task_name: str, group: list[dict]) -> dict:
        """Synthesize N replicated task entries into a single grouped entry.

        For single-task with workers > 1, all N side_infos share the same
        task_name. This method collapses them into one entry with aggregate
        stats and an LLM-synthesized summary of failure/success patterns.
        """
        scores = [e["score"] for e in group]
        n_passed = sum(1 for s in scores if s > 0)
        n_failed = len(group) - n_passed
        mean_score = sum(scores) / len(scores)

        failed_entries = [e for e in group if e["score"] <= 0]
        passed_entries = [e for e in group if e["score"] > 0]

        # Build synthesis input from individual trajectory analyses
        synthesis_parts = []
        if failed_entries:
            synthesis_parts.append(f"FAILED RUNS ({n_failed}/{len(group)}):")
            for i, e in enumerate(failed_entries):
                analysis = e.get("trajectory_analysis") or e.get("agent_trace", "")[:2000]
                synthesis_parts.append(f"  Run {i+1} (score={e['score']:.2f}): {analysis}")
        if passed_entries:
            synthesis_parts.append(f"\nPASSED RUNS ({n_passed}/{len(group)}):")
            for i, e in enumerate(passed_entries):
                analysis = e.get("trajectory_analysis") or e.get("agent_trace", "")[:2000]
                synthesis_parts.append(f"  Run {i+1} (score={e['score']:.2f}): {analysis}")

        synthesis_input = "\n".join(synthesis_parts)

        # LLM call to synthesize the group
        try:
            synthesis_prompt = (
                f"Synthesize these {len(group)} runs of the same task '{task_name}' "
                f"into a concise analysis. Focus on:\n"
                f"1. Failure patterns — common root causes across failed runs\n"
                f"2. Success patterns — what the successful runs did differently\n"
                f"3. Variance signal — what changed between pass/fail on the same task\n\n"
                f"{synthesis_input}\n\n"
                f"Respond with a concise synthesis (under 2000 chars)."
            )
            response = litellm.completion(
                model=self.model,
                max_tokens=2000,
                messages=[{"role": "user", "content": synthesis_prompt}],
            )
            synthesis = response.choices[0].message.content or ""
        except Exception:
            # Fallback: concatenate individual analyses
            synthesis = synthesis_input[:3000]

        # Pick the most representative error from failed entries
        representative_error = None
        for e in failed_entries:
            if e.get("error"):
                representative_error = e["error"]
                break

        return {
            "task_name": task_name,
            "status": "mixed" if n_passed > 0 and n_failed > 0 else ("passed" if n_passed > 0 else "failed"),
            "score": mean_score,
            "rewards": {},
            "error": representative_error,
            "agent_trace": "",
            "trajectory_analysis": synthesis,
            "duration": None,
            "n_runs": len(group),
            "n_passed": n_passed,
            "n_failed": n_failed,
        }

    @staticmethod
    def _ensure_frontmatter(original: str, proposed: str) -> str:
        """Ensure the proposal preserves the original's frontmatter and core content.

        If the LLM rewrote the skill from scratch (dropping frontmatter or key
        sections), fall back to appending only the NEW content to the original.
        """
        if not original.strip():
            return proposed

        # Extract frontmatter from original
        has_frontmatter = original.startswith("---")
        if has_frontmatter:
            parts = original.split("---", 2)
            if len(parts) >= 3:
                frontmatter = f"---{parts[1]}---"
            else:
                frontmatter = None
                has_frontmatter = False
        else:
            frontmatter = None

        # Check if the proposed skill preserved the original structure
        # by looking for key lines from the original
        original_lines = set(l.strip() for l in original.splitlines() if l.strip() and len(l.strip()) > 10)
        proposed_lines = set(l.strip() for l in proposed.splitlines() if l.strip())
        overlap = len(original_lines & proposed_lines)
        preservation_ratio = overlap / max(len(original_lines), 1)

        # If most of the original is preserved (>50%), accept the proposal
        if preservation_ratio > 0.5:
            # Just make sure frontmatter is there
            if has_frontmatter and not proposed.startswith("---"):
                return f"{frontmatter}\n\n{proposed}"
            return proposed

        # The LLM rewrote from scratch — extract only new content and append to original
        # Find lines in proposed that aren't in the original
        new_lines = []
        in_new_section = False
        for line in proposed.splitlines():
            stripped = line.strip()
            if stripped.startswith("#") and stripped not in original_lines:
                in_new_section = True
            if in_new_section:
                new_lines.append(line)

        if new_lines:
            addition = "\n".join(new_lines)
            return f"{original.rstrip()}\n\n{addition}\n"

        # Nothing salvageable — return original unchanged
        return original

    def _call_llm(self, prompt: str) -> str:
        """Call LLM via litellm to generate improved skill content."""
        try:
            response = litellm.completion(
                model=self.model,
                max_tokens=16000,
                messages=[{"role": "user", "content": prompt}],
            )

            text = response.choices[0].message.content or ""

            match = re.search(r"<new_skill>(.*?)</new_skill>", text, re.DOTALL)
            return match.group(1).strip() if match else text.strip()
        except Exception as e:
            if self.run_dir:
                self.run_dir.mkdir(parents=True, exist_ok=True)
                (self.run_dir / f"error_{self._call_count}.txt").write_text(str(e))
            return ""

    def _log(self, candidate, side_infos, components, results):
        """Save proposal interaction for debugging."""
        if not self.run_dir:
            return
        self.run_dir.mkdir(parents=True, exist_ok=True)

        # Count actual items (side_infos is {component: [items...]})
        if isinstance(side_infos, dict):
            items_per_component = {k: len(v) for k, v in side_infos.items()}
            total_items = sum(items_per_component.values())
        else:
            items_per_component = {}
            total_items = len(side_infos)

        log_path = self.run_dir / f"proposal_{self._call_count}.json"
        log_path.write_text(json.dumps({
            "call_number": self._call_count,
            "components": components,
            "num_side_infos": total_items,
            "items_per_component": items_per_component,
            "result_lengths": {k: len(v) for k, v in results.items()},
        }, indent=2))

        # Save the full prompt for debugging
        for component in components:
            current_value = candidate.get(component, "")
            component_data = side_infos.get(component, []) if isinstance(side_infos, dict) else side_infos
            prompt = self._build_prompt(current_value, component_data)
            prompt_path = self.run_dir / f"proposal_{self._call_count}_prompt.txt"
            prompt_path.write_text(prompt)

        # Save the result text so rejected candidates can still be inspected
        for component, value in results.items():
            result_path = self.run_dir / f"proposal_{self._call_count}_result_{component}.md"
            result_path.write_text(value)
