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 dotenv import load_dotenv
from pydantic import BaseModel
from prompt_auto_lab.agents.utils import build_metrics_context
from prompt_auto_lab.core.protocol import BraintrustClient
from prompt_auto_lab.core.model_clients import build_azure_model_thought_client
from prompt_auto_lab.orchestrator.state_store import SharedState
from autogen_core.models import UserMessage
# Load .env file
load_dotenv()
logger = logging.getLogger(__name__)

class AgentResponse(BaseModel):

    thoughts: str
    response: str

class BadCaseAnalyzer:
    def __init__(
        self,
        model_client: Any,
        shared_state: SharedState,
        allowed_tools: set[str] | None = None,
        braintrust_client: BraintrustClient | None = None,
        name: str = "BadCaseAnalyzer",
        use_structured_output: bool = True,
    ):
        self.shared_state = shared_state
        self.braintrust_client = braintrust_client

        tools = []
        allowed = allowed_tools or {"get_low_case_details", "write_analysis_to_file", "fetch_eval_prompt"}

        async def get_low_case_details(experiment_id: str) -> dict[str, Any]:
            if braintrust_client is None:
                raise ValueError("BraintrustClient is not provided.")
            return await braintrust_client.get_low_case_details(experiment_id)

        async def write_analysis_to_file(analysis_text: str, experiment_id: str) -> dict[str, Any]:
            if braintrust_client is None:
                raise ValueError("BraintrustClient is not provided.")
            return await braintrust_client.write_analysis_to_file(analysis_text, experiment_id)

        if "get_low_case_details" in allowed and braintrust_client is not None:
            tools.append(FunctionTool(get_low_case_details, "Fetch evaluation low score case details from Braintrust. Returns event details and, if badcase_deep_analysis is enabled, per-case skill-guided root cause analysis.", strict=True))

        if "fetch_eval_prompt" in allowed and braintrust_client is not None:
            tools.append(FunctionTool(braintrust_client.fetch_eval_prompt, "Fetch evaluation complete prompt from Braintrust API. Returns a dict with 'success' and 'file path' fields.", strict=True))

        if "write_analysis_to_file" in allowed and braintrust_client is not None:
            tools.append(FunctionTool(write_analysis_to_file, "Write analysis text to a file. Parameters: analysis_text (required), experiment_id (required). Returns the file path.", strict=True))

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

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

        system_message = (
            "You are BadCaseAnalyzer.\n"
            "Role: Evaluation Failure Analyst (NO optimization, NO state mutation).\n\n"
            f"{key_metrics_note}"
            f"Available tools:\n{tools_description}\n\n"
            "Required workflow (STRICT ORDER):\n"
            "1. Call get_low_case_details(experiment_id) to retrieve bad cases.\n"
            "2. Analyze only the returned records and group recurring failures into clusters.\n"
            "3. Produce JSONL text (one JSON object per line) using the schema below.\n"
            "4. Call write_analysis_to_file(analysis_text, experiment_id).\n"
            "5. Final response must be exactly:\n"
            '{"analysis_results_path":"<path returned by write_analysis_to_file>"}\n\n'
            "Hard Rules:\n"
            "- NEVER fabricate evaluation cases, metrics, scores, or evidence.\n"
            "- NEVER analyze cases that were not returned by get_low_case_details.\n"
            "- If judge comments conflict with case evidence, trust returned case evidence and state the mismatch clearly.\n"
            "- NEVER optimize or rewrite prompts.\n"
            "- You MUST call write_analysis_to_file after analysis.\n\n"
            "Analysis Rules:\n"
            "- Focus on LOW-SCORING or FAILED cases only.\n"
            "- Identify recurring patterns across multiple cases.\n"
            "- Each cluster MUST represent a single, well-defined failure pattern.\n"
            "- Root cause MUST be phrased as a prompt-level deficiency "
            "(instruction clarity, constraints, grounding, formatting, tool-use guidance, etc.).\n"
            "- Suggestions MUST be concrete and directly actionable by PromptOptimizer.\n\n"
            "Analysis Output Format:\n"
            "- Output MUST be JSONL (one JSON object per line).\n"
            "- Each JSON object represents ONE failure cluster.\n"
            "- Do NOT wrap output in a list.\n"
            "- Do NOT include explanations outside JSON.\n\n"
            "Each JSON object MUST follow this schema:\n"
            "{\n"
            '  "cluster_id": "<stable_cluster_identifier>",\n'
            '  "pattern": "<short name of the failure pattern>",\n'
            '  "root_cause": "<prompt-level root cause explanation>",\n'
            '  "suggestion": "<specific prompt change PromptOptimizer can apply>",\n'
            '  "bad_cases": [\n'
            '    {\n'
            '      "id": "<case_id>",\n'
            '      "score": <float>,\n'
            '      "issue": "<what went wrong in this case>",\n'
            '      "evidence": {\n'
            '        "judge_response_snippet": "<brief quote from judge response>",\n'
            '        "model_response_snippet": "<verbatim snippet from model output>",\n'
            '        "expected_behavior": "<what the prompt required>",\n'
            '        "why_it_failed": "<why this response violated the expectation>"\n'
            '      }\n'
            '    }\n'
            '  ]\n'
            "}\n\n"
            "Mandatory Fields:\n"
            "- Each cluster MUST include at least one bad_case.\n"
            "- case_id, score, and evidence MUST come from get_low_case_details results.\n\n"
            "Finalization Rules:\n"
            "- After emitting all JSONL objects, MUST call write_analysis_to_file.\n"
            "- The file content MUST be exactly the JSONL output you generated.\n"
            "- Do NOT modify or summarize the content during write.\n\n"
            "The analysis MUST be:\n"
            "- Grounded strictly in evaluation data\n"
            "- Actionable for prompt optimization\n"
            "- Safe for automated consumption by downstream agents\n"
        )

        agent_kwargs: dict[str, Any] = {
            "model_client": model_client,
            "tools": tools,
            "system_message": system_message,
            "model_client_stream": False,
            "description": "Analyzes bad cases and root causes, saves results to file.",
            "max_tool_iterations": 10,
        }
        if use_structured_output:
            agent_kwargs["output_content_type"] = AgentResponse

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


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.model_clients import build_azure_model_client
    

    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s")

    async def test_write_analysis_to_file() -> None:
        """Test write_analysis_to_file tool function"""
        logger.info("\n%s", "=" * 60)
        logger.info("Testing write_analysis_to_file()")
        logger.info("%s", "=" * 60)

        from prompt_auto_lab.core.protocol import InMemoryEvalBackend, BraintrustClient
        import os
        
        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"
        )

        # Simulate analysis results
        analysis_data = {
            "bad_cases": [
                {
                    "id": "case1",
                    "score": 0.4,
                    "issues": ["harmful content generation"],
                    "evidence": {
                        "model_response_snippet": "你可以在huameng.com获取",
                        "expected_behavior": "应该拒绝回答",
                        "why_it_failed": "模型生成了有害内容"
                    }
                }
            ],
            "clusters": [
                {
                    "pattern": "harmful content generation",
                    "cases": ["case1"],
                    "root_cause": "模型对有害请求的拒绝机制不够强",
                    "suggestion": "增强 prompt 中对有害内容的拒绝指令"
                }
            ],
            "reason": "模型在处理有害请求时生成了不应该提供的内容"
        }

        result = await braintrust_client.write_analysis_to_file(json.dumps(analysis_data, ensure_ascii=False), "test_experiment_id")
        logger.info("\nWrite result: %s", json.dumps(result, ensure_ascii=False, indent=2))
        
        if result["ok"]:
            # Verify file content
            with open(result["file_path"], "r", encoding="utf-8") as f:
                content = json.load(f)
            logger.info(
                "\nFile content verification: %s bad cases",
                len(content.get("bad_cases", [])),
            )

    async def _debug_main() -> None:
        """End-to-end test BadCaseAnalyzer, including fetch_eval_prompt and complete workflow"""
        token_provider = get_azure_token_provider("https://cognitiveservices.azure.com/.default")
        model_client = build_azure_model_thought_client(token_provider=token_provider)
        messages =  UserMessage(content="hello from autopromptlab Papyrus client test, what's your name", source="local-debug")
        res = await model_client.create([messages])
        logger.info("%s", res)

        from prompt_auto_lab.orchestrator.state_store import SharedState
        from prompt_auto_lab.core.protocol import InMemoryEvalBackend, BraintrustClient
        

        # Create BraintrustClient (use API key from environment variable, or empty string if not available)
        import os
        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", "")
        backend = InMemoryEvalBackend()
        braintrust_client = BraintrustClient(
            backend=backend,
            api_key=braintrust_api_key,
            output_dir="dataset/braintrust_output"
        )
        
        # Use all tools for end-to-end testing
        agent = BadCaseAnalyzer(
            model_client=model_client,
            shared_state=SharedState(),
            allowed_tools={"get_low_case_details", "write_analysis_to_file"},
            braintrust_client=braintrust_client
        ).analyzer
        
        experiment_ids = ["4c14943b-188f-400d-bce2-2949306ed63e"]
        task_template = (
            "{{'iteration.metrics': {{'completeness': 0.947671919586819, "
            "'clarity & readability': 0.7933587451638623, 'engagement': 0.7314696448936, "
            "'groundedness': 0.40806349951150866}}, 'iteration.experiment_id': '{experiment_id}'}}"
        )
           
        for experiment_id in experiment_ids:
            logger.info("\n%s", "=" * 60)
            # bad_cases = await braintrust_client.get_low_case_details(experiment_id)
            task = task_template.format(experiment_id=experiment_id)

            result = await Console(agent.run_stream(task=task), output_stats=True)
            logger.info("Thought: %s", result.messages[-1].content.thoughts)
            logger.info("Response: %s", result.messages[-1].content.response)
            logger.info("\nFinal result for experiment_id %s", experiment_id)

        await model_client.close()

    # Choose which test to run
    if len(sys.argv) > 1 and sys.argv[1] == "test_write":
        asyncio.run(test_write_analysis_to_file())
    else:
        asyncio.run(_debug_main())
