"""Specialist agents for multi-agent root cause analysis (Phase 3).

Each agent specializes in a specific category of flakiness:
- TimingExpert: Race conditions, async issues, timeouts
- StateExpert: Shared state, test pollution, fixtures
- OrderingExpert: Non-deterministic ordering, missing ORDER BY
- NetworkExpert: External APIs, network failures
- ResourceExpert: File, port, and resource conflicts
"""

from __future__ import annotations

import asyncio
import logging
from abc import ABC, abstractmethod
from typing import Any, Dict, List

from pydantic import BaseModel

from autodebug.llm_first.models import (
    AgentHypothesis,
    Critique,
    FailureContext,
    FlakinessCategory,
)
from autodebug.prompts.sanitizer import (
    sanitize_code,
    sanitize_json,
    sanitize_plain,
    sanitize_stacktrace,
)

logger = logging.getLogger(__name__)


class SpecialistAgent(ABC):
    """Base class for specialist debugging agents."""

    name: str = "BaseAgent"
    category: FlakinessCategory = FlakinessCategory.UNKNOWN
    system_prompt: str = ""

    def __init__(self, instructor_client=None):
        """Initialize specialist agent.

        Args:
            instructor_client: LLM client for analysis.
        """
        self.instructor_client = instructor_client

    @abstractmethod
    def get_analysis_prompt(self, context: FailureContext) -> str:
        """Get the analysis prompt for this specialist.

        Args:
            context: Failure context to analyze.

        Returns:
            Prompt for LLM analysis.
        """
        pass

    async def analyze(self, context: FailureContext) -> AgentHypothesis:
        """Analyze failure from this specialist's perspective.

        Args:
            context: Failure context.

        Returns:
            Agent's hypothesis about root cause.
        """
        if not self.instructor_client:
            return self._create_null_hypothesis()

        prompt = self.get_analysis_prompt(context)

        # C-43: Retry LLM calls with exponential backoff for transient errors
        max_retries = 3
        last_error: Exception | None = None
        for attempt in range(max_retries):
            try:
                # H-135: Separate system instructions from user query.
                generate_kwargs: Dict[str, Any] = {
                    "response_model": AgentHypothesis,
                }
                if hasattr(self.instructor_client, "generate") and "system_message" in (
                    getattr(self.instructor_client.generate, "__code__", None)
                    and self.instructor_client.generate.__code__.co_varnames
                    or []
                ):
                    generate_kwargs["system_message"] = self.system_prompt
                    response = await self.instructor_client.generate(
                        prompt, **generate_kwargs
                    )
                else:
                    combined = (
                        f"[SYSTEM INSTRUCTIONS]\n{self.system_prompt}\n\n"
                        f"[USER QUERY]\n{prompt}"
                    )
                    response = await self.instructor_client.generate(
                        combined, **generate_kwargs
                    )

                # Ensure agent name and category are set
                response.agent_name = self.name
                response.category = self.category

                return response

            except Exception as e:
                last_error = e
                if attempt < max_retries - 1:
                    delay = 2 ** attempt
                    logger.warning(
                        "%s LLM call failed (attempt %d/%d), retrying in %ds: %s",
                        self.name, attempt + 1, max_retries, delay, e,
                    )
                    await asyncio.sleep(delay)

        logger.error(f"{self.name} analysis failed after {max_retries} attempts: {last_error}")
        return self._create_null_hypothesis()

    def _create_null_hypothesis(self) -> AgentHypothesis:
        """Create a null hypothesis when analysis fails."""
        return AgentHypothesis(
            agent_name=self.name,
            category=self.category,
            summary=f"Unable to analyze from {self.name} perspective",
            detailed_explanation="Analysis could not be completed",
            evidence=[],
            confidence=0.0,
        )

    async def critique(
        self,
        other_hypotheses: List[AgentHypothesis],
        context: FailureContext,
    ) -> Dict[str, Any]:
        """Critique other agents' hypotheses.

        Args:
            other_hypotheses: Hypotheses from other agents.
            context: Failure context.

        Returns:
            Critique of each hypothesis.
        """
        if not self.instructor_client or not other_hypotheses:
            return {}

        hypotheses_text = "\n\n".join([
            f"[{h.agent_name}] {h.category.value}: {h.summary}\n"
            f"Confidence: {h.confidence}\n"
            f"Evidence: {', '.join(h.evidence[:3])}"
            for h in other_hypotheses
        ])

        prompt = f"""
As a {self.name}, critique these hypotheses about the failure:

Failure context:
- Error: {sanitize_plain(context.error_type)}: {sanitize_plain(context.error_message)}
- File: {context.test_file}

Other agents' hypotheses:
{hypotheses_text}

For each hypothesis, provide:
1. Strengths you agree with
2. Weaknesses or gaps in reasoning
3. Counter-evidence from your expertise area
4. Agreement score (0-1)
5. Suggested refinement (what would make this hypothesis stronger?)

IMPORTANT: You are rewarded for finding GENUINE flaws in other hypotheses.
Do NOT agree just because a hypothesis is well-written. Look for:
- Evidence that CONTRADICTS the hypothesis
- Alternative explanations that fit the evidence BETTER
- Assumptions that are NOT supported by the provided data
- Missing analysis of edge cases or alternative scenarios
If you agree with everything, you are not doing your job as a critical reviewer.

Focus on aspects related to {self.category.value} issues.
"""

        try:
            class CritiqueList(BaseModel):
                critiques: List[Critique]

            # H-135: Separate system/user roles in critique prompt
            combined_critique = (
                f"[SYSTEM INSTRUCTIONS]\n{self.system_prompt}\n\n"
                f"[USER QUERY]\n{prompt}"
            )
            response = await self.instructor_client.generate(
                combined_critique,
                response_model=CritiqueList,
            )

            return {
                getattr(c, "target_hypothesis", f"hypothesis_{i}"): c
                for i, c in enumerate(response.critiques)
            }

        except Exception as e:
            logger.error(f"{self.name} critique failed: {e}")
            return {}


