#!/usr/bin/env python3
"""Context engineering for large code reviews.

Components:
1) Chunker: Split files into chunks with overlap
2) OffloadStore: Store chunk content externally
3) MemoryStore: Keep only summaries (FileCard/Finding/RepoSummary)
4) Planner: Prioritize which chunk to analyze next
5) ContextBuilder: Build LLM context within budget
6) IterativeLoopRunner: Orchestrate the review loop
"""

from __future__ import annotations

import json
import logging
import os
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

from normal.ai.ai_utils import count_tokens
from normal.review.chunking import Chunk, Chunker, ChunkIngestor
from normal.review.prompts import SYSTEM_PROMPT, SYSTEM_PROMPT_MULTI, build_user_prompt

logger = logging.getLogger(__name__)


# -----------------------------
# Data Models
# -----------------------------


@dataclass
class FileCard:
    """Summary of a file's role and characteristics."""
    file_path: str
    role_summary: str
    key_symbols: List[str] = field(default_factory=list)
    deps_hints: List[str] = field(default_factory=list)
    risk_flags: List[str] = field(default_factory=list)


@dataclass
class Finding:
    """A code review finding."""
    category: str
    subtype: str
    severity: str
    file_path: str
    line_range: Tuple[int, int]
    evidence: str
    message: str
    min_fix: str


@dataclass
class RepoSummary:
    """High-level repository summary."""
    architecture_overview: str = ""
    hotspots: List[str] = field(default_factory=list)
    recurring_patterns: List[str] = field(default_factory=list)


@dataclass
class ReviewResult:
    """Result of a single chunk review."""
    chunk_id: str
    success: bool
    raw_response: str
    parsed_data: Optional[Dict] = None
    error_message: Optional[str] = None


# -----------------------------
# Chunker (imported from chunking.py)
# -----------------------------


# -----------------------------
# Offload Store
# -----------------------------


class OffloadStore:
    """Abstract store for chunk content."""
    
    def put(self, chunk_id: str, content: str, meta: Dict) -> None:
        raise NotImplementedError
    
    def get(self, chunk_id: str) -> Optional[Tuple[str, Dict]]:
        raise NotImplementedError


class InMemoryOffloadStore(OffloadStore):
    """In-memory implementation of OffloadStore."""
    
    def __init__(self):
        self._store: Dict[str, Tuple[str, Dict]] = {}
    
    def put(self, chunk_id: str, content: str, meta: Dict) -> None:
        self._store[chunk_id] = (content, meta)
    
    def get(self, chunk_id: str) -> Optional[Tuple[str, Dict]]:
        return self._store.get(chunk_id)


# -----------------------------
# Memory Store
# -----------------------------


class MemoryStore:
    """Store for review summaries with bounded size."""
    
    def __init__(self, max_filecards: int = 200, max_findings: int = 500):
        self.max_filecards = max_filecards
        self.max_findings = max_findings
        self.filecards: Dict[str, FileCard] = {}
        self.findings: List[Finding] = []
        self.repo_summary = RepoSummary()

    def upsert_filecard(self, card: FileCard) -> None:
        """Add or update a file card."""
        self.filecards[card.file_path] = card
        
        # Enforce size limit
        while len(self.filecards) > self.max_filecards:
            oldest = next(iter(self.filecards))
            if oldest != card.file_path:
                del self.filecards[oldest]
            else:
                break

    def add_findings(self, findings: List[Finding]) -> None:
        """Add findings with deduplication."""
        self.findings.extend(findings)
        
        # Deduplicate by (file_path, line_range, message)
        seen: set = set()
        deduped: List[Finding] = []
        for f in reversed(self.findings):
            key = (f.file_path, f.line_range, f.message)
            if key not in seen:
                seen.add(key)
                deduped.append(f)
        
        deduped.reverse()
        self.findings = deduped[:self.max_findings]

    def update_repo_summary(
        self,
        architecture_overview: Optional[str] = None,
        hotspots: Optional[List[str]] = None,
        recurring_patterns: Optional[List[str]] = None,
    ) -> None:
        """Update repository summary."""
        if architecture_overview is not None:
            self.repo_summary.architecture_overview = architecture_overview
        if hotspots is not None:
            self.repo_summary.hotspots = hotspots
        if recurring_patterns is not None:
            self.repo_summary.recurring_patterns = recurring_patterns


