"""Chunking and ingestion helpers for normal review mode."""

from __future__ import annotations

import logging
import os
from dataclasses import dataclass
from typing import List, Optional

from normal.ai.ai_utils import count_tokens

logger = logging.getLogger(__name__)


@dataclass
class Chunk:
    """Represents a code chunk from a file."""
    file_path: str
    start_line: int
    end_line: int
    chunk_id: str
    overlap: int
    token_count: int = 0


class Chunker:
    """Split files into manageable chunks based on token count."""

    def __init__(self, chunk_tokens: int = 30000, overlap_tokens: int = 500):
        """
        Initialize chunker with token-based limits.

        Args:
            chunk_tokens: Maximum tokens per chunk (default 30000, leaves room for context)
            overlap_tokens: Token overlap between chunks (default 500)
        """
        self.chunk_tokens = chunk_tokens
        self.overlap_tokens = overlap_tokens

    def chunk_file(self, content: str, file_path: str) -> List[Chunk]:
        """Split file content into token-based chunks."""
        lines = content.splitlines()
        chunks: List[Chunk] = []
        n = len(lines)

        if n == 0:
            return chunks

        start_line = 0
        while start_line < n:
            # Start with initial estimate: take lines until we hit token limit
            end_line = start_line
            current_content = []
            current_tokens = 0

            # Build chunk line by line until token limit
            while end_line < n:
                line = lines[end_line]
                test_content = "\n".join(current_content + [line])
                test_tokens = count_tokens(test_content)

                if test_tokens > self.chunk_tokens and current_content:
                    # We've exceeded limit, stop here
                    break

                current_content.append(line)
                current_tokens = test_tokens
                end_line += 1

            # Ensure we made progress (at least 1 line)
            if end_line == start_line:
                end_line = start_line + 1
                current_content = [lines[start_line]]
                current_tokens = count_tokens(current_content[0])

            chunk_id = f"{file_path}:{start_line + 1}-{end_line}"
            chunks.append(Chunk(
                file_path=file_path,
                start_line=start_line + 1,
                end_line=end_line,
                chunk_id=chunk_id,
                overlap=self.overlap_tokens,
                token_count=current_tokens,
            ))

            if end_line >= n:
                break

            # Calculate overlap in lines (go back by overlap_tokens worth of lines)
            overlap_start = end_line
            overlap_content = []
            overlap_tokens_counted = 0

            # Walk backwards to find overlap boundary
            while overlap_start > start_line and overlap_tokens_counted < self.overlap_tokens:
                overlap_start -= 1
                overlap_line = lines[overlap_start]
                overlap_content.insert(0, overlap_line)
                overlap_tokens_counted = count_tokens("\n".join(overlap_content))

            # Next chunk starts at overlap boundary (or end_line if overlap is small)
            start_line = max(start_line + 1, overlap_start)

        # Log summary with token info
        total_tokens = sum(c.token_count for c in chunks)
        if len(chunks) == 1:
            logger.info(
                f"Split {file_path} into 1 chunk (max: {self.chunk_tokens} tokens, "
                f"actual: {chunks[0].token_count} tokens)"
            )
        else:
            chunk_tokens_str = ", ".join(f"{c.token_count}" for c in chunks)
            logger.info(
                f"Split {file_path} into {len(chunks)} chunks (max: {self.chunk_tokens} "
                f"tokens/chunk, chunks: [{chunk_tokens_str}] tokens, overlap: "
                f"{self.overlap_tokens} tokens)"
            )

        for i, chunk in enumerate(chunks):
            logger.debug(
                f"  Chunk {i+1}: lines {chunk.start_line}-{chunk.end_line}, "
                f"{chunk.token_count} tokens"
            )

        return chunks


