"""CrossoverAgent for combining optimizations from multiple parent candidates."""
from __future__ import annotations

import logging
from typing import Any, TYPE_CHECKING

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

if TYPE_CHECKING:
    from prompt_auto_lab.core.protocol import GitClient, PicassoClient
    from prompt_auto_lab.orchestrator.state_store import SharedState

logger = logging.getLogger(__name__)


class CrossoverAgent:
    """Agent that crossovers (combines) optimizations from multiple parent candidates.

    This agent reads the git diffs from selected parent candidates and
    intelligently merges their changes to create a new hybrid candidate.
    """

    def __init__(
        self,
        model_client: Any,
        picasso_client: PicassoClient,
        git_client: GitClient,
        shared_state: SharedState,
        allowed_tools: set[str] | None = None,
        name: str = "CrossoverAgent",
    ):
        self.picasso_client = picasso_client
        self.git_client = git_client
        self.shared_state = shared_state

        tools = []
        allowed = allowed_tools or {
            "get_parent_diffs",
            "read_prompt_template",
            "write_prompt_files",
        }

        # Tool to get diffs from parent commits compared to analysis commit
        async def get_parent_diffs(
            parent_commits: list[str],
            analysis_commit: str,
        ) -> dict[str, Any]:
            """Get git diffs for N parent commits compared to the analysis commit.

            Args:
                parent_commits: List of parent commit IDs (N >= 2)
                analysis_commit: The common ancestor commit (analysis baseline)

            Returns:
                Dict with diffs from all parents:
                {
                    "parents": [{"index": int, "commit": str, "diff": str, "ok": bool, ...}, ...],
                    "analysis_commit": str,
                    "diff_files": [str, ...],
                    "parent_count": int
                }
            """
            try:
                results: dict[str, Any] = {"analysis_commit": analysis_commit}
                all_diff_files: set[str] = set()
                parents_list: list[dict[str, Any]] = []

                # Fetch from origin to ensure we have the latest commits
                await git_client.git_fetch()

                for idx, parent_commit in enumerate(parent_commits, 1):
                    # Get diff between analysis commit and parent commit
                    # Use max_chars=50000 to get complete diff for crossover (default 2000 is too small)
                    diff_result = await git_client.git_diff(analysis_commit, parent_commit, max_chars=50000)
                    diff_text = diff_result.get("diff", "")

                    # Use files_changed from git_diff result (extracted before truncation)
                    files_in_diff = set(diff_result.get("files_changed", []))
                    all_diff_files.update(files_in_diff)

                    parents_list.append({
                        "index": idx,
                        "commit": parent_commit,
                        "diff": diff_text,
                        "ok": diff_result.get("ok", False),
                        "error": diff_result.get("error"),
                        "files": list(files_in_diff),
                        "diff_truncated": diff_result.get("diff_truncated", False),
                    })

                results["parents"] = parents_list
                results["parent_count"] = len(parents_list)
                # Add union of all files modified by all parents
                results["diff_files"] = sorted(all_diff_files)
                results["ok"] = True
                return results

            except Exception as e:
                logger.exception("Error getting parent diffs")
                return {"ok": False, "error": str(e)}

        if "get_parent_diffs" in allowed:
            tools.append(FunctionTool(
                get_parent_diffs,
                "Get git diffs for N parent commits compared to analysis commit. Use this to see what changes each parent made.",
                strict=True,
            ))

        if "read_prompt_template" in allowed:
            tools.append(FunctionTool(
                picasso_client.read_prompt_template,
                "Read prompt files from Picasso repo. Input: JSON array of file paths.",
                strict=True,
            ))

        if "write_prompt_files" in allowed:
            tools.append(FunctionTool(
                picasso_client.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,
            ))

        tools_description = ", ".join([f"{tool.name}: {tool._description}" for tool in tools])

        system_message = (
            "You are CrossoverAgent.\n"
            "Goal: Combine optimizations from N parent commits into a SINGLE, unified hybrid candidate.\n\n"
            "CRITICAL: You MUST execute ALL steps automatically. Do NOT ask for confirmation, do NOT pause, do NOT say 'continue' or 'ready to run'. Just execute the tools immediately and complete the entire workflow in one go.\n\n"
            "UNDERSTANDING THE CONTEXT:\n"
            "- The 'analysis_commit' is the BASELINE (original version before any optimizations)\n"
            "- Each parent commit contains DIFFERENT optimizations on top of that baseline\n"
            "- The diffs show what each parent ADDED or CHANGED compared to the baseline\n"
            "- The current files you read are the BASELINE version\n"
            "- Your job is to APPLY the combined changes from all parents to the baseline files\n\n"
            "WORKFLOW:\n"
            "1. Use get_parent_diffs with the provided parent_commits list and analysis_commit to see what changes each parent made\n"
            "   - The result contains a 'parents' array, each with '+' lines (additions) and '-' lines (removals)\n"
            "   - The result includes 'diff_files': a list of all files modified by any parent (union)\n"
            "   - The result includes 'parent_count': the number of parents\n"
            "2. Use read_prompt_template to read the current (baseline) content of files in 'diff_files'\n"
            "   - Pass the diff_files array as a JSON array string: '[\"path1\", \"path2\"]'\n"
            "   - IMPORTANT: Only include the array, no extra brackets or braces!\n"
            "3. Analyze all parent diffs to understand what each parent optimized:\n"
            "   - Look at the '+' lines to see what each parent ADDED\n"
            "   - Identify the INTENT behind each change\n"
            "4. Combine the best aspects of ALL parents' changes:\n"
            "   - If changes are to different files/sections, include all of them\n"
            "   - If changes conflict, intelligently merge or pick the best one\n"
            "   - **CRITICAL: If multiple parents made SIMILAR changes (same intent, similar wording), you MUST synthesize them into ONE unified version** - do NOT keep duplicates!\n"
            "5. Use write_prompt_files to APPLY the combined changes to the baseline files\n"
            "   - You MUST make modifications - do not return empty modifications!\n"
            "   - Use the diff content to determine what text to add/replace\n\n"
            "CRITICAL - DEDUPLICATION RULES:\n"
            "- If multiple parents added similar content (e.g., several added instructions about page context), you MUST:\n"
            "  1. Identify the common intent/purpose across all parent changes\n"
            "  2. Write ONE unified paragraph that captures the best wording from all parents\n"
            "  3. NEVER keep near-duplicate paragraphs - this is unacceptable!\n"
            "- Example: If parent 1 added 'Use page content first before calling tools', parent 2 added 'Answer from page context before using tools', and parent 3 added 'Prioritize page content over tool calls', combine into ONE clear statement.\n\n"
            "IMPORTANT:\n"
            "- The goal is to create ONE clean, non-redundant candidate\n"
            "- Do NOT simply concatenate all parents' changes - actively MERGE and DEDUPLICATE\n"
            "- If some parents made no changes, use the other parents' changes\n"
            "- Ensure the final result is valid Liquid template syntax\n"
            "- You MUST apply modifications - never return empty modifications_json!\n\n"
            f"Available tools: {tools_description}\n\n"
            "CRITICAL - MODIFICATION INSTRUCTIONS FORMAT:\n"
            "Use the file path 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'
            '}\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\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"
            "- Apply modifications in order (they are applied sequentially)\n\n"
            "Return a JSON summary when done:\n"
            '{"status": "success", "modified_files": [...], "merged_from": ["commit1", "commit2", ...], "summary": "..."}\n'
        )

        self.crossover_agent = AssistantAgent(
            name,
            model_client=model_client,
            tools=tools,
            system_message=system_message,
            description="Combines optimizations from multiple parent candidates.",
            max_tool_iterations=10,
        )


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

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

    async def load_state_from_jsonl(jsonl_path: str) -> dict:
        """Load state from version.jsonl file (reads last line)."""
        state = {}
        try:
            with open(jsonl_path, "r", encoding="utf-8") as f:
                for line in f:
                    line = line.strip()
                    if line:
                        state = json.loads(line)
            logger.info("Loaded state from %s", jsonl_path)
        except Exception as e:
            logger.error("Failed to load state: %s", e)
        return state

    async def test_crossover_agent() -> None:
        """Test CrossoverAgent with mock or real state."""
        # Import inside function to avoid circular import
        from autogen_agentchat.ui import Console
        from prompt_auto_lab.core.auth_token import get_azure_token_provider
        from prompt_auto_lab.core.model_clients import build_azure_model_thought_client
        from prompt_auto_lab.core.protocol import GitClient, PicassoClient
        from prompt_auto_lab.orchestrator.state_store import SharedState

        # Option 1: Load from a real version.jsonl file
        # Uncomment and modify the path to use real data:
        # state_data = await load_state_from_jsonl("path/to/your/version.jsonl")

        # Option 2: Use mock state for testing
        # This simulates iteration 1 after OptimizeTeam has run, with 3 candidates
        # Each candidate has a different commit_id representing their optimizations
        state_data = {
            "iteration_idx": "1",
            "iteration": {
                "root": {
                    "analysis_commit_id": "480bc3bfb34ca4e44d385b9bed8875f2cff9b5f1",
                },
                "1": {
                    "analysis_source_node": "root",
                    "analysis_commit_id": "480bc3bfb34ca4e44d385b9bed8875f2cff9b5f1",
                    "candidates": {
                        "c1": {
                            "commit_id": "c5b948d",
                            "branch": "autopromptlab/ga/480bc3bf_1_32s8q1",
                        },
                        "c2": {
                            "commit_id": "1863f84",
                            "branch": "autopromptlab/ga/480bc3bf_2_6vlsg8",
                        },
                        "c3": {
                            "commit_id": "6521c11",
                            "branch": "autopromptlab/ga/480bc3bf_3_sao93s",
                        },
                    },
                },
            },
        }

        shared_state = SharedState(initial=state_data)

        # Create clients
        picasso_client = PicassoClient()
        git_client = GitClient(repo_path=picasso_client.repo_path)

        token_provider = get_azure_token_provider("https://cognitiveservices.azure.com/.default")
        model_client = build_azure_model_thought_client(token_provider=token_provider)

        # Create CrossoverAgent
        agent = CrossoverAgent(
            model_client=model_client,
            picasso_client=picasso_client,
            git_client=git_client,
            shared_state=shared_state,
        ).crossover_agent

        # Test commit info (from your version.jsonl iteration["1"])
        parent_commits = ["c5b948d", "1863f84", "6521c11"]  # c1, c2, and c3 commit IDs
        analysis_commit = "480bc3bfb34ca4e44d385b9bed8875f2cff9b5f1"  # root's analysis_commit_id

        task = f"""Combine optimizations from {len(parent_commits)} parent commits. Execute ALL steps immediately without asking for confirmation.

Use get_parent_diffs with:
  - parent_commits: {json.dumps(parent_commits)}
  - analysis_commit: {analysis_commit}

Then read the files, merge the changes intelligently, and write the combined result.
"""

        logger.info("Starting CrossoverAgent test with %d parents: %s", len(parent_commits), parent_commits)

        res = await Console(agent.run_stream(task=task), output_stats=True)

        logger.info("\nFinal Agent output:")
        if res.messages:
            logger.info("%s", res.messages[-1].content)

        await model_client.close()

    # Run test
    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_crossover_agent())
    elif len(sys.argv) > 1 and sys.argv[1] == "load":
        # Load state from file and print
        if len(sys.argv) > 2:
            async def _load_and_print():
                state = await load_state_from_jsonl(sys.argv[2])
                print(json.dumps(state, indent=2, ensure_ascii=False))
            asyncio.run(_load_and_print())
        else:
            print("Usage: python -m prompt_auto_lab.agents.crossover load <path/to/version.jsonl>")
    else:
        print("Usage:")
        print("  python -m prompt_auto_lab.agents.crossover agent")
        print("  python -m prompt_auto_lab.agents.crossover load <path/to/version.jsonl>")
