"""
cace_env/server.py
CACEEnvironment — single unified environment.

One episode flow handles everything:
  reset() → seed posts on network → IC spread → enrich all → agent picks + decides
  step()  → three-track reward + spread bonus

No mode switching. No V1/V2 branching. One clean class.
"""

import os, uuid, random
from typing import Optional, List

from openenv.core import Environment, create_fastapi_app

from cace_env.models import CACEAction, CACEObservation, CACEState
from cace_env.dataset import CaseDataset, ACTION_MAP
from cace_env.pipeline import enrich, build_observation
from cace_env.reward import compute_reward

# ── Network (optional — graceful fallback if networkx not installed) ──────────
try:
    import networkx as nx
    _HAS_NX = True
except ImportError:
    _HAS_NX = False

DATASET_PATH    = os.environ.get("DATASET_PATH", "data/all_cases.json")
BATCH_SIZE      = int(os.environ.get("BATCH_SIZE", "20"))   # posts per episode
REVIEW_BUDGET   = int(os.environ.get("REVIEW_BUDGET", "8")) # posts agent must review
NETWORK_STEPS   = int(os.environ.get("NETWORK_STEPS", "3")) # IC cascade steps


# ── Network helpers ───────────────────────────────────────────────────────────

def _build_graph() -> Optional[object]:
    """
    Watts-Strogatz small-world graph approximating SNAP Facebook ego topology.
    ~1000 nodes, avg degree 6, rewiring 0.1.
    Falls back to None if networkx not available.
    """
    if not _HAS_NX:
        return None
    return nx.watts_strogatz_graph(n=1000, k=6, p=0.1, seed=42)


def _ic_spread(G, seed_nodes: List[int], steps: int) -> List[dict]:
    """
    Independent Cascade spread simulation.
    Returns per-seed spread metrics: share_velocity, network_reach, position.
    Falls back to synthetic signals if G is None.
    """
    if G is None or not _HAS_NX:
        # Synthetic spread signals when networkx unavailable
        signals = []
        for i in range(len(seed_nodes)):
            v = random.uniform(0.05, 0.6)
            signals.append({
                "share_velocity": round(v, 3),
                "network_reach": round(v * 0.5, 3),
                "network_position": random.choice(["hub", "bridge", "edge"]),
            })
        return signals

    results = []
    avg_deg = sum(dict(G.degree()).values()) / G.number_of_nodes()

    for seed in seed_nodes:
        active, newly = {seed}, {seed}
        for _ in range(steps):
            nxt = set()
            for node in newly:
                p = 1.0 / max(1, G.degree(node))
                nxt |= {nb for nb in G.neighbors(node)
                        if nb not in active and random.random() < p}
            active |= nxt
            newly = nxt

        reach    = len(active) / G.number_of_nodes()
        velocity = min(1.0, reach * 2.0)
        deg      = G.degree(seed)
        pos      = "hub" if deg > 2*avg_deg else "bridge" if deg > avg_deg else "edge"
        results.append({
            "share_velocity":   round(velocity, 3),
            "network_reach":    round(reach, 3),
            "network_position": pos,
        })
    return results


# ── Environment ───────────────────────────────────────────────────────────────

