from __future__ import annotations

import json
import logging
from pathlib import Path
from typing import Any

from autogen_agentchat.agents import AssistantAgent
from autogen_core.tools import FunctionTool

from prompt_auto_lab.agents.utils import build_metrics_context
from prompt_auto_lab.core.optimizer_papyrus_client import OptimizerPapyrusClient
from prompt_auto_lab.core.protocol import BraintrustClient, PicassoClient
from prompt_auto_lab.orchestrator.state_store import SharedState
from pydantic import BaseModel
from autogen_agentchat.messages import StructuredMessage
from autogen_core.model_context import BufferedChatCompletionContext

logger = logging.getLogger(__name__)

class AgentResponse(BaseModel):
    thoughts: str
    response: str


class OptimizedPrompt:
    def __init__(
        self,
        model_client: Any,
        client: PicassoClient,
        shared_state: SharedState,
        braintrust_client: BraintrustClient | None = None,
        analysis_agent_client: Any | None = None,
        allowed_tools: set[str] | None = None,
        name: str = "PromptOptimizer",
        AgentResponse: Any = None,
        use_structured_output: bool = True,
    ):
        self.client = client
        self.model_client = model_client
        self.braintrust_client = braintrust_client
        tools = []
        allowed = allowed_tools or {
            "read_analysis_txt",
            "read_prompt_template",
            "write_prompt_files",
            "get_prompt_optimization_guidelines",
        }
        
        # Read tool_description.txt and include it in system_message
        tool_description_content = ""
        try:
            # optimizer.py -> agents -> prompt_auto_lab -> autoPromptLab2
            repo_root = Path(__file__).resolve().parents[2]
            tool_desc_path = repo_root / "tool_description.txt"
            if tool_desc_path.exists():
                with open(tool_desc_path, "r", encoding="utf-8") as f:
                    tool_description_content = f.read()
        except Exception as e:
            # If file not found or error reading, continue without it
            pass
        
        if "read_analysis_txt" in allowed and self.braintrust_client:
            tools.append(FunctionTool(self.braintrust_client.read_analysis_txt, "Read analysis artifact for a given experiment_id. Returns data as a list containing JSON objects from all lines.", strict=True))
        if "read_prompt_template" in allowed:
            tools.append(FunctionTool(client.read_prompt_template, "Read prompt files from Picasso repo. Input: JSON array of file paths relative to repo (e.g., '[\"Service.Chat/Inference/Prompts/Templates/responding-common/idiolect.liquid\"]' or '[\"Orchestration/Extensions/WebSearch/LiquidTemplates/actions/Search.liquid\"]'). Returns file contents.", strict=True))
        if "liquid_check" in allowed:
            tools.append(FunctionTool(client.liquid_check, "Run Liquid check on prompt", strict=True))
        if "make_diff" in allowed:
            tools.append(FunctionTool(client.make_diff, "Generate unified diff", strict=True))
        if "write_prompt_files" in allowed:
            async def review_and_write_prompt_files(modifications_json: str) -> dict[str, Any]:
                """Apply modifications with LLM dedup review before writing."""
                reviewed_json = await self._review_modifications(modifications_json)
                return await client.write_prompt_files(reviewed_json)

            tools.append(FunctionTool(review_and_write_prompt_files, "Apply modifications to prompt files using modification instructions. Input: JSON with file paths as keys, each value is a list of modification operations (add_after, add_before, replace).", strict=True))
        
        if "read_memory_summaries" in allowed and self.braintrust_client:
            tools.append(FunctionTool(self.braintrust_client.read_memory_summaries, "Read memory summaries from memory.jsonl file. Returns all historical memory summaries from previous iterations, sorted by iteration_id.", strict=True))

        if "read_priority_templates" in allowed and self.braintrust_client:
            tools.append(FunctionTool(self.braintrust_client.read_priority_templates, "Read priority templates file generated by LiquidStructureAgent. Returns 2-3 priority templates identified from bad case analysis that should be focused on for optimization.", strict=True))

        if "get_prompt_optimization_guidelines" in allowed and analysis_agent_client is not None:
            async def get_prompt_optimization_guidelines(goal_description: str) -> dict[str, Any]:
                """Retrieve relevant prompt optimization guidelines for this goal."""
                return await analysis_agent_client.find_relevant_skills(
                    goal_description, model_client
                )
            tools.append(FunctionTool(
                get_prompt_optimization_guidelines,
                "Retrieve relevant Picasso prompt writing and optimization guidelines for a given "
                "prompt improvement goal. Returns structured guidance (e.g. meta-prompt standards, "
                "failure taxonomy) most relevant to the described optimization goal.",
                strict=True,
            ))

        tools_description = ", ".join([f"{tool.name}: {tool._description}" for tool in tools])
        
        # Build tool description section for system_message
        tool_description_section = ""
        if tool_description_content:
            tool_description_section = (
                "\nTOOL DESCRIPTIONS (Reference for prompt.liquid modifications):\n"
                "The following tool definitions are available. Use this information when modifying prompt.liquid to add tool-related instructions:\n"
                f"{tool_description_content}\n"
                "When modifying prompt.liquid based on tool-related analysis, refer to these tool definitions to understand tool capabilities and usage.\n\n"
            )

        # Build metrics context using shared utility
        key_metrics_note = build_metrics_context(self.braintrust_client, context_type="optimizer")

        system_message = (
            "You are PromptOptimizer.\n"
            "Goal: Improve prompt files based on bad case analysis and long-term memory insights, write changes.\n\n"
            "Please always read read_analysis_txt even though experiment_id you have procced before.\n\n"
            f"{key_metrics_note}"
            "LONG-TERM MEMORY:\n"
            "- Use read_memory_summaries to read historical memory summaries from memory.jsonl file\n"
            "- This reads directly from persistent storage (memory.jsonl), not from shared_state\n"
            "- IMPORTANT: Only call read_memory_summaries if memory_summary_path is not empty. If memory_summary_path is empty, skip this tool call.\n"
            "- Memory summaries contain:\n"
            "  * effectiveness: json format, it including two parts: 1.actual overall score delta. 2.bad case overall_score delta (positive = better, negative = worse)\n"
            "  * change_summary: What changed in previous iterations\n"
            "  * impact_analysis: How changes affected scores\n"
            "  * lessons_learned: Key insights from past optimizations\n"
            "  * git_diff: Diff of prompt file changes made\n"
            "  * recommendations: Actionable suggestions based on historical patterns\n"
            # "- IMPORTANT:"
            #  "- IMPORTANT: Memory is based on LLM judge's score changes, which have inherent variability and uncertainty\n"
            
            "- CRITICAL: MEMORY INTERPRETATION\n"
            "  * lessons_learned describe NOT ONLY what worked, but also which DIRECTIONS caused harm.\n"
            "  * If lessons_learned explicitly warn against over-restriction, over-generalization, or excessive constraints,\n"
            "  * you MUST treat these as hard boundaries, not optional suggestions.\n"
            "  * You MUST NOT introduce changes that move further in a direction that lessons_learned identify as harmful,\n"
            "  * even if that direction previously improved some metrics.\n"
            "  * Memory comparisons happen BEFORE the revert (during the memory stage)\n"
            "  * The memory reflects the effectiveness of changes that were later reverted\n"
            "  * The current code state (what you see when reading prompt files) is AFTER the revert\n"
            "  * Memory is a historical reference about what was tried, not about the current code state\n"
            "  * Use memory to understand what was attempted and why it didn't work, but don't expect to see those changes in the current files\n"
            "- Use this memory to:\n"
            "  * Look for patterns across multiple iterations rather than relying on single data points\n"
            "  * Avoid repeating ineffective changes\n"
            "  * Apply patterns that worked well before\n"
            "  * Learn from past mistakes (including reverted changes)\n"
            "  * Build on successful strategies\n"
            "- Review all historical summaries to identify patterns across multiple iterations\n"
            "- Prioritize recommendations from iterations with high effectiveness_score\n"
            "- STRATEGY: If previous changes were ineffective but analyzer still reports the same problem:\n"
            "  * Consider rewriting the concept with different wording or phrasing\n"
            "  * Try alternative approaches to address the same underlying issue\n"
            "  * Experiment with different ways to convey the same principle\n\n"
            f"{tool_description_section}"
            "CRITICAL - GENERALIZATION RULE:\n"
            "- The bad case analysis may be from a SPECIFIC task (e.g., page summary, code generation)\n"
            # "- You MUST abstract the suggestions into GENERAL principles that apply to ALL chat scenarios\n"
            "- You MUST abstract suggestions into GENERAL principles, but prefer CONDITIONAL or SCOPED rules"
            " over absolute global prohibitions unless explicitly required by safety or correctness.\n\n"


            "- DO NOT add task-specific rules (e.g., 'when summarizing pages...')\n"
            "- DO add general principles (e.g., 'base responses on provided content, avoid adding ungrounded information')\n"
            "- Example: 'groundedness issue in page summary' → general rule: 'stick to user-provided content unless explicitly asked to expand'\n\n"
            f"Available tools: {tools_description}\n\n"
            "Workflow:\n"
            "0. get_prompt_optimization_guidelines - (Optional) Retrieve relevant optimization guidelines\n"
            "   - Call with a brief description of the goal (e.g. 'fix tool-call grounding issue')\n"
            "   - Returns Picasso prompt writing standards and taxonomy relevant to your optimization\n"
            "   - Use this to align your changes with official prompt engineering guidelines\n"
            "1. read_priority_templates - Read the priority templates file (from priority_templates_path)\n"
            "   - This tells you which 2-3 templates to focus on based on bad case analysis\n"
            "   - Each template includes: path, purpose, DO/DON'T modify guidelines, and suggested changes\n"
            "   - IMPORTANT: Extract the 'path' field from each priority template for use in step 4\n"
            "2. read_memory_summaries - Read historical memory summaries from memory.jsonl file\n"
            "   - This reads all previous memory summaries from persistent storage\n"
            "   - Review lessons_learned and recommendations from past iterations\n"
            "   - Use this to guide your optimization strategy\n"
            # "   - SKIP THIS STEP if memory_summary_path is empty (memory agent is disabled)\n"
            "3. read_analysis_txt - Read analysis results using experiment_id\n"
            "4. read_prompt_template - Read the priority template files from Picasso repo\n"
            "   - Extract paths from step 1's priority_templates (each has a 'path' field)\n"
            "   - Pass these paths as a JSON array: '[\"Service.Chat/...\", \"Orchestration/Extensions/...\"]'\n"
            "   - This reads ONLY the files you need to modify\n"
            "5. Analyze bad cases and ABSTRACT into general principles\n"
            "   - Consider memory insights when analyzing: what worked/didn't work before?\n"
            "   - Apply lessons learned from previous iterations\n"
            "   - Focus changes on the priority templates identified in step 1\n"
            "6. Generate modification instructions for the priority template file(s)\n"
            "   - Use the file paths from priority_templates as keys in the modification JSON\n"
            "   - ONLY modify files listed in priority_templates\n"
            "   - If the analysis involves tool-related issues, refer to the TOOL DESCRIPTIONS section above to understand tool capabilities\n"
            "   - Place tool-related modifications in prompt.liquid\n"
            "7. write_prompt_files - Apply modifications to files using modification instructions\n\n"
            "CRITICAL - MODIFICATION INSTRUCTIONS FORMAT:\n"
            "Use the file path (from priority_templates) as the key:\n"
            '{\n'
            '  "Service.Chat/Inference/Prompts/Templates/responding-common/idiolect.liquid": [\n'
            '    {"operation": "add_after", "search": "existing text", "content": "\\n- New rule here\\n"},\n'
            '    {"operation": "replace", "search": "old text", "content": "new text"}\n'
            '  ],\n'
            '  "Orchestration/Extensions/WebSearch/LiquidTemplates/actions/Search.liquid": [...]\n'
            '}\n\n'
            "JSON FORMATTING RULES FOR write_prompt_files:\n"
            "- The modifications_json parameter must be a SINGLE, COMPLETE, VALID JSON string.\n"
            "- The JSON string must start with '{' and end with '}' with no extra characters after.\n"
            "- All special characters must be properly escaped: use \\\" for quotes, \\\\ for backslashes, \\n for newlines.\n"
            "- The JSON string must be a complete object: {\"file_path\": [operations], ...}\n"
            "- DO NOT add extra quotes, brackets, or any other characters after the closing '}'.\n"
            "- Before calling write_prompt_files, verify your JSON string is valid by mentally checking:\n"
            "  1. It starts with '{'\n"
            "  2. It ends with '}'\n"
            "  3. All quotes inside are escaped as \\\"\n"
            "  4. All newlines are escaped as \\n\n"
            "  5. There are NO characters after the final '}'\n\n"
            "MODIFICATION RULES:\n"
            "- Use 'add_after' to add new content after existing text\n"
            "- Use 'add_before' to add new content before existing text\n"
            "- Use 'replace' to replace existing text with new text\n"
            "- The 'search' field must be an EXACT match of existing text in the file (use enough context to make it unique)\n"
            "- The 'content' field is the new content to add or replace\n"
            "- CRITICAL for add_after/add_before: The 'content' field should contain ONLY the NEW content to add.\n"
            "  DO NOT include the search text in the content field - it will be duplicated!\n"
            "  WRONG: {\"operation\": \"add_after\", \"search\": \"old text\", \"content\": \"old text\\nnew line\"}\n"
            "  RIGHT: {\"operation\": \"add_after\", \"search\": \"old text\", \"content\": \"new line\"}\n"
            "- Apply modifications in order (they are applied sequentially)\n"
            "- Only include files that need modifications\n"
            "- Make MINIMAL changes - only modify what's necessary\n"
            "- DO NOT delete content unless explicitly replacing it\n"
            "- NEVER modify, add, or remove {%% include %%} statements - these are protected\n\n"
            "FORMAT CONSISTENCY:\n"
            "- When adding a new rule or item to an existing numbered/bulleted list, APPEND it at the end with the next sequential number. Do NOT insert before existing items or break the numbering.\n"
            "- Match the surrounding formatting style (indentation, bullet type, tone) of the section you are editing.\n\n"
            "AVOID DUPLICATION:\n"
            "- Each rule or constraint should appear in only ONE file. Do NOT add the same or similar rule to multiple files.\n"
            "- Before adding a rule, check if an equivalent rule already exists in any file.\n\n"
            "CONCISENESS:\n"
            "- Keep added content short and direct. Prefer a single concise sentence over a lengthy paragraph.\n"
            "- Do NOT add verbose explanations, examples, or justifications inside the prompt template unless strictly necessary.\n\n"
            "LANGUAGE:\n"
            "- All modifications to prompt files MUST be written in English only.\n"
            "- Do NOT output Chinese, Japanese, or any other non-English language in the modified content.\n\n"
            "- In a single iteration, you MUST NOT introduce more than ONE new absolute prohibition (MUST NOT).\n"
            "- Prefer conditional or scoped rules over global prohibitions whenever possible.\n"
            "CRITICAL - OVER-CONSTRAINT CHECK:\n"
            "- If the previous iteration shows score regression AND lessons_learned mention over-restriction,\n"
            "  you MUST explore alternative formulations or conditional relaxations instead of adding stricter rules.\n"

            "IMPORTANT: Your final output MUST be valid JSON indicating completion:\n"
            '{"status": "success", "modified_files": [...], "summary": "..."}\n'
            "Do not fabricate tool outputs. Use actual results from tool calls.\n\n"
        )
        agent_kwargs: dict[str, Any] = {
            "model_client": model_client,
            "tools": tools,
            "system_message": system_message,
            "model_client_stream": False,
            "description": "Optimizes prompt and pushes to repo.",
            "max_tool_iterations": 10,
        }
        if use_structured_output:
            agent_kwargs["output_content_type"] = AgentResponse

        self.optimizedPrompt = AssistantAgent(name, **agent_kwargs)

    async def _review_modifications(self, modifications_json: str) -> str:
        """Use LLM to review and deduplicate modifications across files.

        Condenses verbose additions and removes cross-file duplication
        before the modifications are applied to prompt files.
        """
        from autogen_core.models import SystemMessage, UserMessage

        system_prompt = (
            "You are a prompt modification reviewer.\n"
            "Your job: review modifications intended for multiple prompt template files and remove redundancy.\n\n"
            "Modification format reference:\n"
            "- Format: {\"file\": [{\"operation\":..., \"search\":..., \"content\":...}]}\n"
            "- Each modification has an \"operation\" field: 'add_after' (insert new content after the matched text), "
            "'add_before' (insert new content before the matched text), or 'replace' (replace the matched text with new content).\n"
            "- The \"search\" field is an exact substring match of existing text in the target file, used to locate where the modification applies.\n"
            "- The \"content\" field is the new text to insert or use as replacement.\n\n"
            "Rules:\n"
            "1. If similar or overlapping content is being added to multiple files, "
            "keep it in only the MOST relevant file and remove from others.\n"
            "2. Condense verbose additions — prefer 1-2 concise sentences over long bullet-point lists.\n"
            "3. Preserve the JSON structure exactly: same keys, same operations, same search fields.\n"
            "4. Only modify the \"content\" fields — never change \"operation\" or \"search\".\n"
            "5. To skip a modification entirely, set its \"content\" to \"\" (empty string).\n"
            "6. Return ONLY the cleaned JSON, no explanation.\n"
        )

        user_prompt = f"Review and deduplicate these modifications:\n{modifications_json}"

        logger.info("Reviewing modifications with LLM before writing...")
        logger.info("[review] BEFORE:\n%s", modifications_json)
        try:
            result = await self.model_client.create(
                messages=[
                    SystemMessage(content=system_prompt),
                    UserMessage(content=user_prompt, source="user"),
                ]
            )

            reviewed = result.content if isinstance(result.content, str) else str(result.content)
            # Strip markdown code fences if present
            reviewed = reviewed.strip()
            if reviewed.startswith("```"):
                lines = reviewed.split("\n")
                lines = lines[1:]  # remove opening ```json
                if lines and lines[-1].strip() == "```":
                    lines = lines[:-1]
                reviewed = "\n".join(lines).strip()

            # Validate it's still valid JSON; fallback to original if not
            json.loads(reviewed)
            changed = reviewed != modifications_json
            logger.info("[review] AFTER:\n%s", reviewed)
            logger.info("[review] Changed: %s", changed)
            return reviewed
        except json.JSONDecodeError:
            logger.warning("LLM review returned invalid JSON, using original modifications")
            return modifications_json
        except Exception as e:
            logger.warning("LLM review failed (%s), using original modifications", e)
            return modifications_json


