"""ALFWorld environment adapter for ReflACT.

Connects the ReflACT training loop to ALFWorld by implementing
:class:`~skillopt.envs.base.EnvAdapter`.
"""
from __future__ import annotations

from dataclasses import dataclass
import json
import os

from skillopt.gradient.deep_probe import generate_deep_probe_instruction
from skillopt.datasets.base import BatchSpec
from skillopt.envs.base import EnvAdapter
from skillopt.envs.alfworld.dataloader import ALFWorldDataLoader
from skillopt.envs.alfworld.rollout import (
    build_alfworld_env,
    run_alfworld_batch,
    TASKS,
)
from skillopt.gradient.reflect import run_minibatch_reflect
from skillopt.utils import compute_score


@dataclass(frozen=True)
class ALFWorldBatchRun:
    """Lazy ALFWorld batch description.

    The adapter materializes this in rollout chunks so a large evaluation set
    does not keep every ALFWorld simulator open at once.
    """

    env_num: int
    eval_dataset: str
    seed: int
    is_train: bool
    workers: int
    specific_gamefiles: list[str] | None = None
    result_ids: list[str] | None = None
    items: list[dict] | None = None

    def __iter__(self):
        return iter(self.items or [])

    def __len__(self) -> int:
        return int(self.env_num or 0)