# -----------------------------
# Planner
# -----------------------------


class Planner:
    """Plan which chunks to analyze next."""
    
    ENTRYPOINT_HINTS = ["main", "cli", "app", "server", "run"]
    CORE_HINTS = ["service", "api", "db", "auth", "storage", "core"]

    def __init__(self):
        self.queue: List[Chunk] = []
        self.visited: set[str] = set()

    def add_chunks(self, chunks: List[Chunk]) -> None:
        """Add chunks to the queue."""
        self.queue.extend(chunks)

    def select_next(self) -> Optional[Chunk]:
        """Select the next chunk to analyze."""
        candidates = [c for c in self.queue if c.chunk_id not in self.visited]
        if not candidates:
            return None
        
        candidates.sort(key=self._score, reverse=True)
        chosen = candidates[0]
        self.visited.add(chosen.chunk_id)
        self.queue = [c for c in self.queue if c.chunk_id != chosen.chunk_id]
        return chosen

    def _score(self, chunk: Chunk) -> int:
        """Score a chunk for priority."""
        path = chunk.file_path.lower()
        score = 0
        
        if any(h in path for h in self.ENTRYPOINT_HINTS):
            score += 20
        if any(h in path for h in self.CORE_HINTS):
            score += 10
        
        # Prefer smaller chunks
        chunk_size = chunk.end_line - chunk.start_line
        score -= chunk_size // 100
        
        return score


# -----------------------------
# Context Builder
# -----------------------------


class ContextBuilder:
    """Build context for LLM calls."""
    
    def __init__(self, budget_chars: int = 6000, top_k_cards: int = 3):
        self.budget_chars = budget_chars
        self.top_k_cards = top_k_cards

    def build(
        self,
        repo_summary: RepoSummary,
        filecards: List[FileCard],
        chunk_content: str,
    ) -> Dict:
        """Build context dictionary."""
        selected_cards = filecards[:self.top_k_cards]
        
        context = {
            "repo_summary": {
                "architecture_overview": repo_summary.architecture_overview[:500],
                "hotspots": repo_summary.hotspots[:5],
                "recurring_patterns": repo_summary.recurring_patterns[:3],
            },
            "file_cards": [
                {
                    "file_path": fc.file_path,
                    "role_summary": fc.role_summary[:200],
                    "key_symbols": fc.key_symbols[:5],
                }
                for fc in selected_cards
            ],
            "current_chunk": chunk_content,
        }
        
        return context


# -----------------------------
# Iterative Loop Runner
# -----------------------------