if __name__ == "__main__":
    import asyncio
    import sys
    from pathlib import Path

    from autogen_agentchat.ui import Console
    from prompt_auto_lab.core.auth_token import get_azure_token_provider

    repo_root = Path(__file__).resolve().parents[2]
    sys.path.append(str(repo_root))

    from prompt_auto_lab.core.protocol import BraintrustClient, GitClient, InMemoryEvalBackend, PicassoClient
    from prompt_auto_lab.core.model_clients import build_azure_model_thought_client
    import os

    async def test_agent() -> None:
        """Test complete PromptOptimizer Agent"""
        # Create PicassoClient instance for testing
        picasso_client = PicassoClient()
       
        token_provider = get_azure_token_provider("https://cognitiveservices.azure.com/.default")
        model_client = build_azure_model_thought_client(token_provider=token_provider)
 
        from prompt_auto_lab.core.protocol import BraintrustClient, InMemoryEvalBackend
        from prompt_auto_lab.orchestrator.state_store import SharedState
        import os
 
        # Use the specified analysis.txt file
        experiment_id = "645dec68-2592-4f6a-aa64-820430030641"
        analysis_file_path = "dataset/braintrust_test_output/experiments_645dec68-2592-4f6a-aa64-820430030641/analysis.txt"
       
        # Create BraintrustClient for reading analysis
        backend = InMemoryEvalBackend()
        braintrust_org = os.getenv("BRAINTRUST_ORG", "non-pme")
        braintrust_api_key = os.getenv("BRAINTRUST_AZURE_API_KEY", "") if braintrust_org == "bt-azure" else os.getenv("BRAINTRUST_API_KEY", "")
        braintrust_client = BraintrustClient(
            backend=backend,
            api_key=braintrust_api_key,
            output_dir="dataset/braintrust_test_output",
            local_debug=True
        )
 
        agent = OptimizedPrompt(
            model_client=model_client,
            client=picasso_client,
            shared_state=SharedState(),
            braintrust_client=braintrust_client,
            AgentResponse=AgentResponse,
        ).optimizedPrompt
       
        # Read analysis from the specified file
        result = braintrust_client.read_analysis_txt(experiment_id)
       
        if not result["ok"]:
            logger.error("Failed to read analysis file: %s", result.get("error"))
            logger.error("Attempted to read: %s", analysis_file_path)
            return
       
        logger.info("Successfully read analysis from: %s", result.get("file_path", analysis_file_path))
        if result.get("data"):
            logger.info("Analysis data contains %s record(s)", len(result["data"]))
       
        branch = "autopromptlab2/v1"
       
        # Simple task - let the agent's system_message guide the workflow
        task = f"""Improve the prompt using the analysis results from experiment_id: {experiment_id}
 
branch: {branch}
experiment_id: {experiment_id}
"""
        # Use RoundRobinGroupChat to allow agent to execute multiple rounds
        from autogen_agentchat.teams import RoundRobinGroupChat
        from autogen_agentchat.conditions import TextMentionTermination, MaxMessageTermination
       
        # Terminate when agent output contains "status", or after max 20 messages
        # termination = TextMentionTermination('"status"') | MaxMessageTermination(max_messages=20)
       
        # team = RoundRobinGroupChat(
        #     participants=[agent],
        #     termination_condition=termination,
        # )
       
        res = await Console(agent.run_stream(task=task), output_stats=True)
        logger.info("\nFinal Agent thought:")
        logger.info("%s", res.messages[-1].content.thoughts)
        logger.info("%s", res.messages[-1].content)
        await model_client.close()
 
    # Choose which test to run
    import sys
    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s")
    if len(sys.argv) > 1 and sys.argv[1] == "agent":
        asyncio.run(test_agent())
    else:
        logger.info("Usage: python -m prompt_auto_lab.agents.optimizer agent")
 