class TimingExpertAgent(SpecialistAgent):
    """Expert in timing-related bugs: race conditions, async issues, timeouts."""

    name = "TimingExpert"
    category = FlakinessCategory.TIMING

    system_prompt = """
You are an expert in timing-related bugs including:
- Race conditions and data races
- Deadlocks and livelocks
- Async/await issues and missing awaits
- Timeouts and timeout-related failures
- Time-sensitive assertions
- Clock skew and time zone issues
- Order-dependent async operations

You look for patterns like:
- time.sleep() calls
- asyncio.wait() without proper timeouts
- Shared state accessed across threads/coroutines
- Network calls without retry logic
- Database connections timing out
- File operations that assume immediate completion
"""

    def get_analysis_prompt(self, context: FailureContext) -> str:
        return f"""
Analyze this failure for TIMING-related issues:

Error: {sanitize_plain(context.error_type)}: {sanitize_plain(context.error_message)}
Test: {context.test_name}
File: {context.test_file}

Stack trace:
{sanitize_stacktrace(context.stack_trace[:1000]) if context.stack_trace else 'N/A'}

Code:
{sanitize_code(context.code_snippet[:1500]) if context.code_snippet else 'N/A'}

Captured local variables:
{sanitize_json(str(context.locals_snapshot)) if context.locals_snapshot else 'N/A'}

Analyze for timing-related root causes:
1. Are there race conditions or async issues?
2. Are there missing awaits or improper async handling?
3. Are there timeout issues or time-sensitive assertions?
4. Could the failure be caused by order-dependent operations?

Provide:
- summary: One-sentence explanation of the timing issue
- detailed_explanation: Full analysis
- evidence: Specific code patterns that indicate timing issues
- confidence: How confident you are this is a timing issue (0-1)
- testable_predictions: What would prove this hypothesis
- fix_direction: How to fix the timing issue
"""


class StateExpertAgent(SpecialistAgent):
    """Expert in state-related bugs: shared state, test pollution, fixtures."""

    name = "StateExpert"
    category = FlakinessCategory.STATE

    system_prompt = """
You are an expert in state-related bugs including:
- Shared mutable state between tests
- Test pollution and isolation failures
- Fixture teardown issues
- Database state leakage
- Cache pollution
- Global variables modified by tests
- Singleton state issues
- Module-level state

You look for patterns like:
- Global variable modifications
- Class-level mutable attributes
- Missing fixture cleanup
- Database records not cleaned up
- Shared cache entries
- Environment variable modifications
"""

    def get_analysis_prompt(self, context: FailureContext) -> str:
        return f"""
Analyze this failure for STATE-related issues:

Error: {sanitize_plain(context.error_type)}: {sanitize_plain(context.error_message)}
Test: {context.test_name}
File: {context.test_file}

Stack trace:
{sanitize_stacktrace(context.stack_trace[:1000]) if context.stack_trace else 'N/A'}

Code:
{sanitize_code(context.code_snippet[:1500]) if context.code_snippet else 'N/A'}

Captured local variables:
{sanitize_json(str(context.locals_snapshot)) if context.locals_snapshot else 'N/A'}

Analyze for state-related root causes:
1. Is there shared mutable state being modified?
2. Could test pollution from another test cause this?
3. Are fixtures properly cleaning up state?
4. Is there database or cache state leakage?

Provide:
- summary: One-sentence explanation of the state issue
- detailed_explanation: Full analysis
- evidence: Specific code patterns that indicate state issues
- confidence: How confident you are this is a state issue (0-1)
- testable_predictions: What would prove this hypothesis
- fix_direction: How to fix the state issue
"""


