"""
IterSynthEngine: Arsitektur Role-Decoupled Iterative Synthesis untuk Deep Search Agents.
Berdasarkan formulasi riset Tencent AI Lab & Zhejiang University (arXiv:2609.29444, Wu et al., 24 September 2026).

Komponen Arsitektur Utama:
1. Role-Decoupled State Machine: Pergantian deterministik antara Planner (\pi_plan) dan Synthesizer (\pi_synth).
2. Evolving Workspace Memory (M_t): Ruang kerja terstruktur yang mengondensasi bukti relevan dan membuang raw observation noise.
3. Bounded Context Constructor: Membatasi panjang riwayat (H_t) agar tidak terjadi context explosion (mengatasi kegagalan 64K).
4. Role-Decoupled Policy Optimization (RDPO) Credit Allocator: Memisahkan keunggulan kelompok (group advantage) Planner vs Synthesizer.
5. Query-Conditioned Dual Tool Bridge: SerpAPI multi-query search dan query-conditioned document reader.
"""

from dataclasses import dataclass, field
from enum import Enum
from typing import List, Dict, Any, Optional, Tuple, Callable
import json
import re

class AgentRole(str, Enum):
    PLANNER = "planner"
    SYNTHESIZER = "synthesizer"

@dataclass
class PlannerAction:
    thought: str
    action_type: str               # "search", "browse", "answer"
    queries: List[str] = field(default_factory=list)
    urls: List[str] = field(default_factory=list)
    final_answer: Optional[str] = None

@dataclass
class WorkspaceMemory:
    round_idx: int
    core_question: str
    known_facts: List[str] = field(default_factory=list)
    unresolved_subgoals: List[str] = field(default_factory=list)
    active_hypotheses: List[str] = field(default_factory=list)
    key_evidence_snippets: List[Dict[str, str]] = field(default_factory=list)

    def render_markdown(self) -> str:
        lines = [
            f"### Workspace State (Round {self.round_idx})",
            f"**Core Target:** {self.core_question}",
            "**Verified Evidence & Entities:**"
        ]
        if not self.known_facts:
            lines.append("- (Belum ada fakta terverifikasi)")
        else:
            for f in self.known_facts:
                lines.append(f"- {f}")
        
        lines.append("**Pending Knowledge Gaps:**")
        if not self.unresolved_subgoals:
            lines.append("- (Semua celah informasi terpenuhi)")
        else:
            for g in self.unresolved_subgoals:
                lines.append(f"- [ ] {g}")
                
        return "\n".join(lines)

class DualRolePromptTemplate:
    @staticmethod
    def build_planner_prompt(question: str, memory: WorkspaceMemory) -> str:
        return f"""[ROLE: PLANNER]
Anda adalah Planner khusus untuk pemecahan masalah mendalam. Tugas Anda:
1. Menganalisis kondisi memori terkini dan menentukan celah bukti kritis selanjutnya.
2. Memilih tindakan: SEARCH (multi-query), BROWSE (URL spesifik), atau FINAL_ANSWER.
JANGAN membuat ringkasan naratif panjang; fokus pada pemilihan aksi deterministik.

Target Pertanyaan: {question}

Kondisi Memori:
{memory.render_markdown()}

Keluarkan instruksi dalam format JSON:
{{"thought": "...", "action": "search"|"browse"|"answer", "queries": [...], "urls": [...], "answer": null}}
"""

    @staticmethod
    def build_synthesizer_prompt(question: str, current_memory: WorkspaceMemory, raw_observation: str) -> str:
        return f"""[ROLE: SYNTHESIZER]
Anda adalah Synthesizer dan Information Filter. Tugas Anda:
1. Ekstraksi fakta konkret, angka, kutipan, dan data baru dari observasi mentah.
2. Eliminasi total noise eksplorasi, boilerplate HTML, dan informasi tidak relevan.
3. Rekonstruksi dan mutasikan Workspace Memory untuk putaran berikutnya.

Target Pertanyaan: {question}

Memori Sebelumnya:
{current_memory.render_markdown()}

Observasi Mentah Lingkungan (Raw Environment Output):
{raw_observation}

Keluarkan hasil sintesis dalam format JSON yang memperbarui fakta terverifikasi dan sisa celah informasi:
{{"verified_facts_added": [...], "resolved_gaps": [...], "remaining_gaps": [...]}}
"""

