from __future__ import annotations

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

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.clients.picasso.liquid_function_client import LiquidFunctionClient
from prompt_auto_lab.orchestrator.state_store import SharedState

# Load .env file
load_dotenv()
logger = logging.getLogger(__name__)


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


class LiquidStructureAgent:
    """
    Agent that analyzes Liquid template structure from an entry point template.
    Builds include tree, analyzes template purposes, and generates documentation
    to guide the PromptOptimizer on where to make changes.

    All core functionality is delegated to LiquidFunctionClient.
    """

    def __init__(
        self,
        model_client: Any,
        shared_state: SharedState,
        allowed_tools: set[str] | None = None,
        name: str = "LiquidStructureAgent",
        repo_path: str | Path | None = None,
        store_dir: str | Path | None = None,
        use_structured_output: bool = True,
        braintrust_client: Any = None,
    ):
        self.shared_state = shared_state
        self.model_client = model_client
        self.braintrust_client = braintrust_client

        # Initialize LiquidFunctionClient for core functionality
        self.liquid_client = LiquidFunctionClient(
            repo_path=repo_path,
            store_dir=store_dir,
        )

        # Reference key paths from the client
        self.repo_path = self.liquid_client.repo_path
        self.store_dir = self.liquid_client.store_dir
        self.templates_base = self.liquid_client.templates_base

        # Build tools - delegate to LiquidFunctionClient
        tools = []
        allowed = allowed_tools or {
            "analyze_bad_cases_for_templates",
        }

        if "analyze_bad_cases_for_templates" in allowed:
            tools.append(FunctionTool(
                self._analyze_bad_cases_for_templates,
                "Analyze bad case clusters to identify 2-3 priority templates for optimization. Takes the analysis file path (from BadCaseAnalyzer) plus the experiment_id, maps failure patterns to specific liquid templates, and reads eval_prompt.jsonl for condition-aware filtering.",
                strict=True
            ))

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

        system_message = (
            "You are LiquidStructureAgent.\n"
            "Role: Liquid Template Structure Analyst.\n\n"

            "Goal:\n"
            "- Read the pre-generated liquid structure documentation (liquid_structure.md)\n"
            "- Map bad case analysis to priority templates for optimization\n\n"

            "Workflow (1 step):\n"
            "1. Use analyze_bad_cases_for_templates with the analysis file path and experiment_id\n"
            "   - The liquid structure documentation has already been generated during initialization\n"
            "   - The tool reads the bad case analysis file and uses LLM to identify 2-3 priority templates\n"
            "   - The tool then loads eval_prompt.jsonl for the given experiment_id and runs condition-aware filtering\n"
            "   - Cross-references bad cases with the existing structure documentation\n"
            "   - Returns priority_templates_path and priority_templates list\n\n"

            f"Available tools:\n{tools_description}\n\n"

            "Hard Rules:\n"
            "- NEVER stop to ask the user for input - complete the workflow autonomously\n"
            "- The liquid structure documentation is already available from initialization\n"
            "- CRITICAL: Your response MUST include the priority_templates_path from the tool result\n"
            "- Response format: Return a summary mentioning the identified templates AND the priority_templates_path\n"
            "- Example response: 'priority_templates_path: {priority_templates_path}'\n"
        )

        agent_kwargs: dict[str, Any] = {
            "model_client": model_client,
            "tools": tools,
            "system_message": system_message,
            "model_client_stream": False,
            "description": "Analyzes Liquid template structure and generates documentation for prompt optimization.",
            "max_tool_iterations": 4,
        }
        if use_structured_output:
            agent_kwargs["output_content_type"] = AgentResponse

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

    async def _build_and_write_structure_doc(self, entry_point: str | None = None, liquid_entries: list[dict[str, Any]] | None = None) -> dict[str, Any]:
        """
        Build include tree and automatically generate + write structure documentation.
        Delegates to LiquidFunctionClient and updates SharedState.

        Supports two modes:
        1. Single entry_point (backward compatibility): Analyzes one liquid file and its includes
        2. Multiple liquid_entries (new feature): Processes multiple files, analyzing some and listing others

        Args:
            entry_point: Path to the entry point template (legacy, for backward compatibility)
            liquid_entries: List of {"path": str, "analyze_structure": bool} entries (new feature)

        Returns:
            {"ok": bool, "file_path": str, "total_templates": int, "size": int}
        """
        # Determine which mode to use
        if liquid_entries:
            # New mode: multiple entries with analysis flags
            result = await self.liquid_client.build_and_write_structure_doc_from_entries(
                liquid_entries, self.model_client
            )
        elif entry_point:
            # Legacy mode: single entry point
            result = await self.liquid_client.build_and_write_structure_doc(entry_point, self.model_client)
        else:
            return {"ok": False, "error": "Either entry_point or liquid_entries must be provided"}

        if result.get("ok"):
            # Update shared state with the path
            await self.shared_state.patch({"liquid_structure_doc_path": result.get("file_path")})

        return result

    @staticmethod
    def _resolve_eval_prompt_path(analysis_path: Path, experiment_id: str) -> Path:
        """Resolve eval_prompt.jsonl path for the given experiment_id."""
        experiment_dir_name = f"experiments_{experiment_id}"
        if analysis_path.parent.name == experiment_dir_name:
            return analysis_path.parent / "eval_prompt.jsonl"
        return analysis_path.parent.parent / experiment_dir_name / "eval_prompt.jsonl"

    @staticmethod
    def _load_rendered_prompt_from_eval_file(prompt_path: Path) -> str:
        """Load rendered prompt text from eval_prompt.jsonl.

        The file is a single-line JSONL with structure:
        {"data": [{"prompt_text": "..."}], "schema": ..., ...}
        We extract the prompt_text field to get the actual rendered prompt.
        """
        if not prompt_path.exists():
            raise FileNotFoundError(f"eval_prompt.jsonl not found: {prompt_path}")

        raw = prompt_path.read_text(encoding="utf-8")
        try:
            first_line = raw.split("\n", 1)[0]
            parsed = json.loads(first_line)
            rows = parsed.get("data", [])
            if rows and isinstance(rows, list) and "prompt_text" in rows[0]:
                return rows[0]["prompt_text"]
        except (json.JSONDecodeError, KeyError, IndexError):
            pass
        # Fallback: return raw content if parsing fails
        return raw

    async def _analyze_bad_cases_for_templates(self, analysis_file_path: str, experiment_id: str) -> dict[str, Any]:
        """
        Analyze bad case clusters to identify 2-3 priority templates for optimization.
        Uses LLM to analyze bad cases and map them to relevant templates.

        Args:
            analysis_file_path: Path to the analysis file from BadCaseAnalyzer
            experiment_id: Experiment ID used to locate eval_prompt.jsonl

        Returns:
            {
                "ok": bool,
                "priority_templates_path": str,
                "priority_templates": [...],
                "total_clusters_analyzed": int
            }
        """
        # Read the analysis file
        try:
            analysis_path = Path(analysis_file_path)
            # Convert absolute path starting with / to relative path starting with .autopromptlab2
            if str(analysis_path).startswith("/.autopromptlab2"):
                analysis_path = Path("." + str(analysis_path))
            if not analysis_path.exists():
                return {"ok": False, "error": f"Analysis file not found: {analysis_file_path}"}

            bad_case_analysis = analysis_path.read_text(encoding="utf-8")
            # Get the directory of the analysis file to save priority_templates.md there
            output_dir = analysis_path.parent
        except Exception as e:
            return {"ok": False, "error": f"Failed to read analysis file: {e}"}

        # Read structure doc if available
        structure_doc = self.liquid_client.read_structure_doc()

        if not structure_doc:
            return {
                "ok": False,
                "error": "Structure documentation not available. Please run build_include_tree and write_structure_doc first."
            }

        # Use LLM-based analysis (model_client is required)
        if self.model_client is None:
            return {"ok": False, "error": "model_client is required for bad case analysis"}

        result = await self.liquid_client.analyze_bad_cases_for_templates_llm(
            bad_case_analysis=bad_case_analysis,
            model_client=self.model_client,
            structure_doc=structure_doc,
            output_dir=output_dir,
        )

        if result.get("ok"):
            # Update shared state with priority templates path under iteration namespace
            current_state = await self.shared_state.read(None)
            iteration_idx = current_state.get("iteration_idx", "root")

            # Get current iteration dict and update it
            iteration_data = current_state.get("iteration", {})
            if iteration_idx in iteration_data:
                iteration_data[iteration_idx]["priority_templates_path"] = result.get("priority_templates_path", "")

            await self.shared_state.patch({
                "iteration": iteration_data,
            })

            # Condition-aware filtering: load eval_prompt.jsonl by explicit experiment_id input
            try:
                prompt_path = self._resolve_eval_prompt_path(analysis_path, experiment_id)
                logger.info(
                    "Condition-aware filtering: loading rendered prompt from %s for experiment %s",
                    prompt_path,
                    experiment_id,
                )
                rendered_prompt = self._load_rendered_prompt_from_eval_file(prompt_path)
                logger.info(
                    "Condition-aware filtering: rendered prompt loaded (%d chars), running filter pipeline",
                    len(rendered_prompt),
                )
                filter_result = await self.liquid_client.filter_priority_templates(
                    rendered_prompt=rendered_prompt,
                    priority_templates=result.get("priority_templates", []),
                )

                if filter_result.get("ok"):
                    self.liquid_client.write_filtered_sidecars(
                        filtered_templates=filter_result["filtered_templates"],
                    )
                    logger.info("Condition-aware filtering: successfully wrote sidecars")
                else:
                    logger.warning(
                        "Condition-aware filtering: filter pipeline returned error: %s",
                        filter_result.get("error"),
                    )
            except Exception as e:
                logger.warning("Condition-aware filtering failed (non-fatal): %s", e)

        return result


