import asyncio
import logging
import uuid
from collections.abc import Callable
from pathlib import Path
from time import time
from typing import TYPE_CHECKING

from anyio import Path as AsyncPath

from exponent.core.remote_execution import files
from exponent.core.remote_execution.cli_rpc_types import (
    BashToolInput,
    BashToolResult,
    DownloadArtifactToolInput,
    DownloadArtifactToolResult,
    EditToolInput,
    EditToolResult,
    ErrorToolResult,
    GlobToolInput,
    GlobToolResult,
    GrepToolInput,
    GrepToolResult,
    ListToolInput,
    ListToolResult,
    ReadToolArtifactResult,
    ReadToolInput,
    ReadToolResult,
    ToolInputType,
    ToolResultType,
    UploadArtifactToolInput,
    UploadArtifactToolResult,
    WriteToolInput,
    WriteToolResult,
)

if TYPE_CHECKING:
    from exponent.core.remote_execution.client import RemoteExecutionClient
import urllib.request

import aiohttp

from exponent.core.remote_execution.cli_rpc_types import (
    StreamingCodeExecutionRequest,
    StreamingCodeExecutionResponse,
)
from exponent.core.remote_execution.code_execution import (
    execute_code_streaming,
)
from exponent.core.remote_execution.file_write import execute_full_file_rewrite
from exponent.core.remote_execution.truncation import truncate_tool_result
from exponent.core.remote_execution.utils import (
    assert_unreachable,
    safe_get_file_metadata,
    safe_read_file,
)

logger = logging.getLogger(__name__)


async def execute_tool(
    tool_input: ToolInputType,
    working_directory: str,
    upload_client: "RemoteExecutionClient | None" = None,
) -> ToolResultType:
    if isinstance(tool_input, ReadToolInput):
        return await execute_read_file(tool_input, working_directory, upload_client)
    elif isinstance(tool_input, WriteToolInput):
        return await execute_write_file(tool_input, working_directory)
    elif isinstance(tool_input, ListToolInput):
        return await execute_list_files(tool_input, working_directory)
    elif isinstance(tool_input, GlobToolInput):
        return await execute_glob_files(tool_input, working_directory)
    elif isinstance(tool_input, GrepToolInput):
        return await execute_grep_files(tool_input, working_directory)
    elif isinstance(tool_input, EditToolInput):
        return await execute_edit_file(tool_input, working_directory)
    elif isinstance(tool_input, DownloadArtifactToolInput):
        return await execute_download_artifact(tool_input, working_directory)
    elif isinstance(tool_input, UploadArtifactToolInput):
        return await execute_upload_artifact(tool_input, working_directory)
    elif isinstance(tool_input, BashToolInput):
        raise ValueError("Bash tool input should be handled by execute_bash_tool")
    else:
        assert_unreachable(tool_input)


def truncate_result[T: ToolResultType](tool_result: T) -> T:
    return truncate_tool_result(tool_result)


def is_image_file(file_path: str) -> tuple[bool, str | None]:
    ext = Path(file_path).suffix.lower()
    if ext == ".png":
        return (True, "image/png")
    elif ext in [".jpg", ".jpeg"]:
        return (True, "image/jpeg")
    return (False, None)


