"""Analysis tools for the autonomous debugging agent.

This module provides tools for analyzing failures:
- Reading execution transcripts from storage
- Classifying errors into flakiness categories
"""

from __future__ import annotations

import asyncio
import json
import logging
from typing import TYPE_CHECKING, Any, Dict, List, Optional

import requests
import zstandard as zstd

from autodebug.llm_first.autonomous import BaseTool
from autodebug.llm_first.models import FlakinessCategory, FlakinessPrediction

if TYPE_CHECKING:
    pass

logger = logging.getLogger(__name__)


class ReadTranscriptTool(BaseTool):
    """Read I/O transcript from a test run.

    This tool loads captured I/O transcripts (HTTP, SQL, Redis, etc.)
    from storage to help analyze what the test was doing when it failed.
    """

    name = "read_transcript"
    description = """
    Read the captured I/O transcript from a test run.

    Parameters:
    - run_id: ID of the run to read transcript for
    - dep_type: Type of dependency transcript ("HTTP", "SQL", "REDIS", "RNG", "TIME")
    - limit: Maximum number of records to return (default 100)
    - start_index: Start index for pagination (default 0)

    Returns: List of transcript records with timestamps, requests, and responses

    Use this tool to:
    - See what HTTP requests the test made
    - Examine SQL queries and their results
    - Check Redis operations
    - Find the divergence point in recorded behavior
    """

    def __init__(
        self,
        db=None,
        storage_client=None,
        supabase_url: Optional[str] = None,
        supabase_key: Optional[str] = None,
    ):
        """Initialize ReadTranscriptTool.

        Args:
            db: Database client for run metadata.
            storage_client: Storage client for transcript files.
            supabase_url: Supabase URL for storage.
            supabase_key: Supabase service role key.
        """
        self.db = db
        self.storage_client = storage_client
        self.supabase_url = supabase_url
        self.supabase_key = supabase_key

    async def execute(
        self,
        run_id: str,
        dep_type: str = "HTTP",
        limit: int = 100,
        start_index: int = 0,
    ) -> List[Dict[str, Any]]:
        """Read transcript records for a run.

        Args:
            run_id: Run ID to read transcript for.
            dep_type: Type of dependency ("HTTP", "SQL", "REDIS", "RNG", "TIME").
            limit: Maximum records to return.
            start_index: Starting index for pagination.

        Returns:
            List of transcript records.
        """
        try:
            # Normalize dep_type
            dep_type = dep_type.upper()
            valid_types = {"HTTP", "SQL", "REDIS", "RNG", "TIME"}

            if dep_type not in valid_types:
                return [{
                    "error": f"Invalid dep_type: {dep_type}. "
                             f"Valid types: {', '.join(valid_types)}"
                }]

            # Get transcript from storage
            records = await self._fetch_transcript(run_id, dep_type)

            if not records:
                return [{
                    "info": f"No {dep_type} transcript found for run {run_id}"
                }]

            # Apply pagination
            paginated = records[start_index:start_index + limit]

            # Format records for readability
            return self._format_records(paginated, dep_type)

        except Exception as e:
            logger.error(f"Error reading transcript: {e}")
            return [{"error": f"Failed to read transcript: {e}"}]

    async def _fetch_transcript(
        self,
        run_id: str,
        dep_type: str,
    ) -> List[Dict[str, Any]]:
        """Fetch transcript from storage.

        Args:
            run_id: Run ID.
            dep_type: Dependency type.

        Returns:
            List of raw transcript records.
        """
        # Map dep_type to file name
        file_map = {
            "HTTP": "deps_http.jsonl.zst",
            "SQL": "deps_sql.jsonl.zst",
            "REDIS": "deps_redis.jsonl.zst",
            "RNG": "time_rng.jsonl.zst",
            "TIME": "time_rng.jsonl.zst",
        }

        filename = file_map.get(dep_type)
        if not filename:
            return []

        # Try to fetch from storage
        content = await self._download_from_storage(run_id, filename)

        if not content:
            return []

        # Parse JSONL content
        return self._parse_jsonl(content)

    async def _download_from_storage(
        self,
        run_id: str,
        filename: str,
    ) -> Optional[bytes]:
        """Download file from Supabase storage.

        Args:
            run_id: Run ID.
            filename: File name to download.

        Returns:
            Decompressed file content or None.
        """
        # Use storage client if available
        if self.storage_client:
            try:
                return await self.storage_client.download(
                    bucket="run-bundles",
                    path=f"{run_id}/{filename}",
                )
            except Exception as e:
                logger.warning(f"Storage client download failed: {e}")

        # Fall back to direct Supabase API
        if self.supabase_url and self.supabase_key:
            try:
                url = f"{self.supabase_url}/storage/v1/object/run-bundles/{run_id}/{filename}"
                headers = {
                    "Authorization": f"Bearer {self.supabase_key}",
                    "apikey": self.supabase_key,
                }

                loop = asyncio.get_event_loop()
                response = await loop.run_in_executor(
                    None,
                    lambda: requests.get(url, headers=headers, timeout=30),
                )

                if response.status_code == 200:
                    # Decompress zstd
                    if filename.endswith(".zst"):
                        dctx = zstd.ZstdDecompressor()
                        return dctx.decompress(response.content)
                    return response.content

            except Exception as e:
                logger.warning(f"Direct storage download failed: {e}")

        # Try database if available
        if self.db:
            try:
                # Query for transcript records from database
                query = """
                    SELECT data FROM run_transcripts
                    WHERE run_id = $1 AND dep_type = $2
                    ORDER BY sequence_num
                """
                results = await self.db.fetch(query, run_id, filename.split(".")[0])
                if results:
                    # Combine all data
                    all_data = [r["data"] for r in results]
                    return json.dumps(all_data).encode()
            except Exception as e:
                logger.warning(f"Database transcript fetch failed: {e}")

        return None

    def _parse_jsonl(self, content: bytes) -> List[Dict[str, Any]]:
        """Parse JSONL content.

        Args:
            content: Raw JSONL bytes.

        Returns:
            List of parsed records.
        """
        records = []
        text = content.decode("utf-8", errors="replace")

        for line in text.strip().split("\n"):
            if not line.strip():
                continue
            try:
                records.append(json.loads(line))
            except json.JSONDecodeError:
                continue

        return records

    def _format_records(
        self,
        records: List[Dict[str, Any]],
        dep_type: str,
    ) -> List[Dict[str, Any]]:
        """Format records for readability.

        Args:
            records: Raw records.
            dep_type: Dependency type.

        Returns:
            Formatted records.
        """
        formatted = []

        for i, record in enumerate(records):
            formatted_record = {
                "index": i,
                "timestamp": record.get("timestamp", record.get("ts")),
            }

            if dep_type == "HTTP":
                formatted_record.update({
                    "method": record.get("method"),
                    "url": record.get("url"),
                    "status": record.get("status"),
                    "request_headers": self._truncate_dict(
                        record.get("request_headers", {}), 3
                    ),
                    "response_preview": self._truncate_str(
                        str(record.get("response_body", "")), 200
                    ),
                    "duration_ms": record.get("duration_ms"),
                })

            elif dep_type == "SQL":
                formatted_record.update({
                    "query": self._truncate_str(record.get("query", ""), 500),
                    "params": record.get("params"),
                    "row_count": record.get("row_count"),
                    "duration_ms": record.get("duration_ms"),
                })

            elif dep_type == "REDIS":
                formatted_record.update({
                    "command": record.get("command"),
                    "args": record.get("args"),
                    "result_preview": self._truncate_str(
                        str(record.get("result", "")), 200
                    ),
                })

            elif dep_type in ("RNG", "TIME"):
                formatted_record.update({
                    "type": record.get("type"),
                    "value": record.get("value"),
                    "source": record.get("source"),
                })

            formatted.append(formatted_record)

        return formatted

    def _truncate_str(self, s: str, max_len: int) -> str:
        """Truncate string to max length."""
        if len(s) <= max_len:
            return s
        return s[:max_len] + "..."

    def _truncate_dict(self, d: Dict, max_keys: int) -> Dict:
        """Truncate dict to max keys."""
        if len(d) <= max_keys:
            return d
        keys = list(d.keys())[:max_keys]
        result = {k: d[k] for k in keys}
        result["..."] = f"({len(d) - max_keys} more)"
        return result