class IterativeLoopRunner:
    """Orchestrate the iterative review process."""

    def __init__(
        self,
        offload: Optional[OffloadStore] = None,
        memory: Optional[MemoryStore] = None,
        chunker: Optional[Chunker] = None,
        planner: Optional[Planner] = None,
        context_builder: Optional[ContextBuilder] = None,
        progress_handle=None,
    ):
        self.offload = offload or InMemoryOffloadStore()
        self.memory = memory or MemoryStore()
        self.chunker = chunker or Chunker()
        self.planner = planner or Planner()
        self.context_builder = context_builder or ContextBuilder()
        self.ingestor = ChunkIngestor(
            offload=self.offload,
            planner=self.planner,
            chunker=self.chunker,
            progress_handle=progress_handle,
        )
        self.results: List[ReviewResult] = []

    def ingest_files(self, file_paths: List[str]) -> int:
        """Ingest files and create chunks. Returns number of chunks."""
        return self.ingestor.ingest_files(file_paths)
    
    def ingest_files_batched(self, file_paths: List[str], chunk_tokens: int) -> int:
        """Ingest multiple files into batched chunks within token limit."""
        return self.ingestor.ingest_files_batched(file_paths, chunk_tokens)

    def ingest_file_entries(self, file_paths: List[str]) -> int:
        """[DEPRECATED] Old ingest_files implementation - kept for reference."""
        total_chunks = 0
        
        for path in file_paths:
            # Check file existence
            if not os.path.exists(path):
                logger.warning(f"Path does not exist: {path}")
                continue
            
            if not os.path.isfile(path):
                logger.warning(f"Path is not a file: {path}")
                continue
            
            # Check file permissions
            if not os.access(path, os.R_OK):
                logger.error(f"No read permission for file: {path}")
                continue
            
            # Check file size (skip very large files)
            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}")
                    continue
                if file_size == 0:
                    logger.debug(f"Empty file, skipping: {path}")
                    continue
            except OSError as e:
                logger.error(f"Failed to get file size for {path}: {e}")
                continue
            
            try:
                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}")
                    continue
                
                chunks = self.chunker.chunk_file(content, file_path=path)
                
                if not chunks:
                    logger.debug(f"No chunks created for file: {path}")
                    continue
                
                for chunk in chunks:
                    try:
                        lines = content.splitlines()
                        # Validate line ranges
                        if chunk.start_line < 1 or chunk.end_line > len(lines):
                            logger.warning(f"Invalid chunk range [{chunk.start_line}:{chunk.end_line}] for {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)
                total_chunks += len(chunks)
                
                # Log ingestion with token info
                if len(chunks) == 1:
                    logger.debug(f"Ingested {path}: 1 chunk ({chunks[0].token_count} tokens)")
                else:
                    chunk_tokens_str = ", ".join(f"{c.token_count}" for c in chunks)
                    logger.debug(f"Ingested {path}: {len(chunks)} chunks ([{chunk_tokens_str}] tokens)")
                
            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[str, Any]]) -> int:
        """Ingest files from in-memory entries (path/content)."""
        return self.ingestor.ingest_file_entries(file_entries)
    
    def ingest_file_entries_batched(self, file_entries: List[Dict[str, Any]], chunk_tokens: int) -> int:
        """Ingest multiple file entries into batched chunks within token limit."""
        return self.ingestor.ingest_file_entries_batched(file_entries, chunk_tokens)

    def run(self, max_iterations: int = 50) -> MemoryStore:
        """Run the review loop."""
        from normal.ai.ai import generate_review, LLMError
        
        model = os.getenv("GAR_MODEL", "qwen3-coder:30b")
        logger.debug(f"Starting review with model: {model}")
        
        iterations = 0
        successful = 0
        failed = 0
        
        while iterations < max_iterations:
            chunk = self.planner.select_next()
            if not chunk:
                logger.debug("No more chunks to process")
                break
            
            stored = self.offload.get(chunk.chunk_id)
            if not stored:
                iterations += 1
                continue
            
            chunk_content, meta = stored
            logger.debug(f"Processing chunk {iterations + 1}: {chunk.chunk_id}")
            
            # Build context
            context = self.context_builder.build(
                repo_summary=self.memory.repo_summary,
                filecards=list(self.memory.filecards.values()),
                chunk_content=chunk_content,
            )
            
            # Call LLM
            result = self._review_chunk(model, context, chunk.chunk_id, meta)
            self.results.append(result)
            
            if result.success and result.parsed_data:
                self._ingest_result(result.parsed_data)
                successful += 1
            else:
                failed += 1
            
            iterations += 1
            
            # Periodic synthesis
            if iterations % 10 == 0:
                self._synthesize()
        
        # Final synthesis
        self._synthesize()
        
        logger.debug(f"Review complete: {successful} successful, {failed} failed out of {iterations} chunks")
        return self.memory

    def _review_chunk(
        self,
        model: str,
        context: Dict,
        chunk_id: str,
        meta: Optional[Dict] = None,
    ) -> ReviewResult:
        """Review a single chunk."""
        from normal.ai.ai import generate_review, LLMError
        
        user_prompt = build_user_prompt(context, chunk_id, meta)
        system_prompt = SYSTEM_PROMPT_MULTI if meta and meta.get("file_paths") else SYSTEM_PROMPT
        
        try:
            response = generate_review(
                model=model,
                prompt=(system_prompt, user_prompt),
                temperature=0.3,
                max_tokens=8000,  # 긴 코드 분석을 위해 증가
                max_retries=2,
            )
            
            # Try to parse JSON
            try:
                parsed = json.loads(response)
                return ReviewResult(
                    chunk_id=chunk_id,
                    success=True,
                    raw_response=response,
                    parsed_data=parsed,
                )
            except json.JSONDecodeError as e:
                logger.warning(f"JSON parse failed for {chunk_id}: {e}")
                return ReviewResult(
                    chunk_id=chunk_id,
                    success=False,
                    raw_response=response,
                    error_message=f"JSON parse error: {e}",
                )
                
        except LLMError as e:
            logger.error(f"LLM error for {chunk_id}: {e}")
            return ReviewResult(
                chunk_id=chunk_id,
                success=False,
                raw_response="",
                error_message=str(e),
            )

    def _ingest_result(self, data: Dict) -> None:
        """Ingest parsed LLM result into memory."""
        # FileCard(s)
        fc_list = data.get("file_cards")
        if fc_list and isinstance(fc_list, list):
            for fc_data in fc_list:
                if not isinstance(fc_data, dict):
                    continue
                card = FileCard(
                    file_path=fc_data.get("file_path", "unknown"),
                    role_summary=fc_data.get("role_summary", ""),
                    key_symbols=fc_data.get("key_symbols") or [],
                    deps_hints=fc_data.get("deps_hints") or [],
                    risk_flags=fc_data.get("risk_flags") or [],
                )
                self.memory.upsert_filecard(card)
        else:
            fc_data = data.get("file_card")
            if fc_data and isinstance(fc_data, dict):
                card = FileCard(
                    file_path=fc_data.get("file_path", "unknown"),
                    role_summary=fc_data.get("role_summary", ""),
                    key_symbols=fc_data.get("key_symbols") or [],
                    deps_hints=fc_data.get("deps_hints") or [],
                    risk_flags=fc_data.get("risk_flags") or [],
                )
                self.memory.upsert_filecard(card)

        # Findings
        findings_data = data.get("findings") or []
        findings: List[Finding] = []
        for f in findings_data:
            if not isinstance(f, dict):
                continue
            try:
                line_range = f.get("line_range", [0, 0])
                if isinstance(line_range, list) and len(line_range) >= 2:
                    line_range = (line_range[0], line_range[1])
                else:
                    line_range = (0, 0)
                
                findings.append(Finding(
                    category=f.get("category", ""),
                    subtype=f.get("subtype", ""),
                    severity=f.get("severity", "medium"),
                    file_path=f.get("file_path", ""),
                    line_range=line_range,
                    evidence=f.get("evidence", "")[:500],
                    message=f.get("message", ""),
                    min_fix=f.get("min_fix", ""),
                ))
            except Exception as e:
                logger.debug(f"Failed to parse finding: {e}")
                continue
        
        if findings:
            self.memory.add_findings(findings)

        # Repo summary
        rs = data.get("repo_summary")
        if rs and isinstance(rs, dict):
            self.memory.update_repo_summary(
                architecture_overview=rs.get("architecture_overview", ""),
                hotspots=rs.get("hotspots") or [],
                recurring_patterns=rs.get("recurring_patterns") or [],
            )

    def _synthesize(self) -> None:
        """Synthesize findings into hotspots."""
        freq: Dict[str, int] = {}
        for f in self.memory.findings:
            freq[f.file_path] = freq.get(f.file_path, 0) + 1
        
        hotspots = sorted(freq.items(), key=lambda x: x[1], reverse=True)
        self.memory.repo_summary.hotspots = [fp for fp, _ in hotspots[:5]]