class CACEEnvironment(Environment[CACEAction, CACEObservation, CACEState]):
    """
    Cultural Context Arbitration Environment — unified V1+V2.

    Episode flow (always the same):

    reset()
      1. Sample BATCH_SIZE posts from dataset
      2. Seed on social graph, run IC spread (3 steps)
      3. Attach spread signals to each post
      4. Pre-enrich all posts via 4 frozen agents (uses cache for speed)
      5. Return batch observation — agent sees ALL posts with signals

    step(action)
      action.selected_indices: which REVIEW_BUDGET posts to review (V2 prioritisation)
      action.action_int:        moderation decision (same for all selected posts)
      → compute 3-track reward + spread bonus per post → return avg reward

    If action.selected_indices is None (simple single-case use):
      → treat action.action_int as decision for the first (primary) case
      → compute 3-track reward without spread bonus
    """

    SUPPORTS_CONCURRENT_SESSIONS = True

    def __init__(self):
        super().__init__()
        self._dataset = CaseDataset(DATASET_PATH)
        self._graph   = _build_graph()

        # Episode state
        self._episode_id: str = ""
        self._batch: List[dict] = []         # [{case, enriched, signals, obs_str}]
        self._step_count: int = 0

        # Metrics
        self._total_episodes: int = 0
        self._correct:        int = 0
        self._rewards:   List[float] = []

    # ── reset ─────────────────────────────────────────────────────────────────

    def reset(
        self,
        seed: Optional[int] = None,
        episode_id: Optional[str] = None,
        **kwargs,
    ) -> CACEObservation:
        """
        Start a new episode.

        1. Sample BATCH_SIZE cases.
        2. Run IC spread on social graph.
        3. Enrich all cases via 4-agent pipeline (from cache when possible).
        4. Return batch observation.
        """
        if seed is not None:
            random.seed(seed)

        self._episode_id = episode_id or str(uuid.uuid4())
        self._step_count = 0
        self._total_episodes += 1

        # 1. Sample batch
        cases = self._dataset.sample_batch(BATCH_SIZE)

        # 2. Network spread signals
        seed_nodes = random.sample(
            list(self._graph.nodes()) if self._graph else list(range(BATCH_SIZE)),
            min(BATCH_SIZE, 1000 if self._graph else BATCH_SIZE)
        )[:BATCH_SIZE]
        signals = _ic_spread(self._graph, seed_nodes, NETWORK_STEPS)

        # 3. Enrich all cases (from pipeline cache when available — fast)
        self._batch = []
        for i, case in enumerate(cases):
            cached = self._dataset.get_cache(case["id"])
            enriched = enrich(
                case["post_text"],
                cache={case["id"]: cached},
                case_id=case["id"]
            )
            # Ensure post_text is always in enriched state for build_observation
            enriched["post_text"] = case["post_text"]
            sig = signals[i] if i < len(signals) else {
                "share_velocity": 0.1, "network_reach": 0.05, "network_position": "edge"
            }
            sig["harm_probability"] = 1.0 if case["board_outcome"] == "REMOVE" else 0.0
            self._batch.append({
                "case":     case,
                "enriched": enriched,
                "signals":  sig,
                "obs_str":  build_observation(enriched, sig),
            })

        # 4. Build unified observation
        obs_str = self._build_observation()

        return CACEObservation(
            observation=obs_str,
            case_id=f"BATCH-{self._episode_id[:8]}",
            language=self._batch[0]["enriched"].get("language", "Unknown"),
            region=self._batch[0]["enriched"].get("region", "Unknown"),
            complexity=self._batch[0]["case"].get("complexity", "medium"),
            culture_flag=self._batch[0]["case"].get("culture_flag", False),
            batch_posts=self._batch_summaries(),
            network_step=0,
            done=False,
            reward=None,
        )

    # ── step ──────────────────────────────────────────────────────────────────

    def step(
        self,
        action: CACEAction,
        timeout_s: Optional[float] = None,
        **kwargs,
    ) -> CACEObservation:
        """
        Apply moderation decisions.

        If action.selected_indices provided:
            → review those REVIEW_BUDGET posts, compute per-post 3-track + spread reward
        Else (single-case fallback):
            → apply decision to first post only, compute 3-track reward
        """
        if not self._batch:
            raise RuntimeError("Call reset() before step().")

        self._step_count += 1
        indices = action.selected_indices

        if indices:
            reward, breakdown = self._step_batch(action.action_str, indices)
        else:
            reward, breakdown = self._step_single(action.action_str)

        self._rewards.append(reward)
        if breakdown.get("correct"):
            self._correct += 1

        primary_case = self._batch[0]["case"]
        return CACEObservation(
            observation=self._build_observation(),
            case_id=f"BATCH-{self._episode_id[:8]}",
            language=self._batch[0]["enriched"].get("language", "Unknown"),
            region=self._batch[0]["enriched"].get("region", "Unknown"),
            complexity=primary_case.get("complexity", "medium"),
            culture_flag=primary_case.get("culture_flag", False),
            mode="batch" if indices else "single",
            done=True,
            reward=reward,
            reward_breakdown={
                **breakdown,
                "ground_truth": primary_case.get("board_outcome", "?"),
                "case_id": primary_case.get("id", "?"),
            },
        )

    def _step_single(self, decision: str) -> tuple[float, dict]:
        """Single-case decision (first post in batch). No spread bonus."""
        item = self._batch[0]
        case = item["case"]
        r, bd = compute_reward(
            decision=decision,
            ground_truth=case["board_outcome"],
            complexity=case["complexity"],
            culture_flag=case.get("culture_flag", False),
            share_velocity=0.0,
            network_reach=0.0,
        )
        return r, bd

    def _step_batch(self, decision: str, indices: List[int]) -> tuple[float, dict]:
        """Batch review: compute reward per selected post, return average."""
        selected = [self._batch[i] for i in indices if i < len(self._batch)]
        total, breakdowns = 0.0, []

        for item in selected[:REVIEW_BUDGET]:
            case = item["case"]
            sig  = item["signals"]
            r, bd = compute_reward(
                decision=decision,
                ground_truth=case["board_outcome"],
                complexity=case["complexity"],
                culture_flag=case.get("culture_flag", False),
                share_velocity=sig["share_velocity"],
                network_reach=sig["network_reach"],
            )
            total += r
            breakdowns.append(bd)

        avg = total / max(1, len(breakdowns))
        correct_count = sum(1 for bd in breakdowns if bd["correct"])
        # Average the three tracks so they show in training logs
        avg_t1 = sum(bd.get("track1_cultural", 0) for bd in breakdowns) / max(1, len(breakdowns))
        avg_t2 = sum(bd.get("track2_harm",     0) for bd in breakdowns) / max(1, len(breakdowns))
        avg_t3 = sum(bd.get("track3_policy",   0) for bd in breakdowns) / max(1, len(breakdowns))
        return avg, {
            "avg_reward":      round(avg, 4),
            "correct":         correct_count == len(breakdowns),
            "correct_count":   correct_count,
            "total_reviewed":  len(breakdowns),
            "track1_cultural": round(avg_t1, 2),
            "track2_harm":     round(avg_t2, 2),
            "track3_policy":   round(avg_t3, 2),
            "per_post":        breakdowns,
        }

    # ── Observation builder ───────────────────────────────────────────────────

    def _build_observation(self) -> str:
        """
        Unified observation: network batch summary + primary case enrichment.
        Agent sees both the spread signals (for prioritisation) and the full
        deliberation context (for the moderation decision).
        """
        # Part 1: Network batch summary (for prioritisation)
        lines = [
            "═══ CULTURAL CONTEXT ARBITRATION ENVIRONMENT ═══",
            f"Episode: {self._episode_id[:8]} | Posts: {BATCH_SIZE} | Review budget: {REVIEW_BUDGET}",
            "",
            "── NETWORK BATCH (select your review queue) ──",
        ]
        for i, item in enumerate(self._batch):
            s   = item["signals"]
            c   = item["case"]
            tag = "⚠️ " if s["harm_probability"] > 0.5 else "   "
            lines.append(
                f"[{i:02d}] {tag}{c['id']} | {item['enriched'].get('language','?')} | "
                f"{item['enriched'].get('region','?')} | "
                f"velocity={s['share_velocity']:.2f} reach={s['network_reach']:.2f} "
                f"pos={s['network_position']}"
            )
            lines.append(f"      {c['post_text'][:90]}...")

        # Part 2: Primary case full deliberation (for decision)
        primary = self._batch[0]
        lines += [
            "",
            "── PRIMARY CASE (full deliberation) ──",
            build_observation(primary["enriched"], primary["signals"]),
            "",
            "── YOUR TASK ──",
            f"1. SELECT {REVIEW_BUDGET} indices to review (comma-separated): e.g. 0,3,5,7,9,11,14,17",
            "2. DECIDE for each selected post:",
            "   ALLOW | REMOVE | ALLOW_WITH_LABEL | ESCALATE | RESTRICT_DISTRIBUTION",
            "",
            "Format: INDICES: 0,3,5,... | DECISION: ALLOW",
        ]
        return "\n".join(lines)

    def _batch_summaries(self) -> List[dict]:
        return [
            {
                "index":           i,
                "case_id":         item["case"]["id"],
                "post_preview":    item["case"]["post_text"][:100],
                "language":        item["enriched"].get("language", "Unknown"),
                "region":          item["enriched"].get("region", "Unknown"),
                "share_velocity":  item["signals"]["share_velocity"],
                "network_reach":   item["signals"]["network_reach"],
                "harm_probability":item["signals"]["harm_probability"],
                "network_position":item["signals"]["network_position"],
                "ground_truth":    item["case"]["board_outcome"],
            }
            for i, item in enumerate(self._batch)
        ]

    # ── state property ────────────────────────────────────────────────────────

    @property
    def state(self) -> CACEState:
        primary = self._batch[0] if self._batch else {}
        avg_50  = (
            sum(self._rewards[-50:]) / min(50, len(self._rewards))
            if self._rewards else 0.0
        )
        return CACEState(
            episode_id=self._episode_id,
            step_count=self._step_count,
            case_id=(primary.get("case") or {}).get("id", ""),
            post_text=(primary.get("case") or {}).get("post_text", "")[:200],
            language=(primary.get("enriched") or {}).get("language", "Unknown"),
            region=(primary.get("enriched") or {}).get("region", "Unknown"),
            policy_clause=(primary.get("enriched") or {}).get("policy_clause", "Unknown"),
            cultural_brief=(primary.get("enriched") or {}).get("cultural_brief", "")[:150],
            challenge_brief=(primary.get("enriched") or {}).get("challenge_brief", "")[:150],
            policy_anchor=(primary.get("enriched") or {}).get("policy_anchor", "")[:150],
            ground_truth=(primary.get("case") or {}).get("board_outcome", ""),
            complexity=(primary.get("case") or {}).get("complexity", "medium"),
            mode="unified",
            total_episodes=self._total_episodes,
            correct_decisions=self._correct,
            accuracy=round(self._correct / max(1, self._total_episodes), 4),
            avg_reward_last_50=round(avg_50, 4),
            network_nodes=self._graph.number_of_nodes() if self._graph else None,
            network_edges=self._graph.number_of_edges() if self._graph else None,
            posts_in_batch=len(self._batch),
            posts_selected=REVIEW_BUDGET,
        )


# ── FastAPI app ───────────────────────────────────────────────────────────────

# Use singleton — OpenEnv creates new instance per request by default
# which breaks stateful environments. We use a single shared instance.
_ENV_INSTANCE = CACEEnvironment()

def env_factory():
    return _ENV_INSTANCE

app = create_fastapi_app(
    env=env_factory,
    action_cls=CACEAction,
    observation_cls=CACEObservation,
    max_concurrent_envs=1,  # single instance = single concurrent session
)

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)