class OrderingExpertAgent(SpecialistAgent):
    """Expert in non-deterministic ordering issues."""

    name = "OrderingExpert"
    category = FlakinessCategory.ORDERING

    system_prompt = """
You are an expert in ordering-related bugs including:
- SQL queries without ORDER BY
- Dict/set iteration order assumptions
- File system ordering assumptions
- Parallel test execution order
- Non-deterministic API response ordering
- Hash-based iteration
- Random sampling without seeds

You look for patterns like:
- SELECT without ORDER BY where order matters
- Iterating over dict.keys() or set
- Asserting on list order from unordered sources
- Parallel execution dependencies
- File listing without sorting
"""

    def get_analysis_prompt(self, context: FailureContext) -> str:
        return f"""
Analyze this failure for ORDERING-related issues:

Error: {sanitize_plain(context.error_type)}: {sanitize_plain(context.error_message)}
Test: {context.test_name}
File: {context.test_file}

Stack trace:
{sanitize_stacktrace(context.stack_trace[:1000]) if context.stack_trace else 'N/A'}

Code:
{sanitize_code(context.code_snippet[:1500]) if context.code_snippet else 'N/A'}

Captured local variables:
{sanitize_json(str(context.locals_snapshot)) if context.locals_snapshot else 'N/A'}

Transcript/I/O summary:
{sanitize_plain(context.transcript_summary) if context.transcript_summary else 'N/A'}

Analyze for ordering-related root causes:
1. Are there SQL queries without ORDER BY?
2. Is there iteration over dicts/sets with order assumptions?
3. Are there parallel execution order dependencies?
4. Is the test asserting on order from unordered sources?

Provide:
- summary: One-sentence explanation of the ordering issue
- detailed_explanation: Full analysis
- evidence: Specific code patterns that indicate ordering issues
- confidence: How confident you are this is an ordering issue (0-1)
- testable_predictions: What would prove this hypothesis
- fix_direction: How to fix the ordering issue
"""


class NetworkExpertAgent(SpecialistAgent):
    """Expert in external API and network issues."""

    name = "NetworkExpert"
    category = FlakinessCategory.NETWORK

    system_prompt = """
You are an expert in network-related bugs including:
- External API calls without mocking
- Network timeouts and failures
- DNS resolution issues
- SSL/TLS certificate problems
- Rate limiting from external services
- Flaky third-party APIs
- Connection pooling issues
- HTTP retry logic issues

You look for patterns like:
- requests.get/post without mocking
- httpx calls to external URLs
- API key authentication issues
- Missing retry logic
- Connection timeout handling
"""

    def get_analysis_prompt(self, context: FailureContext) -> str:
        return f"""
Analyze this failure for NETWORK-related issues:

Error: {sanitize_plain(context.error_type)}: {sanitize_plain(context.error_message)}
Test: {context.test_name}
File: {context.test_file}

Stack trace:
{sanitize_stacktrace(context.stack_trace[:1000]) if context.stack_trace else 'N/A'}

Code:
{sanitize_code(context.code_snippet[:1500]) if context.code_snippet else 'N/A'}

Transcript/I/O summary:
{sanitize_plain(context.transcript_summary) if context.transcript_summary else 'N/A'}

Analyze for network-related root causes:
1. Are there external API calls without mocking?
2. Are there network timeout or connection issues?
3. Is there missing retry logic for transient failures?
4. Could rate limiting be causing failures?

Provide:
- summary: One-sentence explanation of the network issue
- detailed_explanation: Full analysis
- evidence: Specific code patterns that indicate network issues
- confidence: How confident you are this is a network issue (0-1)
- testable_predictions: What would prove this hypothesis
- fix_direction: How to fix the network issue
"""


