import os
from typing import Optional, List, Tuple, Union
from autocoder.common import AutoCoderArgs, SourceCode
from autocoder.common.autocoderargs_parser import AutoCoderArgsParser
from autocoder.common.v2.agent.agentic_edit_tools.base_tool_resolver import (
    BaseToolResolver,
)
from autocoder.common.v2.agent.agentic_edit_types import ReadFileTool, ToolResult
from autocoder.common.pruner.context_pruner import PruneContext
from autocoder.common.tokens import count_string_tokens as count_tokens
from autocoder.common.wrap_llm_hint.utils import add_hint_to_text
from loguru import logger
import typing
from autocoder.rag.loaders import (
    extract_text_from_pdf,
    extract_text_from_docx,
    extract_text_from_ppt,
    extract_text_from_excel,
)

if typing.TYPE_CHECKING:
    from autocoder.common.v2.agent.agentic_edit import AgenticEdit


# Token budget reserved for hint message (approx 150 tokens for hint text + formatting)
HINT_TOKEN_RESERVE = 200


class ReadFileToolResolver(BaseToolResolver):
    """文件读取工具解析器，支持多文件/行号范围/query抽取，统一按 token 限制截断"""

    def __init__(
        self, agent: Optional["AgenticEdit"], tool: ReadFileTool, args: AutoCoderArgs
    ):
        super().__init__(agent, tool, args)
        self.tool: ReadFileTool = tool

        self.args_parser = AutoCoderArgsParser()
        self.safe_zone_tokens = self._get_parsed_safe_zone_tokens()

        # 初始化 context_pruner 用于 query 抽取
        self.context_pruner = (
            PruneContext(
                max_tokens=self.safe_zone_tokens,
                args=self.args,
                llm=self.agent.context_prune_llm,
            )
            if self.agent and self.agent.context_prune_llm
            else None
        )

    def _get_parsed_safe_zone_tokens(self) -> int:
        """解析 safe_zone_tokens 参数"""
        return self.args_parser.parse_context_prune_safe_zone_tokens(
            self.args.context_prune_safe_zone_tokens, self.args.code_model
        )

    # -------------------------------------------------------------------------
    # Token 截断相关方法（可单独测试）
    # -------------------------------------------------------------------------

    def _truncate_to_token_budget(self, text: str, max_tokens: int) -> str:
        """
        按 token 预算截断文本。

        使用 count_string_tokens 配合二分查找实现 token 级别截断。

        Args:
            text: 要截断的文本
            max_tokens: 最大 token 数

        Returns:
            截断后的文本
        """
        if max_tokens <= 0:
            return ""

        try:
            current_tokens = count_tokens(text)
            if current_tokens <= max_tokens:
                return text

            # 使用二分查找找到合适的字符截断点
            # 估算初始范围：假设平均 4 字符/token
            low = 0
            high = len(text)

            # 快速缩小范围：先用估算值
            estimated_chars = max_tokens * 4
            if estimated_chars < high:
                high = min(high, estimated_chars + 1000)  # 加一些余量

            result = ""
            while low < high:
                mid = (low + high + 1) // 2
                candidate = text[:mid]
                candidate_tokens = count_tokens(candidate)

                if candidate_tokens <= max_tokens:
                    result = candidate
                    low = mid
                else:
                    high = mid - 1

            return result if result else text[: max_tokens * 4]
        except Exception as e:
            logger.warning(
                f"Token truncation failed, falling back to char truncation: {e}"
            )
            # 安全回退：按字符截断，假设平均 4 字符/token
            char_limit = max_tokens * 4
            return text[:char_limit] if len(text) > char_limit else text

    def _enforce_token_limit(
        self,
        text: str,
        file_path: str,
        limit_tokens: int,
    ) -> Tuple[str, bool]:
        """
        强制执行 token 限制，超限时截断并添加提示。

        Args:
            text: 输入文本
            file_path: 文件路径（用于提示信息）
            limit_tokens: token 上限

        Returns:
            Tuple[处理后的文本, 是否被截断]
        """
        current_tokens = count_tokens(text)

        if current_tokens <= limit_tokens:
            return text, False

        # 超限，需要截断
        total_lines = len(text.split("\n"))

        # 预留 hint 的 token 空间
        content_budget = max(100, limit_tokens - HINT_TOKEN_RESERVE)

        # 按 token 预算截断
        truncated_content = self._truncate_to_token_budget(text, content_budget)

        # 生成提示信息
        hint = (
            f"Cannot read entire file ({current_tokens} tokens, {total_lines} lines). "
            f"Content truncated to approximately {content_budget} tokens. "
            f"Use start_line/end_line for specific range, "
            f"or query to extract relevant content."
        )

        return add_hint_to_text(truncated_content, hint), True

    # -------------------------------------------------------------------------
    # 内容提取方法
    # -------------------------------------------------------------------------

    def _extract_lines_by_range(
        self, content: str, start_line: Optional[int], end_line: Optional[int]
    ) -> str:
        """根据行号范围提取内容（1-based，支持负数索引）"""
        if start_line is None and end_line is None:
            return content

        lines = content.split("\n")
        total_lines = len(lines)

        # 处理起始行
        if start_line is None:
            start_idx = 0
        elif start_line < 0:
            start_idx = max(0, total_lines + start_line)
        else:
            start_idx = max(0, start_line - 1)

        # 处理结束行
        if end_line is None or end_line == -1:
            end_idx = total_lines
        elif end_line < -1:
            end_idx = max(0, total_lines + end_line + 1)
        else:
            end_idx = min(total_lines, end_line)

        # 验证范围
        if start_idx >= total_lines:
            return f"Error: start_line {start_line} exceeds total lines {total_lines}"
        if start_idx >= end_idx:
            return f"Error: invalid line range (start={start_line}, end={end_line})"

        return "\n".join(lines[start_idx:end_idx])

    def _extract_content_by_query(
        self, content: str, query: str, file_path: str
    ) -> str:
        """
        根据 query 从内容中抽取相关部分。

        注意：即使抽取成功，返回的内容仍可能超过 token 限制，
        调用方需要再次调用 _enforce_token_limit() 进行最终限制。
        """
        if not query or not self.context_pruner:
            return content

        try:
            source_code = SourceCode(
                module_name=file_path, source_code=content, tokens=count_tokens(content)
            )
            pruned_sources = self.context_pruner.handle_overflow(
                file_sources=[source_code],
                conversations=[{"role": "user", "content": query}],
                strategy=self.args.context_prune_strategy,
            )
            return pruned_sources[0].source_code if pruned_sources else content
        except Exception as e:
            logger.warning(f"Query extraction failed for '{query}': {e}")
            return content

    def _read_raw_content(self, file_path: str) -> str:
        """读取原始文件内容"""
        ext = os.path.splitext(file_path)[1].lower()

        if ext == ".pdf":
            return extract_text_from_pdf(file_path)
        elif ext == ".docx":
            return extract_text_from_docx(file_path)
        elif ext in (".pptx", ".ppt"):
            slides = extract_text_from_ppt(file_path)
            return "\n\n".join(f"--- Slide {idx} ---\n{text}" for idx, text in slides)
        elif ext in (".xlsx", ".xls"):
            sheets = extract_text_from_excel(file_path)
            return "\n\n".join(
                f"--- Sheet: {sheet_name.split('#')[-1]} ---\n{content}"
                for sheet_name, content in sheets
            )
        else:
            with open(file_path, "r", encoding="utf-8", errors="replace") as f:
                return f.read()

    def _read_file_content(
        self,
        file_path: str,
        start_line: Optional[int],
        end_line: Optional[int],
        query: Optional[str],
    ) -> Tuple[str, bool]:
        """
        读取文件内容，统一流程：读取 → 可选抽取 → 统一限长

        流程：
        1. 读取原始文件内容
        2. 如果指定了行号范围，按范围提取
        3. 否则如果有 query 且 pruner 可用，做 query 抽取
        4. 最后统一执行 token 限制（无论上述哪个分支）

        Returns:
            Tuple[content, is_truncated]: 内容和是否被截断的标记
        """
        # Step 1: 读取原始内容
        content = self._read_raw_content(file_path)

        # Step 2: 可选的内容抽取
        if start_line is not None or end_line is not None:
            # 按行号范围提取
            content = self._extract_lines_by_range(content, start_line, end_line)
        elif query:
            # 按 query 抽取（即使 pruner 不可用也会返回原文）
            content = self._extract_content_by_query(content, query, file_path)

        # Step 3: 统一执行 token 限制（这是关键修复点）
        # 无论是全文、行号范围还是 query 抽取，最终都要遵守 token 上限
        return self._enforce_token_limit(content, file_path, self.safe_zone_tokens)

    # -------------------------------------------------------------------------
    # 参数解析方法
    # -------------------------------------------------------------------------

    def _parse_single_int(
        self, value: Optional[Union[str, int]], field_name: str
    ) -> Optional[int]:
        """解析单个整数参数"""
        if value is None:
            return None
        if isinstance(value, int):
            return value
        if isinstance(value, str):
            v = value.strip()
            if not v:
                return None
            if "," in v:
                raise ValueError(
                    f"Parameter format error: '{field_name}' contains comma but only one file path is specified. "
                    f"For single file: use a single value like {field_name}='10'. "
                    f"For multiple files: specify multiple paths like path='a.py, b.py' with {field_name}='10, 20'."
                )
            return int(v)
        raise ValueError(f"Invalid type for '{field_name}': {type(value)}")

    def _split_values(self, raw: Optional[Union[str, int]]) -> List[Optional[str]]:
        """将参数拆分为列表"""
        if raw is None:
            return []
        if isinstance(raw, int):
            return [str(raw)]
        if isinstance(raw, str):
            parts = [p.strip() for p in raw.split(",")]
            return [p if p else None for p in parts]
        return []

    def _parse_int_list(
        self, items: List[Optional[str]], field_name: str
    ) -> Tuple[bool, Optional[List[Optional[int]]], Optional[str]]:
        """解析整数列表"""
        if not items:
            return True, None, None
        parsed = []
        for raw in items:
            if raw is None or raw == "":
                parsed.append(None)
            else:
                try:
                    parsed.append(int(raw))
                except ValueError:
                    return False, None, f"'{field_name}' must be integers. Got: {raw}"
        return True, parsed, None

    # -------------------------------------------------------------------------
    # 主入口方法
    # -------------------------------------------------------------------------

    def read_single_file(self, file_path: str) -> ToolResult:
        """读取单个文件"""
        if not os.path.exists(file_path):
            return ToolResult(success=False, message=f"File not found: {file_path}")
        if not os.path.isfile(file_path):
            return ToolResult(success=False, message=f"Not a file: {file_path}")

        try:
            start_line = self._parse_single_int(self.tool.start_line, "start_line")
            end_line = self._parse_single_int(self.tool.end_line, "end_line")
        except ValueError as e:
            return ToolResult(success=False, message=str(e))

        try:
            content, is_truncated = self._read_file_content(
                file_path, start_line, end_line, self.tool.query
            )

            # 构建消息
            msg_parts = [file_path]
            if start_line or end_line:
                msg_parts.append(f"lines {start_line or 1}-{end_line or 'end'}")
            if self.tool.query:
                msg_parts.append(f"query: '{self.tool.query}'")
            if is_truncated:
                msg_parts.append("(truncated)")

            return ToolResult(
                success=True, message=" ".join(msg_parts), content=content
            )
        except Exception as e:
            logger.exception(f"Error reading file '{file_path}'")
            return ToolResult(
                success=False, message=f"Error reading '{file_path}': {e}"
            )

    def resolve(self) -> ToolResult:
        """解析读取文件请求，支持多文件"""
        raw_path = self.tool.path or ""
        paths = [p.strip() for p in raw_path.split(",") if p.strip()]

        if not paths:
            return ToolResult(success=False, message="'path' is required")

        # 单文件模式
        if len(paths) == 1:
            # 检查单文件模式下是否误用了逗号分隔的参数
            for field, value in [
                ("start_line", self.tool.start_line),
                ("end_line", self.tool.end_line),
            ]:
                if isinstance(value, str) and value and "," in value:
                    return ToolResult(
                        success=False,
                        message=(
                            f"Parameter format error: '{field}' has comma-separated values but only one file path is specified. "
                            f"For single file: use a single value like {field}='10'. "
                            f"For multiple files: specify multiple paths like path='a.py, b.py' with {field}='10, 20'."
                        ),
                    )
            return self.read_single_file(paths[0])

        # 多文件模式
        start_items = self._split_values(self.tool.start_line)
        end_items = self._split_values(self.tool.end_line)
        query_items = self._split_values(self.tool.query)

        # 验证参数数量
        for name, items in [
            ("start_line", start_items),
            ("end_line", end_items),
            ("query", query_items),
        ]:
            if items and len(items) != len(paths):
                return ToolResult(
                    success=False,
                    message=(
                        f"Parameter mismatch: '{name}' has {len(items)} value(s) but you specified {len(paths)} file path(s). "
                        f"When reading multiple files with '{name}', you must provide comma-separated values matching the number of paths. "
                        f"Example: path='a.py, b.py' with {name}='value1, value2'. "
                        f"Alternatively, omit '{name}' to read all files without this filter."
                    ),
                )

        # 解析整数列表
        ok, start_lines, err = self._parse_int_list(start_items, "start_line")
        if not ok:
            return ToolResult(success=False, message=err or "Invalid start_line")
        ok, end_lines, err = self._parse_int_list(end_items, "end_line")
        if not ok:
            return ToolResult(success=False, message=err or "Invalid end_line")

        # 读取每个文件
        results = []
        failures = []
        any_truncated = False
        for i, path in enumerate(paths):
            if not os.path.exists(path):
                failures.append(f"[{i+1}] not found: {path}")
                continue
            if not os.path.isfile(path):
                failures.append(f"[{i+1}] not a file: {path}")
                continue

            try:
                s = start_lines[i] if start_lines else None
                e = end_lines[i] if end_lines else None
                q = query_items[i] if query_items else None
                content, is_truncated = self._read_file_content(path, s, e, q)
                if is_truncated:
                    any_truncated = True
                results.append(f"---- File: {path} ----\n{content}")
            except Exception as ex:
                logger.exception(ex)
                failures.append(f"[{i+1}] error: {path}: {ex}")

        if not results:
            return ToolResult(
                success=False,
                message="Failed to read all files. " + "; ".join(failures),
            )

        summary = f"{len(results)}/{len(paths)} files read"
        if any_truncated:
            summary += " (some files truncated)"
        if failures:
            summary += f" | failures: {'; '.join(failures)}"

        return ToolResult(success=True, message=summary, content="\n\n".join(results))