async def execute_read_file(  # noqa: PLR0911, PLR0915
    tool_input: ReadToolInput,
    working_directory: str,
    upload_client: "RemoteExecutionClient | None" = None,
) -> ReadToolResult | ErrorToolResult:
    # Validate absolute path requirement
    if not tool_input.file_path.startswith("/"):
        return ErrorToolResult(
            error_message=f"File path must be absolute, got relative path: {tool_input.file_path}"
        )

    # Validate offset and limit
    offset = tool_input.offset if tool_input.offset is not None else 0
    limit = tool_input.limit if tool_input.limit is not None else 2000

    if limit <= 0:
        return ErrorToolResult(error_message=f"Limit must be positive, got: {limit}")

    file = AsyncPath(working_directory, tool_input.file_path)

    # Check if this is an image file and we have an upload client
    is_image, media_type = is_image_file(tool_input.file_path)
    if is_image and media_type and upload_client is not None:
        try:
            file_name = Path(tool_input.file_path).name
            s3_key = f"images/{uuid.uuid4()}/{file_name}"

            upload_response = await upload_client.request_upload_url(s3_key, media_type)

            f = await file.open("rb")
            async with f:
                file_data = await f.read()

                def _upload() -> int:
                    req = urllib.request.Request(
                        upload_response.upload_url,
                        data=file_data,
                        headers={"Content-Type": media_type},
                        method="PUT",
                    )
                    with urllib.request.urlopen(req) as resp:
                        status: int = resp.status
                        return status

                status = await asyncio.to_thread(_upload)
                if status != 200:
                    raise RuntimeError(f"Upload failed with status {status}")

            return ReadToolResult(
                artifact=ReadToolArtifactResult(
                    s3_uri=upload_response.s3_uri,
                    file_path=tool_input.file_path,
                    media_type=media_type,
                )
            )
        except Exception as e:
            return ErrorToolResult(error_message=f"Failed to upload image to S3: {e!s}")

    try:
        exists = await file.exists()
    except (OSError, PermissionError) as e:
        return ErrorToolResult(error_message=f"Cannot access file: {e!s}")

    if not exists:
        return ErrorToolResult(
            error_message="File not found",
        )

    try:
        if await file.is_dir():
            return ErrorToolResult(
                error_message=f"{await file.absolute()} is a directory",
            )
    except (OSError, PermissionError) as e:
        return ErrorToolResult(error_message=f"Cannot check file type: {e!s}")

    try:
        content = await safe_read_file(file)
    except PermissionError:
        return ErrorToolResult(
            error_message=f"Permission denied: cannot read {tool_input.file_path}"
        )
    except UnicodeDecodeError:
        return ErrorToolResult(
            error_message="File appears to be binary or has invalid text encoding"
        )
    except Exception as e:
        return ErrorToolResult(error_message=f"Error reading file: {e!s}")

    metadata = await safe_get_file_metadata(file)

    # Handle empty files
    if not content:
        return ReadToolResult(
            content="",
            num_lines=0,
            start_line=0,
            total_lines=0,
            metadata=metadata,
        )

    content_lines = content.splitlines(keepends=True)
    total_lines = len(content_lines)

    # Handle offset beyond file length for positive offsets
    if offset >= 0 and offset >= total_lines:
        return ReadToolResult(
            content="",
            num_lines=0,
            start_line=offset,
            total_lines=total_lines,
            metadata=metadata,
        )

    # Use Python's native slicing - it handles negative offsets naturally
    # Handle the case where offset + limit < 0 (can't mix negative and non-negative indices)
    if offset < 0 and offset + limit < 0:
        # Both start and end are negative, use negative end index
        end_index = offset + limit
    elif offset < 0 and offset + limit >= 0:
        # Start is negative but end would be positive/zero, slice to end
        end_index = None
    else:
        # Normal case: both indices are non-negative
        end_index = offset + limit

    content_lines = content_lines[offset:end_index]

    # Calculate the actual start line for the result
    if offset < 0:
        # For negative offsets, calculate where we actually started
        actual_start_line = max(0, total_lines + offset)
    else:
        actual_start_line = offset

    # Apply character-level truncation at line boundaries to ensure consistency
    # This ensures the content field and num_lines field remain in sync
    CHARACTER_LIMIT = 90_000  # Match the limit in truncation.py

    # Join lines and check total size
    final_content = "".join(content_lines)

    if len(final_content) > CHARACTER_LIMIT:
        # Truncate at line boundaries to stay under the limit
        truncated_lines: list[str] = []
        current_size = 0
        truncation_message = "\n[Content truncated due to size limit]"
        truncation_size = len(truncation_message)
        lines_included = 0

        for line in content_lines:
            # Check if adding this line would exceed the limit (accounting for truncation message)
            if current_size + len(line) + truncation_size > CHARACTER_LIMIT:
                final_content = "".join(truncated_lines) + truncation_message
                break
            truncated_lines.append(line)
            current_size += len(line)
            lines_included += 1
        else:
            # All lines fit (shouldn't happen if we got here, but be safe)
            final_content = "".join(truncated_lines)
            lines_included = len(content_lines)

        num_lines = lines_included
    else:
        num_lines = len(content_lines)

    return ReadToolResult(
        content=final_content,
        num_lines=num_lines,
        start_line=actual_start_line,
        total_lines=total_lines,
        metadata=metadata,
    )


async def execute_write_file(
    tool_input: WriteToolInput, working_directory: str
) -> WriteToolResult:
    file_path = tool_input.file_path
    path = Path(working_directory, file_path)
    result = await execute_full_file_rewrite(
        path, tool_input.content, working_directory
    )
    return WriteToolResult(message=result)