class ALFWorldAdapter(EnvAdapter):
    """ALFWorld environment adapter.

    Parameters
    ----------
    max_steps : int
        Maximum steps per ALFWorld episode (default 50).
    max_api_workers : int
        Maximum concurrent API calls during rollout (default 8).
    analyst_workers : int
        Parallel workers for analyst stage (default 16).
    failure_only : bool
        If True, only run error analyst (skip success analyst).
    minibatch_size : int
        Trajectories per analyst group, M (default 8).
    edit_budget : int
        Maximum edits per minibatch, L (default 4).
    """

    def __init__(
        self,
        split_dir: str = "",
        data_path: str = "",
        split_mode: str = "split_dir",
        split_ratio: str = "2:1:7",
        split_seed: int = 42,
        split_output_dir: str = "",
        seed: int = 42,
        limit: int = 0,
        train_size: int = 0,
        max_steps: int = 50,
        workers: int = 8,
        max_api_workers: int = 8,
        analyst_workers: int = 16,
        failure_only: bool = False,
        minibatch_size: int = 8,
        edit_budget: int = 4,
        use_deep_reflect: bool = False,
        deep_reflect_failures: int = 4,
        deep_reflect_successes: int = 2,
    ) -> None:
        self.max_steps = max_steps
        self.workers = max(int(workers or 1), 1)
        self.max_api_workers = max_api_workers
        self.analyst_workers = analyst_workers
        self.failure_only = failure_only
        self.minibatch_size = minibatch_size
        self.edit_budget = edit_budget
        self.use_deep_reflect = use_deep_reflect
        self.deep_reflect_failures = deep_reflect_failures
        self.deep_reflect_successes = deep_reflect_successes
        self.dataloader = ALFWorldDataLoader(
            split_dir=split_dir,
            data_path=data_path,
            split_mode=split_mode,
            split_ratio=split_ratio,
            split_seed=split_seed,
            split_output_dir=split_output_dir,
            seed=seed,
            limit=limit,
            train_size=train_size,
        )
        self._traj_cache: dict[str, dict | None] = {}

    def setup(self, cfg: dict) -> None:
        super().setup(cfg)
        self.dataloader.setup(cfg)

    def _load_traj_data(self, item: dict) -> dict | None:
        gamefile = str(item.get("gamefile") or "").strip()
        if not gamefile:
            return None
        if gamefile in self._traj_cache:
            return self._traj_cache[gamefile]

        traj_path = os.path.join(os.path.dirname(gamefile), "traj_data.json")
        try:
            with open(traj_path, encoding="utf-8") as f:
                data = json.load(f)
        except Exception:
            data = None
        self._traj_cache[gamefile] = data
        return data

    @staticmethod
    def _unique_lines(values: list[str], *, limit: int = 0) -> list[str]:
        lines: list[str] = []
        seen: set[str] = set()
        for raw in values:
            line = str(raw or "").strip()
            if not line or line in seen:
                continue
            seen.add(line)
            lines.append(line)
            if limit > 0 and len(lines) >= limit:
                break
        return lines

    @staticmethod
    def _format_high_pddl(high_pddl: list[dict]) -> list[str]:
        steps: list[str] = []
        for idx, step in enumerate(high_pddl or [], start=1):
            discrete = step.get("discrete_action") or {}
            action = str(discrete.get("action") or "").strip()
            args = [str(arg).strip() for arg in (discrete.get("args") or []) if str(arg).strip()]
            if action and args:
                text = f"{action}({', '.join(args)})"
            elif action:
                text = action
            else:
                planner_action = step.get("planner_action") or {}
                text = str(planner_action.get("action") or "").strip()
            if text:
                steps.append(f"{idx}. {text}")
        return steps

    def _build_reference_bundle(self, item: dict) -> dict:
        data = self._load_traj_data(item)
        if not data:
            return {}

        anns = ((data.get("turk_annotations") or {}).get("anns") or [])
        task_descs = self._unique_lines(
            [ann.get("task_desc", "") for ann in anns],
            limit=3,
        )
        high_descs = self._unique_lines(
            [step for ann in anns for step in (ann.get("high_descs") or [])],
            limit=12,
        )
        pddl_params = {
            key: value
            for key, value in (data.get("pddl_params") or {}).items()
            if value not in ("", None, [], {})
        }
        scene = data.get("scene") or {}
        scene_summary = {
            key: scene.get(key)
            for key in ("floor_plan", "scene_num", "dirty_and_empty")
            if scene.get(key) not in ("", None, [], {})
        }
        high_pddl = self._format_high_pddl((data.get("plan") or {}).get("high_pddl") or [])
        task_type = str(data.get("task_type") or item.get("task_type") or "").strip()
        return {
            "task_type": task_type,
            "task_descs": task_descs,
            "high_descs": high_descs,
            "pddl_params": pddl_params,
            "high_pddl": high_pddl,
            "scene_summary": scene_summary,
        }

    def build_reference_text(self, item: dict) -> str:
        bundle = self._build_reference_bundle(item)
        if not bundle:
            return ""

        parts: list[str] = []
        if bundle["task_type"]:
            parts.append(f"## Reference Task Type\n{bundle['task_type']}")
        if bundle["task_descs"]:
            parts.append(
                "## Reference Human Task Descriptions\n"
                + "\n".join(f"- {line}" for line in bundle["task_descs"])
            )
        if bundle["high_descs"]:
            parts.append(
                "## Reference Human High-Level Steps\n"
                + "\n".join(f"{idx}. {line}" for idx, line in enumerate(bundle["high_descs"], start=1))
            )
        if bundle["pddl_params"]:
            parts.append(
                "## Reference PDDL Params\n"
                + "\n".join(f"- {key}: {value}" for key, value in bundle["pddl_params"].items())
            )
        if bundle["high_pddl"]:
            parts.append(
                "## Reference Planner High-Level Plan\n" + "\n".join(bundle["high_pddl"])
            )
        if bundle["scene_summary"]:
            parts.append(
                "## Reference Scene Summary\n"
                + "\n".join(f"- {key}: {value}" for key, value in bundle["scene_summary"].items())
            )
        return "\n\n".join(parts)

    def get_reference_metadata(self, item: dict) -> dict:
        bundle = self._build_reference_bundle(item)
        if not bundle:
            return {"fields": [], "preview": ""}

        fields: list[str] = []
        previews: list[str] = []
        if bundle["task_type"]:
            fields.append("task_type")
            previews.append(f"[task_type] {bundle['task_type']}")
        if bundle["task_descs"]:
            fields.append("task_desc")
            previews.append("[task_desc]\n" + "\n".join(bundle["task_descs"][:2]))
        if bundle["high_descs"]:
            fields.append("high_descs")
            previews.append("[high_descs]\n" + "\n".join(bundle["high_descs"][:3]))
        if bundle["pddl_params"]:
            fields.append("pddl_params")
            previews.append(
                "[pddl_params]\n"
                + "\n".join(
                    f"{key}: {value}" for key, value in list(bundle["pddl_params"].items())[:4]
                )
            )
        if bundle["high_pddl"]:
            fields.append("plan.high_pddl")
            previews.append("[plan.high_pddl]\n" + "\n".join(bundle["high_pddl"][:3]))
        if bundle["scene_summary"]:
            fields.append("scene")
            previews.append(
                "[scene]\n"
                + "\n".join(
                    f"{key}: {value}" for key, value in bundle["scene_summary"].items()
                )
            )
        return {
            "fields": fields,
            "preview": "\n\n".join(previews)[:600],
        }

    @staticmethod
    def _infer_dataset_from_gamefile(gamefile: str) -> tuple[str, bool]:
        path = str(gamefile or "")
        if "/valid_seen/" in path:
            return "eval_in_distribution", False
        if "/valid_unseen/" in path:
            return "eval_out_of_distribution", False
        return "train", True

    def get_dataloader(self):
        return self.dataloader

    def _comparison_items(self, items: list[dict]) -> list[dict]:
        enriched: list[dict] = []
        for item in items:
            row = dict(item)
            bundle = self._build_reference_bundle(row)
            if bundle.get("task_descs"):
                row["task_description"] = bundle["task_descs"][0]
            elif bundle.get("task_type"):
                row["task_description"] = bundle["task_type"]
            enriched.append(row)
        return enriched

    def requires_ray(self) -> bool:
        return False

    def build_env_from_batch(self, batch: BatchSpec, **kwargs):
        gamefiles = list(batch.metadata.get("gamefiles") or [])
        result_ids = list(batch.metadata.get("result_ids") or [])
        items = self._comparison_items(list(batch.payload or []))
        return ALFWorldBatchRun(
            env_num=batch.batch_size,
            eval_dataset=batch.metadata.get("eval_dataset", batch.split),
            seed=batch.seed,
            is_train=batch.metadata.get("is_train", batch.phase == "train"),
            specific_gamefiles=gamefiles or None,
            result_ids=result_ids or None,
            items=items,
            workers=self.workers,
        )

    def build_train_env(self, batch_size: int, seed: int, **kwargs):
        batch = self.dataloader.build_train_batch(batch_size=batch_size, seed=seed, **kwargs)
        return self.build_env_from_batch(batch, **kwargs)

    def build_eval_env(self, env_num: int, split: str, seed: int, **kwargs):
        batch = self.dataloader.build_eval_batch(env_num=env_num, split=split, seed=seed, **kwargs)
        return self.build_env_from_batch(batch, **kwargs)

    def rollout(
        self,
        env_manager,
        skill_content: str,
        out_dir: str,
        **kwargs,
    ) -> list[dict]:
        results_path = os.path.join(out_dir, "results.jsonl")
        os.makedirs(out_dir, exist_ok=True)

        # Resume support
        if os.path.exists(results_path):
            existing: list[dict] = []
            with open(results_path) as f:
                for line in f:
                    try:
                        existing.append(json.loads(line))
                    except Exception:
                        pass
            if existing:
                return existing

        if isinstance(env_manager, ALFWorldBatchRun):
            results = self._run_batch(
                env_manager,
                skill_content=skill_content,
                out_dir=out_dir,
            )
        else:
            results = run_alfworld_batch(
                env_manager=env_manager,
                skill_content=skill_content,
                max_steps=self.max_steps,
                out_root=out_dir,
                max_api_workers=self.max_api_workers,
                result_ids=getattr(env_manager, "_skillopt_result_ids", None),
            )

        with open(results_path, "w") as f:
            for r in results:
                f.write(json.dumps(r, ensure_ascii=False) + "\n")

        return results

    @staticmethod
    def _close_env(env_manager) -> None:
        close = getattr(env_manager, "close", None)
        if callable(close):
            close()

    def _run_batch(
        self,
        batch: ALFWorldBatchRun,
        skill_content: str,
        out_dir: str,
        *,
        diagnostic_mode: bool = False,
        diagnostic_instruction: str = "",
    ) -> list[dict]:
        total = int(batch.env_num or 0)
        if total <= 0:
            return []

        workers = max(1, min(int(batch.workers or self.workers), total))
        if total > workers:
            print(
                f"    [alfworld rollout] episodes={total} "
                f"env_workers={workers} chunks={(total + workers - 1) // workers}"
            )

        all_results: list[dict] = []
        for start in range(0, total, workers):
            chunk_size = min(workers, total - start)
            chunk_gamefiles = (
                batch.specific_gamefiles[start:start + chunk_size]
                if batch.specific_gamefiles
                else None
            )
            chunk_ids = (
                batch.result_ids[start:start + chunk_size]
                if batch.result_ids
                else [f"env_{idx:03d}" for idx in range(start, start + chunk_size)]
            )
            chunk_env = build_alfworld_env(
                env_num=chunk_size,
                eval_dataset=batch.eval_dataset,
                seed=batch.seed + start,
                is_train=batch.is_train,
                specific_gamefiles=chunk_gamefiles,
            )
            try:
                chunk_results = run_alfworld_batch(
                    env_manager=chunk_env,
                    skill_content=skill_content,
                    max_steps=self.max_steps,
                    out_root=out_dir,
                    max_api_workers=min(self.max_api_workers, chunk_size),
                    diagnostic_mode=diagnostic_mode,
                    diagnostic_instruction=diagnostic_instruction,
                    result_ids=chunk_ids,
                )
            finally:
                self._close_env(chunk_env)
            all_results.extend(chunk_results)
        return all_results

    def reflect(
        self,
        results: list[dict],
        skill_content: str,
        out_dir: str,
        **kwargs,
    ) -> list[dict | None]:
        prediction_dir = kwargs.get("prediction_dir", os.path.join(out_dir, "predictions"))
        patches_dir = kwargs.get("patches_dir", os.path.join(out_dir, "patches"))
        random_seed = kwargs.get("random_seed")
        step_buffer_context = kwargs.get("step_buffer_context", "")
        meta_skill_context = kwargs.get("meta_skill_context", "")

        return run_minibatch_reflect(
            results=results,
            skill_content=skill_content,
            prediction_dir=prediction_dir,
            patches_dir=patches_dir,
            workers=self.analyst_workers,
            failure_only=self.failure_only,
            minibatch_size=self.minibatch_size,
            edit_budget=self.edit_budget,
            random_seed=random_seed,
            error_system=self.get_error_minibatch_prompt(),
            success_system=self.get_success_minibatch_prompt(),
            step_buffer_context=step_buffer_context,
            meta_skill_context=meta_skill_context,
        )

    def deep_reflect(
        self,
        results: list[dict],
        skill_content: str,
        out_dir: str,
        **kwargs,
    ) -> list[dict | None]:
        if not self.use_deep_reflect:
            return []

        prediction_dir = kwargs.get("prediction_dir", os.path.join(out_dir, "predictions"))
        random_seed = kwargs.get("random_seed")
        step_buffer_context = kwargs.get("step_buffer_context", "")
        meta_skill_context = kwargs.get("meta_skill_context", "")
        selected_items = self.select_representative_items(
            results,
            results,
            n_failures=self.deep_reflect_failures,
            n_successes=self.deep_reflect_successes,
            seed=random_seed,
        )
        if not selected_items:
            return []

        selected_ids = {str(item["id"]) for item in selected_items}
        selected_results = [row for row in results if str(row.get("id")) in selected_ids]
        selected_examples = self.attach_reference_context(selected_results, selected_items)

        field_counts: dict[str, int] = {}
        selected_metadata: list[dict] = []
        for item in selected_items:
            meta = self.get_reference_metadata(item)
            for field in meta["fields"]:
                field_counts[field] = field_counts.get(field, 0) + 1
            selected_metadata.append({
                "id": str(item["id"]),
                "task_type": str(item.get("task_type") or "alfworld"),
                "gamefile": str(item.get("gamefile") or ""),
                "reference_fields": meta["fields"],
                "reference_preview": meta["preview"],
            })

        deep_dir = os.path.join(out_dir, "deep_reflect")
        rollout_dir = os.path.join(deep_dir, "rollout")
        patches_dir = os.path.join(deep_dir, "patches")
        os.makedirs(deep_dir, exist_ok=True)
        field_summary = ", ".join(
            f"{field}({count}/{len(selected_items)})"
            for field, count in sorted(field_counts.items())
        ) or "none"
        print(
            f"    [2b/6 DEEP REFLECT setup] selected={len(selected_items)} "
            f"reference_fields={field_summary}"
        )
        probe = generate_deep_probe_instruction(
            skill_content=skill_content,
            items=selected_examples,
            prediction_dir=prediction_dir,
            system_prompt=self.get_deep_probe_prompt(),
            step_buffer_context=step_buffer_context,
            meta_skill_context=meta_skill_context,
            output_requirements=[
                "- Some trajectories may include a hidden Reference block. Use it to target the student's latent subgoal, missing precondition, or next-step intent, but do not reveal or paraphrase that reference to the student.",
                "- The instruction must request a brief diagnostic readout inside the existing <think>...</think> block.",
                "- The student must still output exactly one admissible action inside <action>...</action>.",
                "- Do not ask for exhaustive inventories, full plans, or long chain-of-thought.",
                "- The instruction text should be ready to append directly to the student's prompt.",
            ],
        )
        if not probe:
            return []

        with open(os.path.join(deep_dir, "probe.json"), "w", encoding="utf-8") as f:
            json.dump(
                {
                    **probe,
                    "reference_summary": {
                        "selected_count": len(selected_items),
                        "field_counts": field_counts,
                    },
                    "selected_examples": selected_metadata,
                },
                f,
                ensure_ascii=False,
                indent=2,
            )

        gamefiles = [str(item.get("gamefile") or "") for item in selected_items]
        if any(not gamefile for gamefile in gamefiles):
            return []
        eval_dataset, is_train = self._infer_dataset_from_gamefile(gamefiles[0])
        deep_env = ALFWorldBatchRun(
            env_num=len(selected_items),
            eval_dataset=eval_dataset,
            seed=random_seed or 42,
            is_train=is_train,
            specific_gamefiles=gamefiles,
            workers=min(self.workers, max(len(selected_items), 1)),
            result_ids=[str(item["id"]) for item in selected_items],
        )
        deep_results = self._run_batch(
            deep_env,
            skill_content=skill_content,
            out_dir=rollout_dir,
            diagnostic_mode=True,
            diagnostic_instruction=probe["probe_instruction"],
        )
        deep_results = self.attach_reference_context(deep_results, selected_items)
        return run_minibatch_reflect(
            results=deep_results,
            skill_content=skill_content,
            prediction_dir=os.path.join(rollout_dir, "predictions"),
            patches_dir=patches_dir,
            workers=self.analyst_workers,
            failure_only=self.failure_only,
            minibatch_size=self.minibatch_size,
            edit_budget=self.edit_budget,
            random_seed=random_seed,
            error_system=self.get_error_minibatch_prompt(),
            success_system=self.get_success_minibatch_prompt(),
            step_buffer_context=step_buffer_context,
            meta_skill_context=meta_skill_context,
        )

    def get_task_types(self) -> list[str]:
        return list(TASKS)