class ChunkIngestor:
    """Ingest files into chunk store and planner."""

    def __init__(self, offload, planner, chunker: Chunker, progress_handle=None):
        self.offload = offload
        self.planner = planner
        self.chunker = chunker
        self.progress_handle = progress_handle

    def ingest_files(self, file_paths: List[str]) -> int:
        """Ingest files and create chunks. Returns number of chunks."""
        total_chunks = 0
        from common.ui import update_file_read_progress

        # progress handle이 있으면 total 설정
        if self.progress_handle:
            update_file_read_progress(self.progress_handle, total=len(file_paths))

        for path in file_paths:
            try:
                content = self._read_file_content(path)
                if content is None:
                    continue

                # 파일 읽기 진행률 업데이트
                if self.progress_handle:
                    update_file_read_progress(self.progress_handle)

                total_chunks += self._ingest_content_chunks(path, content)
            except UnicodeDecodeError as e:
                logger.error(f"Failed to decode file {path}: {e}")
            except IOError as e:
                logger.error(f"IO error reading {path}: {e}")
            except Exception as e:
                logger.error(f"Unexpected error ingesting {path}: {e}", exc_info=True)

        return total_chunks

    def ingest_file_entries(self, file_entries: List[dict]) -> int:
        """Ingest in-memory file entries and create chunks."""
        total_chunks = 0
        for entry in file_entries:
            if not isinstance(entry, dict):
                continue
            if entry.get("is_binary"):
                continue
            path = entry.get("path") or "unknown"
            content = entry.get("content", "")
            if not isinstance(content, str) or not content.strip():
                continue
            try:
                total_chunks += self._ingest_content_chunks(path, content)
            except Exception as e:
                logger.error(f"Unexpected error ingesting entry {path}: {e}", exc_info=True)
        return total_chunks

    def ingest_file_entries_batched(self, file_entries: List[dict], chunk_tokens: int) -> int:
        """Ingest multiple file entries into batched chunks within token limit."""
        total_chunks = 0
        batch_index = 0
        current_files: List[str] = []
        current_parts: List[str] = []
        current_tokens = 0

        def flush_batch() -> None:
            nonlocal total_chunks, batch_index, current_files, current_parts, current_tokens
            if not current_parts:
                return
            combined = "\n\n".join(current_parts)
            line_count = len(combined.splitlines())
            chunk_id = f"batch:{batch_index + 1}"
            chunk = Chunk(
                file_path="batch",
                start_line=1,
                end_line=line_count,
                chunk_id=chunk_id,
                overlap=0,
                token_count=current_tokens,
            )
            self.offload.put(chunk.chunk_id, combined, {"file_paths": list(current_files)})
            self.planner.add_chunks([chunk])
            total_chunks += 1
            logger.info(
                f"Ingested {chunk_id}: {len(current_files)} files ({current_tokens} tokens)"
            )
            batch_index += 1
            current_files = []
            current_parts = []
            current_tokens = 0

        for entry in file_entries:
            if not isinstance(entry, dict):
                continue
            if entry.get("is_binary"):
                continue
            path = entry.get("path") or "unknown"
            content = entry.get("content", "")
            if not isinstance(content, str) or not content.strip():
                continue

            try:
                formatted = self._format_file_with_line_numbers(path, content)
                file_tokens = count_tokens(formatted)

                if file_tokens > chunk_tokens:
                    logger.info(
                        f"File exceeds batch token limit ({file_tokens} tokens): {path}"
                    )
                    flush_batch()
                    total_chunks += self._ingest_content_chunks(path, content)
                    continue

                if current_tokens + file_tokens > chunk_tokens and current_parts:
                    flush_batch()

                current_files.append(path)
                current_parts.append(formatted)
                current_tokens += file_tokens
            except Exception as e:
                logger.error(f"Unexpected error batching entry {path}: {e}", exc_info=True)

        flush_batch()
        return total_chunks

    def ingest_files_batched(self, file_paths: List[str], chunk_tokens: int) -> int:
        """Ingest multiple files into batched chunks within token limit."""
        total_chunks = 0
        batch_index = 0
        current_files: List[str] = []
        current_parts: List[str] = []
        current_tokens = 0

        def flush_batch() -> None:
            nonlocal total_chunks, batch_index, current_files, current_parts, current_tokens
            if not current_parts:
                return
            combined = "\n\n".join(current_parts)
            line_count = len(combined.splitlines())
            chunk_id = f"batch:{batch_index + 1}"
            chunk = Chunk(
                file_path="batch",
                start_line=1,
                end_line=line_count,
                chunk_id=chunk_id,
                overlap=0,
                token_count=current_tokens,
            )
            self.offload.put(chunk.chunk_id, combined, {"file_paths": list(current_files)})
            self.planner.add_chunks([chunk])
            total_chunks += 1
            logger.info(
                f"Ingested {chunk_id}: {len(current_files)} files ({current_tokens} tokens)"
            )
            batch_index += 1
            current_files = []
            current_parts = []
            current_tokens = 0

        for path in file_paths:
            try:
                content = self._read_file_content(path)
                if content is None:
                    continue

                formatted = self._format_file_with_line_numbers(path, content)
                file_tokens = count_tokens(formatted)

                if file_tokens > chunk_tokens:
                    logger.info(
                        f"File exceeds batch token limit ({file_tokens} tokens): {path}"
                    )
                    flush_batch()
                    total_chunks += self._ingest_content_chunks(path, content)
                    continue

                if current_tokens + file_tokens > chunk_tokens and current_parts:
                    flush_batch()

                current_files.append(path)
                current_parts.append(formatted)
                current_tokens += file_tokens
            except Exception as e:
                logger.error(f"Unexpected error batching {path}: {e}", exc_info=True)

        flush_batch()
        return total_chunks

    def _read_file_content(self, path: str) -> Optional[str]:
        """Read file content with validation."""
        if not os.path.exists(path):
            logger.warning(f"Path does not exist: {path}")
            return None

        if not os.path.isfile(path):
            logger.warning(f"Path is not a file: {path}")
            return None

        if not os.access(path, os.R_OK):
            logger.error(f"No read permission for file: {path}")
            return None

        try:
            file_size = os.path.getsize(path)
            if file_size > 10 * 1024 * 1024:  # 10MB limit
                logger.warning(f"File too large ({file_size} bytes), skipping: {path}")
                return None
            if file_size == 0:
                logger.debug(f"Empty file, skipping: {path}")
                return None
        except OSError as e:
            logger.error(f"Failed to get file size for {path}: {e}")
            return None

        with open(path, "r", encoding="utf-8", errors="replace") as f:
            content = f.read()

        if not content.strip():
            logger.debug(f"File contains no content: {path}")
            return None

        return content

    def _ingest_content_chunks(self, path: str, content: str) -> int:
        """Chunk content for a single file and store it."""
        chunks = self.chunker.chunk_file(content, file_path=path)

        if not chunks:
            logger.debug(f"No chunks created for file: {path}")
            return 0

        lines = content.splitlines()
        for chunk in chunks:
            try:
                if chunk.start_line < 1 or chunk.end_line > len(lines):
                    logger.warning(
                        f"Invalid chunk range [{chunk.start_line}:{chunk.end_line}] for "
                        f"{path} (total lines: {len(lines)})"
                    )
                    continue

                chunk_content = "\n".join(lines[chunk.start_line - 1:chunk.end_line])
                meta = {
                    "file_path": chunk.file_path,
                    "start_line": chunk.start_line,
                    "end_line": chunk.end_line,
                }
                self.offload.put(chunk.chunk_id, chunk_content, meta)
            except Exception as e:
                logger.error(f"Failed to process chunk {chunk.chunk_id}: {e}")
                continue

        self.planner.add_chunks(chunks)

        if len(chunks) == 1:
            logger.info(f"Ingested {path}: 1 chunk ({chunks[0].token_count} tokens)")
        else:
            chunk_tokens_str = ", ".join(f"{c.token_count}" for c in chunks)
            logger.info(f"Ingested {path}: {len(chunks)} chunks ([{chunk_tokens_str}] tokens)")

        return len(chunks)

    def _format_file_with_line_numbers(self, path: str, content: str) -> str:
        """Format file content with line numbers for batched analysis."""
        lines = content.splitlines()
        numbered = "\n".join(f"{i + 1}: {line}" for i, line in enumerate(lines))
        return f"=== FILE: {path} ===\n{numbered}"