async def execute_edit_file(  # noqa: PLR0911
    tool_input: EditToolInput, working_directory: str
) -> EditToolResult | ErrorToolResult:
    # Validate absolute path requirement
    if not tool_input.file_path.startswith("/"):
        return ErrorToolResult(
            error_message=f"File path must be absolute, got relative path: {tool_input.file_path}"
        )

    file = AsyncPath(working_directory, tool_input.file_path)

    try:
        exists = await file.exists()
    except (OSError, PermissionError) as e:
        return ErrorToolResult(error_message=f"Cannot access file: {e!s}")

    if not exists:
        return ErrorToolResult(error_message="File not found")

    if tool_input.last_known_modified_timestamp is not None:
        metadata = await safe_get_file_metadata(file)
        if (
            metadata is not None
            and metadata.modified_timestamp > tool_input.last_known_modified_timestamp
        ):
            return ErrorToolResult(
                error_message="File has been modified since last read/write"
            )

    try:
        if await file.is_dir():
            return ErrorToolResult(
                error_message=f"{await file.absolute()} is a directory"
            )
    except (OSError, PermissionError) as e:
        return ErrorToolResult(error_message=f"Cannot check file type: {e!s}")

    try:
        # Read the entire file without truncation limits
        content = await safe_read_file(file)
    except PermissionError:
        return ErrorToolResult(
            error_message=f"Permission denied: cannot read {tool_input.file_path}"
        )
    except UnicodeDecodeError:
        return ErrorToolResult(
            error_message="File appears to be binary or has invalid text encoding"
        )
    except Exception as e:
        return ErrorToolResult(error_message=f"Error reading file: {e!s}")

    # Check if search text exists
    if tool_input.old_string not in content:
        return ErrorToolResult(
            error_message=f"Search text not found in {tool_input.file_path}"
        )

    # Check if old_string and new_string are identical
    if tool_input.old_string == tool_input.new_string:
        return ErrorToolResult(error_message="Old string and new string are identical")

    # Check uniqueness if replace_all is False
    if not tool_input.replace_all:
        occurrences = content.count(tool_input.old_string)
        if occurrences > 1:
            return ErrorToolResult(
                error_message=f"String '{tool_input.old_string}' appears {occurrences} times in file. Use a larger context or replace_all=True"
            )

    # Perform replacement
    if tool_input.replace_all:
        new_content = content.replace(tool_input.old_string, tool_input.new_string)
    else:
        # Replace only the first occurrence
        new_content = content.replace(tool_input.old_string, tool_input.new_string, 1)

    # Write back to file
    try:
        path = Path(working_directory, tool_input.file_path)
        await execute_full_file_rewrite(path, new_content, working_directory)
        return EditToolResult(
            message=f"Successfully replaced text in {tool_input.file_path}",
            metadata=await safe_get_file_metadata(path),
        )
    except Exception as e:
        return ErrorToolResult(error_message=f"Error writing file: {e!s}")


async def execute_list_files(
    tool_input: ListToolInput, working_directory: str
) -> ListToolResult | ErrorToolResult:
    path = AsyncPath(tool_input.path)

    try:
        exists = await path.exists()
    except (OSError, PermissionError) as e:
        return ErrorToolResult(error_message=f"Cannot access path: {e!s}")

    if not exists:
        return ErrorToolResult(error_message=f"Directory not found: {tool_input.path}")

    try:
        is_dir = await path.is_dir()
    except (OSError, PermissionError) as e:
        return ErrorToolResult(
            error_message=f"Cannot check if path is directory: {e!s}"
        )

    if not is_dir:
        return ErrorToolResult(
            error_message=f"Path is not a directory: {tool_input.path}"
        )

    try:
        filenames = [entry.name async for entry in path.iterdir()]
    except (OSError, PermissionError) as e:
        return ErrorToolResult(error_message=f"Cannot list directory contents: {e!s}")

    return ListToolResult(
        files=[filename for filename in filenames],
    )


async def execute_glob_files(
    tool_input: GlobToolInput, working_directory: str
) -> GlobToolResult:
    # async timer
    start_time = time()
    results = await files.glob(
        path=working_directory if tool_input.path is None else tool_input.path,
        glob_pattern=tool_input.pattern,
    )
    duration_ms = int((time() - start_time) * 1000)
    return GlobToolResult(
        filenames=results,
        duration_ms=duration_ms,
        num_files=len(results),
        truncated=len(results) >= files.GLOB_MAX_COUNT,
    )


async def execute_grep_files(
    tool_input: GrepToolInput, working_directory: str
) -> GrepToolResult | ErrorToolResult:
    return await files.search_files(
        path_str=working_directory if tool_input.path is None else tool_input.path,
        file_pattern=tool_input.include,
        regex=tool_input.pattern,
        working_directory=working_directory,
        multiline=tool_input.multiline,
    )