class ClassifyErrorTool(BaseTool):
    """Classify error into flakiness category.

    This tool uses the FlakinessCategorizer to classify errors into
    categories like TIMING, ORDERING, RESOURCE, STATE, etc.
    """

    name = "classify_error"
    description = """
    Classify an error into a flakiness category.

    Parameters:
    - error_type: Exception type (e.g., "AssertionError", "TimeoutError")
    - error_message: The error message
    - stack_trace: Optional stack trace for more context
    - code_snippet: Optional code snippet for analysis

    Returns: Category classification with confidence and explanation

    Categories:
    - TIMING: Race conditions, timeouts, async issues
    - ORDERING: Non-deterministic ordering of results
    - RESOURCE: Resource exhaustion, connection issues
    - STATE: Shared state pollution
    - NETWORK: Network-related flakiness
    - CONCURRENCY: Threading/multiprocessing issues
    - UNKNOWN: Cannot determine category

    Use this tool to understand what type of fix is needed.
    """

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

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

    async def execute(
        self,
        error_type: str,
        error_message: str,
        stack_trace: str = "",
        code_snippet: str = "",
    ) -> Dict[str, Any]:
        """Classify error into flakiness category.

        Args:
            error_type: Exception type.
            error_message: Error message.
            stack_trace: Optional stack trace.
            code_snippet: Optional code snippet.

        Returns:
            Classification result with category and confidence.
        """
        try:
            # Try LLM-based classification first
            if self.instructor_client:
                return await self._classify_with_llm(
                    error_type, error_message, stack_trace, code_snippet
                )

            # Fall back to rule-based classification
            return self._classify_rule_based(
                error_type, error_message, stack_trace, code_snippet
            )

        except Exception as e:
            logger.error(f"Classification error: {e}")
            return {
                "category": FlakinessCategory.UNKNOWN.value,
                "confidence": 0.0,
                "explanation": f"Classification failed: {e}",
                "risk_factors": [],
            }

    async def _classify_with_llm(
        self,
        error_type: str,
        error_message: str,
        stack_trace: str,
        code_snippet: str,
    ) -> Dict[str, Any]:
        """Classify using LLM.

        Args:
            error_type: Exception type.
            error_message: Error message.
            stack_trace: Stack trace.
            code_snippet: Code snippet.

        Returns:
            Classification result.
        """
        prompt = f"""
Classify this test failure into a flakiness category.

Error Type: {error_type}
Error Message: {error_message[:500]}

Stack Trace (last 20 lines):
{self._truncate_stack(stack_trace, 20)}

Code Context:
{code_snippet[:1000] if code_snippet else "Not available"}

Categories:
- TIMING: Race conditions, timeouts, async/await issues, sleep dependencies
- ORDERING: Non-deterministic ordering (dict/set iteration, SQL without ORDER BY)
- RESOURCE: Resource exhaustion, file/port conflicts, connection limits
- STATE: Shared state pollution, global variables, database state
- NETWORK: Network flakiness, DNS, external service issues
- CONCURRENCY: Threading issues, deadlocks, process synchronization
- UNKNOWN: Cannot determine

Analyze the error and classify it. Consider:
1. What is the direct cause of the error?
2. What environmental factors could cause non-determinism?
3. Is this a true flaky test or a consistent bug?

Respond with:
- category: The most likely category
- confidence: How confident you are (0.0-1.0)
- explanation: Why you chose this category
- risk_factors: List of specific risk factors identified
- recommendations: Suggested fix directions
"""

        response = await self.instructor_client.generate(
            prompt,
            response_model=FlakinessPrediction,
        )

        return {
            "category": response.category.value,
            "confidence": response.category_confidence,
            "is_flaky": response.is_flaky,
            "flaky_probability": response.flaky_probability,
            "explanation": response.explanation,
            "risk_factors": response.risk_factors,
            "recommendations": response.recommendations,
        }

    def _classify_rule_based(
        self,
        error_type: str,
        error_message: str,
        stack_trace: str,
        code_snippet: str,
    ) -> Dict[str, Any]:
        """Rule-based fallback classification.

        Args:
            error_type: Exception type.
            error_message: Error message.
            stack_trace: Stack trace.
            code_snippet: Code snippet.

        Returns:
            Classification result.
        """
        combined = f"{error_type} {error_message} {stack_trace} {code_snippet}".lower()

        # Timing patterns
        timing_patterns = [
            "timeout", "timed out", "asyncio", "await", "sleep",
            "deadline", "expired", "too slow", "took too long",
            "event loop", "async def", "coroutine",
        ]

        # Ordering patterns
        ordering_patterns = [
            "order", "sorted", "dict", "set", "hash",
            "random", "shuffle", "different order",
            "expected:", "actual:",  # assertion mismatch
        ]

        # Resource patterns
        resource_patterns = [
            "connection refused", "address in use", "port",
            "file", "disk", "memory", "oom", "resource",
            "exhausted", "limit", "quota", "too many",
        ]

        # State patterns
        state_patterns = [
            "state", "global", "singleton", "shared",
            "database", "cache", "session", "fixture",
            "setup", "teardown", "cleanup",
        ]

        # Network patterns
        network_patterns = [
            "network", "dns", "http", "socket", "connection",
            "unreachable", "refused", "reset", "ssl", "tls",
        ]

        # Concurrency patterns
        concurrency_patterns = [
            "thread", "lock", "deadlock", "race",
            "concurrent", "parallel", "process", "fork",
            "mutex", "semaphore",
        ]

        # Score each category
        scores = {
            FlakinessCategory.TIMING: sum(1 for p in timing_patterns if p in combined),
            FlakinessCategory.ORDERING: sum(1 for p in ordering_patterns if p in combined),
            FlakinessCategory.RESOURCE: sum(1 for p in resource_patterns if p in combined),
            FlakinessCategory.STATE: sum(1 for p in state_patterns if p in combined),
            FlakinessCategory.NETWORK: sum(1 for p in network_patterns if p in combined),
            FlakinessCategory.CONCURRENCY: sum(1 for p in concurrency_patterns if p in combined),
        }

        # Special case: TimeoutError
        if "timeout" in error_type.lower():
            scores[FlakinessCategory.TIMING] += 3

        # Find best category
        best_category = max(scores.items(), key=lambda x: x[1])

        if best_category[1] == 0:
            category = FlakinessCategory.UNKNOWN
            confidence = 0.3
        else:
            category = best_category[0]
            total_score = sum(scores.values())
            confidence = min(best_category[1] / max(total_score, 1), 0.9)

        # Generate explanation
        matched_patterns = []
        pattern_maps = {
            FlakinessCategory.TIMING: timing_patterns,
            FlakinessCategory.ORDERING: ordering_patterns,
            FlakinessCategory.RESOURCE: resource_patterns,
            FlakinessCategory.STATE: state_patterns,
            FlakinessCategory.NETWORK: network_patterns,
            FlakinessCategory.CONCURRENCY: concurrency_patterns,
        }

        if category in pattern_maps:
            matched_patterns = [p for p in pattern_maps[category] if p in combined]

        explanation = f"Classified as {category.value} based on patterns: {', '.join(matched_patterns[:5])}"

        return {
            "category": category.value,
            "confidence": confidence,
            "is_flaky": confidence > 0.4,
            "flaky_probability": confidence,
            "explanation": explanation,
            "risk_factors": matched_patterns[:5],
            "recommendations": self._get_recommendations(category),
        }

    def _get_recommendations(self, category: FlakinessCategory) -> List[str]:
        """Get recommendations based on category.

        Args:
            category: Flakiness category.

        Returns:
            List of recommendations.
        """
        recommendations = {
            FlakinessCategory.TIMING: [
                "Add explicit waits or polling instead of fixed sleeps",
                "Use async/await properly with timeout handling",
                "Mock time-dependent operations in tests",
            ],
            FlakinessCategory.ORDERING: [
                "Add ORDER BY to SQL queries",
                "Sort results before comparison",
                "Use ordered collections or convert to sorted lists",
            ],
            FlakinessCategory.RESOURCE: [
                "Ensure proper resource cleanup in finally blocks",
                "Use unique ports/files per test",
                "Implement retry with backoff for transient failures",
            ],
            FlakinessCategory.STATE: [
                "Use fresh fixtures per test (function scope)",
                "Reset global state in teardown",
                "Isolate database state between tests",
            ],
            FlakinessCategory.NETWORK: [
                "Mock external HTTP calls",
                "Use testcontainers for local services",
                "Add retry logic for network operations",
            ],
            FlakinessCategory.CONCURRENCY: [
                "Use proper synchronization primitives",
                "Avoid shared mutable state between threads",
                "Use thread-safe data structures",
            ],
            FlakinessCategory.UNKNOWN: [
                "Gather more context about the failure",
                "Check if failure is consistent or intermittent",
                "Review recent code changes",
            ],
        }

        return recommendations.get(category, recommendations[FlakinessCategory.UNKNOWN])

    def _truncate_stack(self, stack_trace: str, max_lines: int) -> str:
        """Truncate stack trace to last N lines."""
        lines = stack_trace.strip().split("\n")
        if len(lines) <= max_lines:
            return stack_trace
        return "\n".join(lines[-max_lines:])