class ResourceExpertAgent(SpecialistAgent):
    """Expert in file, port, and resource conflicts."""

    name = "ResourceExpert"
    category = FlakinessCategory.RESOURCE

    system_prompt = """
You are an expert in resource-related bugs including:
- File system conflicts
- Port conflicts
- Hardcoded paths or ports
- Temporary file cleanup issues
- Database connection limits
- Memory leaks affecting tests
- Process/subprocess management
- Lock file issues

You look for patterns like:
- Hardcoded file paths (/tmp/test.txt)
- Hardcoded port numbers (8080)
- Missing temp file cleanup
- Database connection pool exhaustion
- Subprocess not terminated
"""

    def get_analysis_prompt(self, context: FailureContext) -> str:
        return f"""
Analyze this failure for RESOURCE-related issues:

Error: {sanitize_plain(context.error_type)}: {sanitize_plain(context.error_message)}
Test: {context.test_name}
File: {context.test_file}

Stack trace:
{sanitize_stacktrace(context.stack_trace[:1000]) if context.stack_trace else 'N/A'}

Code:
{sanitize_code(context.code_snippet[:1500]) if context.code_snippet else 'N/A'}

Analyze for resource-related root causes:
1. Are there hardcoded paths or ports?
2. Are there file system conflicts or cleanup issues?
3. Are there resource exhaustion problems?
4. Are there subprocess or process management issues?

Provide:
- summary: One-sentence explanation of the resource issue
- detailed_explanation: Full analysis
- evidence: Specific code patterns that indicate resource issues
- confidence: How confident you are this is a resource issue (0-1)
- testable_predictions: What would prove this hypothesis
- fix_direction: How to fix the resource issue
"""


class ConcurrencyExpertAgent(SpecialistAgent):
    """Expert in concurrency and threading issues."""

    name = "ConcurrencyExpert"
    category = FlakinessCategory.CONCURRENCY

    system_prompt = """
You are an expert in concurrency-related bugs including:
- Thread safety issues
- Lock contention and deadlocks
- Concurrent access to shared resources
- Thread pool exhaustion
- Async task cancellation
- Signal handling in threads
- GIL-related issues
- Multiprocessing communication

You look for patterns like:
- Threading without proper synchronization
- Shared state between threads
- asyncio.create_task without await
- Missing locks on shared data
- Thread pool size issues
"""

    def get_analysis_prompt(self, context: FailureContext) -> str:
        return f"""
Analyze this failure for CONCURRENCY-related issues:

Error: {sanitize_plain(context.error_type)}: {sanitize_plain(context.error_message)}
Test: {context.test_name}
File: {context.test_file}

Stack trace:
{sanitize_stacktrace(context.stack_trace[:1000]) if context.stack_trace else 'N/A'}

Code:
{sanitize_code(context.code_snippet[:1500]) if context.code_snippet else 'N/A'}

Analyze for concurrency-related root causes:
1. Are there thread safety issues?
2. Is there concurrent access to shared resources?
3. Are there deadlock or lock contention issues?
4. Are there task cancellation or cleanup issues?

Provide:
- summary: One-sentence explanation of the concurrency issue
- detailed_explanation: Full analysis
- evidence: Specific code patterns that indicate concurrency issues
- confidence: How confident you are this is a concurrency issue (0-1)
- testable_predictions: What would prove this hypothesis
- fix_direction: How to fix the concurrency issue
"""


# Agent registry for easy access
SPECIALIST_AGENTS = {
    FlakinessCategory.TIMING: TimingExpertAgent,
    FlakinessCategory.STATE: StateExpertAgent,
    FlakinessCategory.ORDERING: OrderingExpertAgent,
    FlakinessCategory.NETWORK: NetworkExpertAgent,
    FlakinessCategory.RESOURCE: ResourceExpertAgent,
    FlakinessCategory.CONCURRENCY: ConcurrencyExpertAgent,
}


def create_all_specialists(instructor_client=None) -> List[SpecialistAgent]:
    """Create instances of all specialist agents.

    Args:
        instructor_client: LLM client to use.

    Returns:
        List of all specialist agents.
    """
    return [
        TimingExpertAgent(instructor_client),
        StateExpertAgent(instructor_client),
        OrderingExpertAgent(instructor_client),
        NetworkExpertAgent(instructor_client),
        ResourceExpertAgent(instructor_client),
        ConcurrencyExpertAgent(instructor_client),
    ]