# -----------------------------
# Public API
# -----------------------------


def run_context_review(
    file_paths: List[str],
    max_iterations: int = 50,
    chunking_mode: str = "fast",
    progress_handle=None,
) -> Tuple[MemoryStore, List[ReviewResult]]:
    """Run context-aware code review.
    
    Args:
        file_paths: List of file paths to review
        max_iterations: Maximum number of chunks to process
        chunking_mode: Chunking strategy ("fast" or "thinking")
        progress_handle: Progress handle for file reading progress tracking
        
    Returns:
        Tuple of (MemoryStore with results, list of ReviewResult)
    """
    # Input validation
    if not file_paths:
        logger.error("Empty file_paths list provided")
        raise ValueError("file_paths cannot be empty")
    
    # Validate all paths are strings
    invalid_paths = [p for p in file_paths if not isinstance(p, str)]
    if invalid_paths:
        logger.error(f"Non-string paths found: {invalid_paths}")
        raise TypeError("All file paths must be strings")
    
    if chunking_mode not in ("fast", "thinking"):
        logger.error(f"Invalid chunking_mode: {chunking_mode}")
        raise ValueError(f"chunking_mode must be 'fast' or 'thinking', got {chunking_mode}")

    chunk_tokens = 30000 if chunking_mode == "fast" else 6000
    chunker = Chunker(chunk_tokens=chunk_tokens, overlap_tokens=500)
    runner = IterativeLoopRunner(chunker=chunker, progress_handle=progress_handle)

    if chunking_mode == "fast":
        num_chunks = runner.ingest_files_batched(file_paths, chunk_tokens=chunk_tokens)
    else:
        num_chunks = runner.ingest_files(file_paths)
    
    if num_chunks == 0:
        logger.warning("No chunks created from provided files")
    
    logger.debug(
        f"Prepared {num_chunks} chunks from {len(file_paths)} files (mode: {chunking_mode})"
    )
    
    memory = runner.run(max_iterations=max_iterations)
    return memory, runner.results