async def execute_bash_tool(
    tool_input: BashToolInput, working_directory: str, should_halt: Callable[[], bool]
) -> BashToolResult:
    start_time = time()
    result = None
    async for result in execute_code_streaming(
        StreamingCodeExecutionRequest(
            language="shell",
            content=tool_input.command,
            timeout=120 if tool_input.timeout is None else tool_input.timeout,
            correlation_id=str(uuid.uuid4()),
        ),
        working_directory=working_directory,
        session=None,  # type: ignore
        should_halt=should_halt,
    ):
        pass

    assert isinstance(result, StreamingCodeExecutionResponse)

    return BashToolResult(
        shell_output=result.content,
        exit_code=result.exit_code,
        duration_ms=int((time() - start_time) * 1000),
        timed_out=result.cancelled_for_timeout,
        stopped_by_user=result.halted,
    )


async def execute_download_artifact(
    tool_input: DownloadArtifactToolInput, working_directory: str
) -> DownloadArtifactToolResult | ErrorToolResult:
    """Download an artifact from S3 using a pre-signed URL."""

    # Validate absolute path
    if not tool_input.file_path.startswith("/"):
        return ErrorToolResult(
            error_message=f"File path must be absolute, got relative path: {tool_input.file_path}"
        )

    # Check if file exists and overwrite is False
    file_path = Path(tool_input.file_path)
    if file_path.exists() and not tool_input.overwrite:
        return ErrorToolResult(
            error_message=f"File already exists: {tool_input.file_path}. Set overwrite=True to replace it."
        )

    try:
        # Download from pre-signed URL
        async with aiohttp.ClientSession() as session:
            async with session.get(tool_input.presigned_url) as response:
                if response.status != 200:
                    error_text = await response.text()
                    return ErrorToolResult(
                        error_message=f"Failed to download artifact: HTTP {response.status} - {error_text}"
                    )

                # Create parent directory if needed
                file_path.parent.mkdir(parents=True, exist_ok=True)

                # Write file
                content = await response.read()
                file_path.write_bytes(content)

        file_size = len(content)

        # Try to generate content preview for text files
        # Attempt to decode as UTF-8 to determine if it's a text file
        content_preview = None
        num_lines = None
        total_lines = None
        truncated = False

        try:
            text_content = content.decode("utf-8")
            lines = text_content.splitlines()
            total_lines = len(lines)

            # Show first 50 lines
            preview_limit = 50
            if len(lines) > preview_limit:
                preview_lines = lines[:preview_limit]
                truncated = True
                num_lines = preview_limit
            else:
                preview_lines = lines
                num_lines = len(lines)

            content_preview = "\n".join(preview_lines)
        except UnicodeDecodeError:
            # Binary file, skip preview
            pass

        return DownloadArtifactToolResult(
            file_path=tool_input.file_path,
            artifact_id=tool_input.artifact_id,
            file_size_bytes=file_size,
            content_preview=content_preview,
            num_lines=num_lines,
            total_lines=total_lines,
            truncated=truncated,
        )

    except Exception as e:
        logger.exception("Failed to download artifact")
        return ErrorToolResult(error_message=f"Failed to download artifact: {e!s}")


async def execute_upload_artifact(
    tool_input: UploadArtifactToolInput, working_directory: str
) -> UploadArtifactToolResult | ErrorToolResult:
    """Upload an artifact to S3 using a pre-signed URL."""

    # Validate absolute path
    if not tool_input.file_path.startswith("/"):
        return ErrorToolResult(
            error_message=f"File path must be absolute, got relative path: {tool_input.file_path}"
        )

    # Check if file exists
    file_path = Path(tool_input.file_path)
    if not file_path.exists():
        return ErrorToolResult(error_message=f"File not found: {tool_input.file_path}")

    if not file_path.is_file():
        return ErrorToolResult(
            error_message=f"Path is not a file: {tool_input.file_path}"
        )

    try:
        # Read file
        content = file_path.read_bytes()
        file_size = len(content)

        # Upload to pre-signed URL
        async with aiohttp.ClientSession() as session:
            headers = {"Content-Type": tool_input.content_type}
            async with session.put(
                tool_input.presigned_url, data=content, headers=headers
            ) as response:
                if response.status not in (200, 204):
                    error_text = await response.text()
                    return ErrorToolResult(
                        error_message=f"Failed to upload artifact: HTTP {response.status} - {error_text}"
                    )

        return UploadArtifactToolResult(
            artifact_id=tool_input.artifact_id,
            file_size_bytes=file_size,
            content_type=tool_input.content_type,
        )

    except Exception as e:
        logger.exception("Failed to upload artifact")
        return ErrorToolResult(error_message=f"Failed to upload artifact: {e!s}")
