# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Log aggregation for Claude Agent SDK isolated environment.

This module provides automatic discovery and tailing of log files created
in the `.claude` directory during prompt evaluation. Log content is emitted
as DEBUG-level structured log messages for observability.
"""

from __future__ import annotations

import asyncio
import contextlib
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from dataclasses import dataclass, field
from pathlib import Path

from ...runtime.logging import StructuredLogger, get_logger

__all__ = [
    "ClaudeLogAggregator",
]

logger: StructuredLogger = get_logger(
    __name__, context={"component": "claude_log_aggregator"}
)

# Files and directories to exclude from aggregation (generated by WINK)
_EXCLUDED_PATHS: frozenset[str] = frozenset(
    {
        "settings.json",
        "skills",
    }
)

# File extensions that are likely log files
_LOG_EXTENSIONS: frozenset[str] = frozenset(
    {
        ".log",
        ".jsonl",
        ".txt",
        ".json",
    }
)

# Maximum bytes to read per file per poll cycle.
# At 500ms poll intervals, this limits throughput to ~128KB/s per file.
# Log files growing faster than this rate will have the aggregator fall behind.
_MAX_READ_BYTES: int = 64 * 1024  # 64KB

# Poll interval for file discovery
_POLL_INTERVAL_SECONDS: float = 0.5


@dataclass(slots=True)
class _FileState:
    """Tracks the read state of a single file.

    Note: Not frozen because position, inode, and partial_line are mutated
    during file reading. All mutations occur within a single async task
    (no concurrent access) so thread safety is not a concern.
    """

    path: Path
    position: int = 0
    """Current read position in the file."""

    inode: int = 0
    """Inode number to detect file rotation."""

    partial_line: str = ""
    """Buffer for incomplete line from previous read."""


@dataclass(slots=True)
class ClaudeLogAggregator:
    """Aggregates log files from Claude's .claude directory during execution.

    Monitors the `.claude` directory within the ephemeral home for new files
    and new content in existing files. Discovered content is emitted as
    DEBUG-level structured log messages with appropriate context.

    The aggregator runs as an async context manager that polls for changes
    while the SDK query is executing.

    Example:
        >>> aggregator = ClaudeLogAggregator(
        ...     claude_dir=Path("/tmp/claude-agent-xxx/.claude"),
        ...     prompt_name="my-prompt",
        ... )
        >>> async with aggregator.run():
        ...     # SDK query executes here
        ...     result = await sdk.query(...)
    """

    claude_dir: Path
    """Path to the .claude directory within ephemeral home."""

    prompt_name: str
    """Name of the prompt being evaluated (for log context)."""

    poll_interval: float = _POLL_INTERVAL_SECONDS
    """Interval between file system polls in seconds."""

    _file_states: dict[Path, _FileState] = field(default_factory=dict)
    """Tracks read state for each discovered file."""

    _running: bool = field(default=False, init=False)
    """Whether the aggregator is currently running."""

    _total_bytes_read: int = field(default=0, init=False)
    """Total bytes read across all files."""

    _total_lines_emitted: int = field(default=0, init=False)
    """Total log lines emitted."""

    @asynccontextmanager
    async def run(self) -> AsyncIterator[None]:
        """Run the log aggregator as an async context manager.

        Yields control to the caller while polling for log file changes
        in the background. Stops polling when the context exits.

        Yields:
            None - control is yielded to allow SDK execution.
        """
        self._running = True
        poll_task: asyncio.Task[None] | None = None

        logger.debug(
            "claude_log_aggregator.start",
            event="log_aggregator.start",
            context={
                "claude_dir": str(self.claude_dir),
                "prompt_name": self.prompt_name,
                "poll_interval": self.poll_interval,
            },
        )

        try:
            # Start background polling task
            poll_task = asyncio.create_task(self._poll_loop())

            # Yield control to caller
            yield
        finally:
            # Stop polling
            self._running = False

            if poll_task is not None:  # pragma: no branch
                poll_task.cancel()
                with contextlib.suppress(asyncio.CancelledError):
                    await poll_task

            # Final poll to capture any remaining content
            await self._poll_once()

            logger.debug(
                "claude_log_aggregator.stop",
                event="log_aggregator.stop",
                context={
                    "claude_dir": str(self.claude_dir),
                    "prompt_name": self.prompt_name,
                    "total_bytes_read": self._total_bytes_read,
                    "total_lines_emitted": self._total_lines_emitted,
                    "files_monitored": len(self._file_states),
                },
            )

    async def _poll_loop(self) -> None:
        """Background polling loop for file changes."""
        while self._running:  # pragma: no branch
            await self._poll_once()
            await asyncio.sleep(self.poll_interval)

    async def _poll_once(self) -> None:
        """Perform a single poll cycle: discover files and read new content."""
        if not self.claude_dir.exists():
            return

        # Discover new files
        await self._discover_files()

        # Read new content from tracked files
        await self._read_new_content()

        # Clean up state for deleted files to prevent unbounded memory growth
        self._cleanup_deleted_files()

    async def _discover_files(self) -> None:
        """Discover new log files in the .claude directory."""
        try:
            for path in self.claude_dir.rglob("*"):
                # Skip directories
                if path.is_dir():
                    continue

                # Skip excluded paths
                relative = path.relative_to(self.claude_dir)
                if self._is_excluded(relative):
                    continue

                # Skip non-log extensions
                if not self._is_log_file(path):
                    continue

                # Track new files
                if path not in self._file_states:
                    try:
                        stat = path.stat()
                        self._file_states[path] = _FileState(
                            path=path,
                            position=0,
                            inode=stat.st_ino,
                        )
                        logger.debug(
                            "claude_log_aggregator.file_discovered",
                            event="log_aggregator.file_discovered",
                            context={
                                "file": str(relative),
                                "prompt_name": self.prompt_name,
                                "size_bytes": stat.st_size,
                            },
                        )
                    except OSError:  # pragma: no cover
                        # File may have been deleted between rglob and stat
                        pass
        except OSError:  # pragma: no cover
            # Directory may not exist or be inaccessible
            pass

    async def _read_new_content(self) -> None:
        """Read new content from all tracked files."""
        for file_state in list(self._file_states.values()):
            await self._read_file_content(file_state)

    def _cleanup_deleted_files(self) -> None:
        """Remove state for files that no longer exist.

        Prevents unbounded memory growth when log files are created and deleted
        repeatedly during long-running sessions.
        """
        deleted = [path for path in self._file_states if not path.exists()]
        for path in deleted:
            del self._file_states[path]

    async def _read_file_content(self, state: _FileState) -> None:
        """Read new content from a single file."""
        try:
            # Check if file still exists and hasn't been rotated
            if not state.path.exists():
                return

            stat = state.path.stat()

            # Detect file rotation (inode changed) or in-place truncation
            if stat.st_ino != state.inode:
                # File was rotated (new inode)
                state.position = 0
                state.inode = stat.st_ino
                state.partial_line = ""
            elif stat.st_size < state.position:
                # File was truncated in place (e.g., copytruncate)
                state.position = 0
                state.partial_line = ""

            # No new content
            if stat.st_size <= state.position:
                return

            # Read new content
            bytes_to_read = min(stat.st_size - state.position, _MAX_READ_BYTES)

            # Use run_in_executor for file I/O to avoid blocking
            content = await asyncio.get_running_loop().run_in_executor(
                None,
                ClaudeLogAggregator._read_bytes,
                state.path,
                state.position,
                bytes_to_read,
            )

            if content:  # pragma: no branch
                state.position += len(content)
                self._total_bytes_read += len(content)

                # Emit log lines
                relative_path = state.path.relative_to(self.claude_dir)
                await self._emit_content(relative_path, content, state)

        except OSError:  # pragma: no cover
            # File may have been deleted or become inaccessible
            pass

    @staticmethod
    def _read_bytes(path: Path, offset: int, count: int) -> bytes:
        """Read bytes from file at offset (runs in executor)."""
        with path.open("rb") as f:
            f.seek(offset)
            return f.read(count)

    async def _emit_content(
        self, relative_path: Path, content: bytes, state: _FileState
    ) -> None:
        """Emit log lines from file content.

        Handles partial lines by buffering content that doesn't end with a
        newline until the next read completes the line.

        Args:
            relative_path: Path relative to .claude directory for logging.
            content: Raw bytes read from the file.
            state: File state containing partial line buffer.
        """
        # Use errors="replace" to handle invalid UTF-8 gracefully
        text = content.decode("utf-8", errors="replace")

        # Prepend any partial line from previous read
        if state.partial_line:
            text = state.partial_line + text
            state.partial_line = ""

        # Check if content ends with a newline
        ends_with_newline = text.endswith("\n") or text.endswith("\r")

        # Split into lines
        lines = text.splitlines()

        # If doesn't end with newline, buffer the last part as partial
        if lines and not ends_with_newline:
            state.partial_line = lines.pop()

        # Emit complete lines
        for line in lines:
            if not line.strip():
                continue

            self._total_lines_emitted += 1

            logger.debug(
                "claude_log_aggregator.log_line",
                event="log_aggregator.log_line",
                context={
                    "file": str(relative_path),
                    "prompt_name": self.prompt_name,
                    "content": line.rstrip(),
                    "sequence_number": self._total_lines_emitted,
                },
            )

    @staticmethod
    def _is_excluded(relative: Path) -> bool:
        """Check if a path should be excluded from monitoring."""
        # Check first component against exclusion set
        if relative.parts and relative.parts[0] in _EXCLUDED_PATHS:
            return True
        # Check exact match
        return str(relative) in _EXCLUDED_PATHS

    @staticmethod
    def _is_log_file(path: Path) -> bool:
        """Check if a file appears to be a log file."""
        # Check extension
        suffix = path.suffix.lower()
        if suffix in _LOG_EXTENSIONS:
            return True

        # Check common log file patterns
        name = path.name.lower()
        return "log" in name or "transcript" in name or "debug" in name