def run_context_review_from_entries(
    file_entries: List[Dict[str, Any]],
    max_iterations: int = 50,
    chunking_mode: str = "fast",
    progress_handle=None,
) -> Tuple[MemoryStore, List[ReviewResult]]:
    """Run context-aware code review from in-memory file entries.
    
    Args:
        file_entries: List of file entries (dict with "path" and "content")
        max_iterations: Maximum number of chunks to process
        chunking_mode: Chunking strategy ("fast" or "thinking")
        progress_handle: Progress handle for file reading progress tracking (not used for file_entries as they're already in memory)
        
    Returns:
        Tuple of (MemoryStore with results, list of ReviewResult)
    """
    if not file_entries:
        logger.error("Empty file_entries list provided")
        memory = MemoryStore()
        return memory, []

    if chunking_mode not in ("fast", "thinking"):
        logger.error(f"Invalid chunking_mode: {chunking_mode}")
        raise ValueError(f"chunking_mode must be 'fast' or 'thinking', got {chunking_mode}")

    chunk_tokens = 30000 if chunking_mode == "fast" else 6000
    chunker = Chunker(chunk_tokens=chunk_tokens, overlap_tokens=500)
    runner = IterativeLoopRunner(chunker=chunker, progress_handle=progress_handle)

    if chunking_mode == "fast":
        num_chunks = runner.ingest_file_entries_batched(file_entries, chunk_tokens=chunk_tokens)
    else:
        num_chunks = runner.ingest_file_entries(file_entries)
    
    if num_chunks == 0:
        logger.warning("No chunks created from provided files")
    
    logger.debug(
        f"Prepared {num_chunks} chunks from {len(file_entries)} files (mode: {chunking_mode})"
    )
    
    memory = runner.run(max_iterations=max_iterations)
    return memory, runner.results
