"""
DiagnosisRouter — maps StructuredDiagnosis to RemediationPlan.

Pure logic, no LLM calls. Uses fix templates indexed by RootCause.
"""

import logging
from typing import Any

from .kg_client import KnowledgeGraphClient
from .models import (
    EnrichedContext,
    FixTemplate,
    RemediationPlan,
    RootCause,
    StructuredDiagnosis,
)

logger = logging.getLogger("aigie.judge.diagnosis_router")


_FIX_TEMPLATES: dict[str, FixTemplate] = {
    # Hallucination
    "ungrounded_claim": FixTemplate(
        strategy="inject_instruction",
        fix_type="inject_system_message",
        template=(
            "IMPORTANT: Only state facts that are directly supported by the following context. "
            "Do not make claims without evidence.\n\nContext:\n{relevant_context}"
        ),
        placement="prepend_system",
        max_retries=2,
    ),
    "tool_result_mismatch": FixTemplate(
        strategy="inject_instruction",
        fix_type="inject_system_message",
        template=(
            "CORRECTION: The tool returned the following result. Use it exactly as provided, "
            "do not paraphrase or embellish:\n{correct_tool_result}"
        ),
        placement="prepend_system",
        max_retries=1,
    ),
    "fabricated_reference": FixTemplate(
        strategy="modify_prompt",
        fix_type="append_user",
        template=(
            "Do not cite any sources, studies, or references unless they were explicitly "
            "provided in the context above."
        ),
        placement="append_user",
        max_retries=1,
    ),
    "overconfident_assertion": FixTemplate(
        strategy="inject_instruction",
        fix_type="inject_system_message",
        template="Express uncertainty when making claims that are not directly verifiable from the provided context.",
        placement="prepend_system",
        max_retries=1,
    ),
    # Drift
    "task_deviation": FixTemplate(
        strategy="inject_instruction",
        fix_type="inject_system_message",
        template="REDIRECT: Return to the original task. The user asked: {original_query}",
        placement="prepend_system",
        max_retries=2,
    ),
    "topic_shift": FixTemplate(
        strategy="modify_prompt",
        fix_type="append_user",
        template="Stay on topic. The current topic is: {current_topic}",
        placement="append_user",
        max_retries=1,
    ),
    "role_deviation": FixTemplate(
        strategy="inject_instruction",
        fix_type="inject_system_message",
        template="Remember your role and behave consistently with your defined instructions.",
        placement="prepend_system",
        max_retries=1,
    ),
    # Safety
    "harmful_content": FixTemplate(
        strategy="retry_with_context",
        fix_type="retry_clean",
        template="Generate a safe, helpful response. Avoid harmful content.",
        placement="prepend_system",
        max_retries=1,
    ),
    "pii_leak": FixTemplate(
        strategy="modify_prompt",
        fix_type="inject_system_message",
        template=(
            "CRITICAL: Redact all personal information (names, emails, phone numbers, "
            "addresses, SSNs) from your response. Replace with [REDACTED]."
        ),
        placement="prepend_system",
        max_retries=1,
    ),
    "injection_success": FixTemplate(
        strategy="stop",
        fix_type="stop",
        template="Prompt injection detected. Stopping execution.",
        placement="prepend_system",
        max_retries=0,
    ),
    "policy_violation": FixTemplate(
        strategy="inject_instruction",
        fix_type="inject_system_message",
        template="Adhere to the defined policies. {policy_text}",
        placement="prepend_system",
        max_retries=1,
    ),
    # Quality
    "incomplete_response": FixTemplate(
        strategy="retry_with_context",
        fix_type="append_user",
        template="Your previous response was incomplete. Please provide a complete answer covering: {missing_aspects}",
        placement="append_user",
        max_retries=2,
    ),
    "incoherent_logic": FixTemplate(
        strategy="retry_with_context",
        fix_type="append_user",
        template="Your response had logical inconsistencies. Please restructure logically.",
        placement="append_user",
        max_retries=1,
    ),
    "irrelevant_content": FixTemplate(
        strategy="modify_prompt",
        fix_type="append_user",
        template="Only include information directly relevant to the query.",
        placement="append_user",
        max_retries=1,
    ),
    "vague_response": FixTemplate(
        strategy="inject_instruction",
        fix_type="append_user",
        template="Be specific and actionable. Instead of general advice, provide concrete steps.",
        placement="append_user",
        max_retries=1,
    ),
    # Tool usage
    "wrong_tool": FixTemplate(
        strategy="inject_instruction",
        fix_type="inject_system_message",
        template="Use the {correct_tool} tool for this task, not {used_tool}.",
        placement="prepend_system",
        max_retries=1,
    ),
    "bad_parameters": FixTemplate(
        strategy="inject_instruction",
        fix_type="inject_system_message",
        template="Correct the tool parameters. Expected format: {expected_params}",
        placement="prepend_system",
        max_retries=1,
    ),
    "misinterpreted_result": FixTemplate(
        strategy="inject_instruction",
        fix_type="inject_system_message",
        template="The tool returned: {actual_result}. Interpret it correctly.",
        placement="prepend_system",
        max_retries=1,
    ),
    "unhandled_tool_error": FixTemplate(
        strategy="retry_with_context",
        fix_type="inject_system_message",
        template="Handle the error gracefully: {error_message}",
        placement="prepend_system",
        max_retries=1,
    ),
    # Loop
    "stuck_loop": FixTemplate(
        strategy="inject_instruction",
        fix_type="inject_system_message",
        template="You are stuck in a loop. Take a completely different approach. Previous attempts that didn't work: {previous_attempts}",
        placement="prepend_system",
        max_retries=1,
        fallback=["fallback_model"],
    ),
    "circular_reasoning": FixTemplate(
        strategy="modify_prompt",
        fix_type="inject_system_message",
        template="Break the circular reasoning. Step back and approach the problem differently.",
        placement="prepend_system",
        max_retries=1,
        fallback=["skip_step"],
    ),
    "tool_retry_loop": FixTemplate(
        strategy="fallback_model",
        fix_type="inject_system_message",
        template="Stop retrying the same tool. Use an alternative approach.",
        placement="prepend_system",
        max_retries=1,
        fallback=["skip_step"],
    ),
}


