import os
import logging
from typing import Any, Dict

from llama_index.core.agent.workflow import ReActAgent
from llama_index.core.tools import FunctionTool
from llama_index.core.workflow import Context
from llama_index.llms.google_genai import GoogleGenAI

# -----------------------------------------------------------------------------
# Context helper tools ---------------------------------------------------------
# -----------------------------------------------------------------------------

async def write_state(ctx: Context, key: str, value: Any) -> str:
    state_dict = await ctx.get("state")
    state_dict[key] = value
    await ctx.set("state", state_dict)
    return f"state['{key}'] written"

async def read_state(ctx: Context, key: str) -> Any:
    state_dict = await ctx.get("state")
    return state_dict.get(key, "")

write_state_tool = FunctionTool.from_defaults(
    fn=write_state,
    name="write_state",
    description="Store or overwrite a value in the shared workflow state.",
)
read_state_tool = FunctionTool.from_defaults(
    fn=read_state,
    name="read_state",
    description="Retrieve a value from the shared workflow state.",
)

# -----------------------------------------------------------------------------
# Fresh implementation of answer_question -------------------------------------
# -----------------------------------------------------------------------------

def answer_question(question: str) -> str:
    """Return chain‑of‑thought and FINAL ANSWER following strict template."""
    gemini_api_key = os.getenv("GEMINI_API_KEY")
    if not gemini_api_key:
        logging.warning("GEMINI_API_KEY not set – returning fallback answer.")
        return f"Chain of thought: (api key missing)\n\nFINAL ANSWER: {question}"

    meta_prompt = (
        "You are a professional assistant. Respond with two sections:"\
        "\n1. Chain of thought: concise reasoning (3–5 sentences)."\
        "\n2. FINAL ANSWER: the concise answer following these rules:"\
        "\n   • If numeric, no thousands separators or units unless requested."\
        "\n   • If text, as few words as possible, no unnecessary articles."\
        "\n   • If list, comma‑separate applying the above rules."\
        "\n   • Must start exactly with 'FINAL ANSWER:' (uppercase)."\
        f"\n\nQuestion: {question}\n\nAnswer:"
    )

    llm = GoogleGenAI(api_key=gemini_api_key, model="gemini-2.5-pro-preview-03-25", temperature=0.05)
    return llm.complete(meta_prompt).text.strip()

answer_question_tool = FunctionTool.from_defaults(
    fn=answer_question,
    name="answer_question",
    description="Generate reasoning and emit 'FINAL ANSWER: ...' following the strict format rules.",
)

# -----------------------------------------------------------------------------
# System prompt (unchanged) ----------------------------------------------------
# -----------------------------------------------------------------------------

SYNTHESIS_SYSTEM_PROMPT = r"""
You are SynthesisAgent, the final composer in a multi‑agent workflow.
Your goal is to merge validated outputs from specialised agents into a concise
user‑facing answer.

POTENTIAL STATE KEYS TO CONSULT
--------------------------------
objective                 – str  (restated user goal)
plan                      – dict (PlannerAgent JSON plan)
evidence                  – list[str]  (ResearchAgent facts)
calculations              – list[dict] (MathAgent results)
code_outputs              – list[dict] (CodeAgent execution)
image_analysis            – list[dict] (ImageAnalyzerAgent)
figure_interpretation     – list[dict] (FigureInterpretationAgent)
video_analysis            – list[dict] (VideoAnalyzerAgent)
text_analysis             – list[dict] (TextAnalyzerAgent)
role_draft                – str       (RoleAgent draft, optional)
reasoning                 – list[str]  (ReasoningAgent chain‑of‑thought)
validation                – list[dict] (AdvancedValidationAgent)

WORKFLOW
--------
1. Read every relevant key. Create a short internal outline.
2. If contradictions or missing evidence exist, hand off to
   advanced_validation_agent or research_agent.
3. Draft a clear, well‑structured answer (<= 200 words or 7 bullet points).
4. Call the tool `answer_question` with the **user question** to format the
   final output as required.

STYLE
-----
* Formal but approachable language; no internal state leakage.
* Cite numeric values plainly; no inline URLs.
* Prefer paragraph then bullets for details.

HANDOFF POLICY
--------------
Allowed targets when more work required:
  • advanced_validation_agent  – contradictions or doubt
  • research_agent             – missing data
  • reasoning_agent            – reconcile complex logic
  • long_context_management_agent – compress oversized context before answer
  
If your response exceeds the maximum token limit and cannot be completed in a single reply, please conclude your output with the marker [CONTINUE]. In subsequent interactions, I will prompt you with “continue” to receive the next portion of the response.
"""

# -----------------------------------------------------------------------------
# Factory ---------------------------------------------------------------------
# -----------------------------------------------------------------------------

def initialize_synthesis_agent() -> ReActAgent:
    logger = logging.getLogger(__name__)
    logger.info("Initialising SynthesisAgent …")

    gemini_api_key = os.getenv("GEMINI_API_KEY")
    if not gemini_api_key:
        raise ValueError("GEMINI_API_KEY required for SynthesisAgent")

    llm = GoogleGenAI(api_key=gemini_api_key, model="gemini-2.5-pro-preview-03-25", temperature=0.05)

    agent = ReActAgent(
        name="synthesis_agent",
        description=(
            "Aggregates all validated information, resolves residual issues and "
            "produces the final user answer via answer_question, adhering to the "
            "required template."),
        tools=[write_state_tool, read_state_tool, answer_question_tool],
        llm=llm,
        system_prompt=SYNTHESIS_SYSTEM_PROMPT,
        can_handoff_to=[
            "advanced_validation_agent",
            "research_agent",
            "reasoning_agent",
            "long_context_management_agent",
        ],
    )
    return agent

# -----------------------------------------------------------------------------
# Stand‑alone test ------------------------------------------------------------
# -----------------------------------------------------------------------------

if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
    ag = initialize_synthesis_agent()
    print("SynthesisAgent ready.")