class RDPOAdvantageCalculator:
    """
    Role-Decoupled Policy Optimization (RDPO) Advantage Calculator.
    Menghitung keunggulan relatif terpisah antara tindakan Planner dan Synthesizer
    untuk mencegah interferensi gradien (credit diffusion).
    """
    @staticmethod
    def compute_group_advantages(
        rollout_outcomes: List[float],           # Terminal reward r_y (1.0 atau 0.0)
        planner_turn_rubrics: List[List[float]], # Turn-level rubric score per trajectory
        synthesizer_turn_rubrics: List[List[float]],
        alpha: float = 0.6,
        beta: float = 0.4
    ) -> Tuple[List[List[float]], List[List[float]]]:
        n = len(rollout_outcomes)
        mean_outcome = sum(rollout_outcomes) / n if n > 0 else 0.0
        std_outcome = (sum((x - mean_outcome) ** 2 for x in rollout_outcomes) / n) ** 0.5 + 1e-8
        
        normalized_outcomes = [(r - mean_outcome) / std_outcome for r in rollout_outcomes]
        
        planner_adv: List[List[float]] = []
        synthesizer_adv: List[List[float]] = []
        
        for i in range(n):
            traj_outcome_adv = normalized_outcomes[i]
            
            # Planner Advantage: gabungan terminal goal + turn-level query precision
            p_turns = [
                alpha * traj_outcome_adv + beta * p_score
                for p_score in planner_turn_rubrics[i]
            ]
            planner_adv.append(p_turns)
            
            # Synthesizer Advantage: gabungan terminal goal + turn-level information density
            s_turns = [
                alpha * traj_outcome_adv + beta * s_score
                for s_score in synthesizer_turn_rubrics[i]
            ]
            synthesizer_adv.append(s_turns)
            
        return planner_adv, synthesizer_adv

class IterSynthRuntime:
    """
    Runtime Eksekusi IterSynth:
    Menjalankan siklus pencarian mendalam dengan pergantian peran Planner-Synthesizer
    di atas bounded workspace memory terisolasi.
    """
    def __init__(
        self,
        llm_caller: Callable[[str], str],
        search_tool: Callable[[List[str]], str],
        browse_tool: Callable[[List[str]], str],
        max_rounds: int = 12
    ):
        self.llm = llm_caller
        self.search_tool = search_tool
        self.browse_tool = browse_tool
        self.max_rounds = max_rounds

    def execute_deep_search(self, question: str) -> Dict[str, Any]:
        memory = WorkspaceMemory(
            round_idx=0,
            core_question=question,
            known_facts=[],
            unresolved_subgoals=[f"Identifikasi entitas kunci dan data primer untuk: {question}"]
        )
        
        trajectory_log = []
        
        for round_idx in range(1, self.max_rounds + 1):
            memory.round_idx = round_idx
            
            # 1. Fase Planner (Konteks Terbatas)
            planner_prompt = DualRolePromptTemplate.build_planner_prompt(question, memory)
            planner_response = self.llm(planner_prompt)
            
            # Parse tindakan Planner
            try:
                plan_data = json.loads(re.search(r'\{.*\}', planner_response, re.DOTALL).group(0))
            except Exception:
                plan_data = {"action": "search", "queries": [question], "thought": "Fallback search"}
                
            action_type = plan_data.get("action", "search")
            trajectory_log.append({"round": round_idx, "role": "planner", "plan": plan_data})
            
            # Evaluasi terminasi dini jika Planner telah yakin dengan jawaban akhir
            if action_type == "answer" or plan_data.get("answer"):
                return {
                    "status": "success",
                    "rounds": round_idx,
                    "final_answer": plan_data.get("answer"),
                    "final_memory": memory,
                    "trajectory_log": trajectory_log
                }
                
            # 2. Eksekusi Tool Lingkungan
            if action_type == "search":
                queries = plan_data.get("queries", [question])
                raw_obs = self.search_tool(queries)
            elif action_type == "browse":
                urls = plan_data.get("urls", [])
                raw_obs = self.browse_tool(urls)
            else:
                raw_obs = "Aksi tidak dikenali."
                
            # 3. Fase Synthesizer (Ekstraksi Bukti & Denoising Bersih)
            synth_prompt = DualRolePromptTemplate.build_synthesizer_prompt(question, memory, raw_obs)
            synth_response = self.llm(synth_prompt)
            
            try:
                synth_data = json.loads(re.search(r'\{.*\}', synth_response, re.DOTALL).group(0))
            except Exception:
                synth_data = {"verified_facts_added": [raw_obs[:200]], "remaining_gaps": []}
                
            trajectory_log.append({"round": round_idx, "role": "synthesizer", "synthesis": synth_data})
            
            # Mutasi Memori Ruang Kerja (Bounded Update - Tanpa Akumulasi Riwayat Mentah)
            for fact in synth_data.get("verified_facts_added", []):
                if fact not in memory.known_facts:
                    memory.known_facts.append(fact)
                    
            memory.unresolved_subgoals = synth_data.get("remaining_gaps", [])
            
        # Jika mencapai batas round maksimal, minta sintesis jawaban akhir
        final_prompt = f"Berdasarkan seluruh fakta terverifikasi berikut:\n{memory.render_markdown()}\nBerikan jawaban komprehensif dan akurat untuk: {question}"
        final_answer = self.llm(final_prompt)
        
        return {
            "status": "max_rounds_reached",
            "rounds": self.max_rounds,
            "final_answer": final_answer,
            "final_memory": memory,
            "trajectory_log": trajectory_log
        }