class DiagnosisRouter:
    """Routes structured diagnosis to specific remediation plan. Pure logic, no LLM."""

    def __init__(self, kg_client: KnowledgeGraphClient | None = None):
        self._kg_client = kg_client

    def route(self, diagnosis: StructuredDiagnosis, ctx: EnrichedContext) -> RemediationPlan | None:
        """Create remediation plan from diagnosis. Returns None if no action needed."""
        if diagnosis.is_passing:
            return None

        # Check KG for a proven fix first
        if self._kg_client:
            kg_match = self._kg_client.find_cached_fix(
                diagnosis.issue_domain, diagnosis.root_cause.value
            )
            if kg_match and kg_match.is_high_confidence:
                logger.debug(
                    f"KG hit: {diagnosis.root_cause.value} → {kg_match.strategy} "
                    f"(success={kg_match.success_rate:.0%}, confidence={kg_match.confidence:.2f})"
                )
                return RemediationPlan(
                    diagnosis=diagnosis,
                    strategy=kg_match.strategy,
                    fix_type="kg_proven",
                    fix_content=kg_match.fix_content,
                    fix_placement="prepend_system",
                    max_retries=2,
                )

        template = _FIX_TEMPLATES.get(diagnosis.root_cause.value)

        if template:
            fix_content = template.render(diagnosis.fix_context)
            return RemediationPlan(
                diagnosis=diagnosis,
                strategy=template.strategy,
                fix_type=template.fix_type,
                fix_content=fix_content,
                fix_placement=template.placement,
                max_retries=template.max_retries,
                fallback_strategies=list(template.fallback),
            )

        # No template — use the diagnosis's own fix instruction
        return RemediationPlan(
            diagnosis=diagnosis,
            strategy=diagnosis.recommended_strategy or "inject_instruction",
            fix_type="inject_system_message",
            fix_content=diagnosis.fix_instruction or "Please improve your response.",
            fix_placement="prepend_system",
            max_retries=1,
        )