if __name__ == "__main__":
    import asyncio
    import sys

    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_thought_client

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

    async def test_tools() -> None:
        """Test the individual tools using LiquidFunctionClient."""
        logger.info("\n%s", "=" * 60)
        logger.info("Testing LiquidStructureAgent Tools via LiquidFunctionClient")
        logger.info("%s", "=" * 60)

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

        # Use LiquidFunctionClient directly for testing
        client = LiquidFunctionClient()

        # Test build_and_write_structure_doc (combined method)
        logger.info("\n1. Testing build_and_write_structure_doc...")
        result = await client.build_and_write_structure_doc("Service.Chat/Inference/Prompts/Templates/syndication-api/responding/prompt.liquid", model_client)
        logger.info("   Result: ok=%s", result.get("ok"))
        if result.get("ok"):
            logger.info("   Total templates: %s", result.get("total_templates", 0))
            logger.info("   File path: %s", result.get("file_path"))
            logger.info("   File size: %s bytes", result.get("size"))
        else:
            logger.error("   Error: %s", result.get("error"))

        # Test structure_doc_exists and read_structure_doc
        logger.info("\n2. Testing structure_doc_exists and read_structure_doc...")
        exists = client.structure_doc_exists()
        logger.info("   Structure doc exists: %s", exists)
        if exists:
            content = client.read_structure_doc()
            logger.info("   Content length: %s chars", len(content) if content else 0)
            # Print first 2000 chars
            if content:
                logger.info("\n   Content preview (first 2000 chars):")
                logger.info("%s", content[:2000])

        await model_client.close()

    async def test_analyze_bad_cases_llm() -> None:
        """Test the LLM-based analyze_bad_cases_for_templates function."""
        logger.info("\n%s", "=" * 60)
        logger.info("Testing analyze_bad_cases_for_templates_llm")
        logger.info("%s", "=" * 60)

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

        # Use LiquidFunctionClient directly for testing
        client = LiquidFunctionClient()

        # Generate structure doc programmatically from actual templates
        logger.info("Generating structure doc from actual templates...")
        result = await client.build_and_write_structure_doc("Service.Chat/Inference/Prompts/Templates/syndication-api/responding/prompt.liquid", model_client)
        if result.get("ok"):
            logger.info("   Generated doc with %s templates", result.get("total_templates"))
            doc_content = client.read_structure_doc()
        else:
            logger.error("   Error generating doc: %s", result.get("error"))
            return

        # Create sample bad case analysis JSONL
        sample_bad_cases = "{\"cluster_id\":\"claim_content_missing\",\"pattern\":\"Model falsely claims page content is filtered or unavailable\",\"root_cause\":\"Prompt lacks strict grounding requirement prohibiting statements that page content is missing or filtered.\",\"suggestion\":\"Add rule: 'Never state that webpage content is missing, filtered, or inaccessible. Always use the provided context.'\",\"bad_cases\":[{\"id\":\"67925781-9dc2-4b7b-bbaf-eec9b7400f83\",\"score\":0.05,\"issue\":\"Model claimed page was filtered despite full content present.\",\"evidence\":{\"model_response_snippet\":\"当前页面内容被过滤\",\"expected_behavior\":\"Extract facts from provided chapter text.\",\"why_it_failed\":\"Model hallucinated missing content instead of grounding response.\"}},{\"id\":\"6d430d33-ca58-4c5c-99ff-1f962172f955\",\"score\":0.05,\"issue\":\"Model again stated content was filtered and inaccessible.\",\"evidence\":{\"model_response_snippet\":\"由于页面内容被过滤\",\"expected_behavior\":\"Use AI browser article text to generate facts.\",\"why_it_failed\":\"Model refused to use available context and invented absence.\"}}]}\n{\"cluster_id\":\"ungrounded_generics\",\"pattern\":\"Model provides generic domain facts unrelated to page content\",\"root_cause\":\"Prompt allows domain knowledge injection without enforcing page-specific grounding.\",\"suggestion\":\"Add constraint: 'All facts must directly reference elements present in the webpage. Do not add external domain knowledge unless explicitly in context.'\",\"bad_cases\":[{\"id\":\"d7d0c784-af2b-4f89-bf92-817c82fc3b99\",\"score\":0.6,\"issue\":\"Environmental chemistry facts unrelated to journal homepage.\",\"evidence\":{\"model_response_snippet\":\"环境化学是研究化学物质在自然环境中行为...\",\"expected_behavior\":\"Describe journal metadata, submission info, or page elements.\",\"why_it_failed\":\"Model filled content gap with generic field overview.\"}},{\"id\":\"3ba0f406-ca0e-4954-92a6-fbb76ffd2064\",\"score\":0.6,\"issue\":\"Provided history of ancient novels unrelated to login page.\",\"evidence\":{\"model_response_snippet\":\"中国古代小说的发展可以追溯到先秦时期\",\"expected_behavior\":\"Discuss platform details and metadata.\",\"why_it_failed\":\"Model hallucinated topic-based content from title cue.\"}}]}\n{\"cluster_id\":\"fabricated_statistics\",\"pattern\":\"Model invents numeric data not present in page\",\"root_cause\":\"Prompt does not explicitly forbid generation of statistics or numeric claims absent from webpage.\",\"suggestion\":\"Add instruction: 'Do not generate any numbers, statistics, or quantitative claims unless they appear explicitly in the webpage context.'\",\"bad_cases\":[{\"id\":\"a6fe249b-bb1f-465e-8cc6-369679aa66a6\",\"score\":0.5,\"issue\":\"Model fabricated identical 31.46% returns for all global indices.\",\"evidence\":{\"model_response_snippet\":\"所有指数涨幅为31.46%\",\"expected_behavior\":\"Only HSI 1-year return is 31.46%.\",\"why_it_failed\":\"Model generalized one statistic across all entries.\"}},{\"id\":\"5efc6052-5f97-43b7-a390-1d85f96dba00\",\"score\":0.75,\"issue\":\"Inserted unsupported percentages (30% users, 25% network issues).\",\"evidence\":{\"model_response_snippet\":\"约30%用户遇到问题\",\"expected_behavior\":\"Use only issues listed on page (browser, adblock, file corruption).\",\"why_it_failed\":\"Model fabricated research statistics to enrich answer.\"}}]}"

        logger.info("\n1. Testing LLM-based analysis with sample bad case clusters...")
        logger.info("   Input clusters: 3 clusters")

        result = await client.analyze_bad_cases_for_templates_llm(
            bad_case_analysis=sample_bad_cases,
            model_client=model_client,
            structure_doc=doc_content,
        )
        logger.info("\n   Result: ok=%s", result.get("ok"))

        if result.get("ok"):
            logger.info("   Total clusters analyzed: %s", result.get("total_clusters_analyzed"))
            logger.info("   Priority templates path: %s", result.get("priority_templates_path"))
            logger.info("   Analysis summary: %s", result.get("analysis_summary"))
            logger.info("\n   Priority templates:")
            for i, template in enumerate(result.get("priority_templates", []), 1):
                logger.info("\n   %s. %s", i, template.get("template_name"))
                logger.info("      Path: %s", template.get("template_path"))
                logger.info("      Purpose: %s", template.get("purpose", "N/A"))
                logger.info("      DO modify: %s", template.get("do_modify", "N/A"))
                logger.info("      DON'T modify: %s", template.get("dont_modify", "N/A"))
                logger.info("      Reason: %s", template.get("reason"))
                logger.info("      Clusters: %s", template.get("clusters"))
                if template.get("suggested_changes"):
                    logger.info("      Suggestions: %s", template.get("suggested_changes"))
        else:
            logger.error("   Error: %s", result.get("error"))

        await model_client.close()

    async def test_agent() -> None:
        """Test the full agent with LLM."""
        logger.info("\n%s", "=" * 60)
        logger.info("Testing LiquidStructureAgent (Full Agent)")
        logger.info("%s", "=" * 60)

        from prompt_auto_lab.orchestrator.state_store import SharedState

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

        shared_state = SharedState()
        agent = LiquidStructureAgent(
            model_client=model_client,
            shared_state=shared_state,
        ).agent

        task = (
            "Analyze the Liquid template structure starting from the entry point: "
            "Service.Chat/Inference/Prompts/Templates/syndication-api/responding/prompt.liquid\n\n"
            "Build the complete include tree and generate documentation for each template."
            "\"iteration.analysis_results_path\": \"autopromptlab2\\logs\\12d53b86\\experiments_3943ccd9-df70-48d6-afac-36986008b814\\analysis.txt\n\""
        )

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

        # Print final message - handle different message types
        last_msg = result.messages[-1]
        logger.info("\n%s", "=" * 60)
        logger.info("Final Result:")
        logger.info("%s", "=" * 60)

        if hasattr(last_msg, 'content'):
            content = last_msg.content
            if hasattr(content, 'thoughts'):
                logger.info("Thoughts: %s", content.thoughts)
                logger.info("Response: %s", content.response)
            else:
                logger.info("Content: %s", content)
        else:
            logger.info("Message: %s", last_msg)

        # Check SharedState for the doc path
        state = await shared_state.read(None)
        doc_path = state.get("liquid_structure_doc_path")
        logger.info("\nGenerated documentation path: %s", doc_path)

        if doc_path and Path(doc_path).exists():
            logger.info("Documentation file exists: Yes")
            logger.info("File size: %s bytes", Path(doc_path).stat().st_size)

        await model_client.close()

    # Choose which test to run
    if len(sys.argv) > 1:
        if sys.argv[1] == "agent":
            asyncio.run(test_agent())
        elif sys.argv[1] == "badcases":
            asyncio.run(test_analyze_bad_cases_llm())
        else:
            asyncio.run(test_tools())
    else:
        asyncio.run(test_tools())
