#!/usr/bin/env python3
"""Context engineering for security 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 (SecurityPlanner for security mode)
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 security.prompts import SYSTEM_PROMPT_JSON, SYSTEM_PROMPT_JSON_MULTI, build_security_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


# -----------------------------
# 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


class SecurityPlanner(Planner):
    """보안 모드용 Planner - 보안 관련 키워드에 높은 가중치 부여."""
    
    # 보안 관련 키워드 (높은 우선순위)
    SECURITY_KEYWORDS_HIGH = [
        "auth", "authentication", "authorization", "login", "logout",
        "password", "token", "secret", "key", "credential", "session",
        "security", "secure", "encrypt", "decrypt", "hash", "salt",
        "injection", "sql", "xss", "csrf", "ssl", "tls",
        "validate", "sanitize", "escape", "whitelist", "blacklist",
        "permission", "role", "access", "privilege", "admin",
    ]
    
    # 보안 관련 키워드 (중간 우선순위)
    SECURITY_KEYWORDS_MEDIUM = [
        "api", "endpoint", "route", "handler", "controller",
        "middleware", "filter", "guard", "policy",
        "user", "account", "profile", "settings",
        "payment", "transaction", "order", "checkout",
    ]

    def _score(self, chunk: Chunk) -> int:
        """보안 모드용 점수 계산 - 보안 키워드에 높은 가중치."""
        path = chunk.file_path.lower()
        score = 0
        
        # 보안 키워드 체크 (높은 우선순위)
        for keyword in self.SECURITY_KEYWORDS_HIGH:
            if keyword in path:
                score += 50  # 높은 가중치
                break  # 첫 번째 매치만 사용
        
        # 보안 키워드 체크 (중간 우선순위)
        for keyword in self.SECURITY_KEYWORDS_MEDIUM:
            if keyword in path:
                score += 25  # 중간 가중치
                break
        
        # 기존 우선순위 로직
        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,
        system_prompt: Optional[str] = None,
        build_user_prompt_fn: Optional[Any] = None,
    ):
        self.offload = offload or InMemoryOffloadStore()
        self.memory = memory or MemoryStore()
        self.chunker = chunker or Chunker()
        self.planner = planner if planner is not None else SecurityPlanner()  # 기본값을 SecurityPlanner로
        self.context_builder = context_builder or ContextBuilder()
        self.ingestor = ChunkIngestor(
            offload=self.offload,
            planner=self.planner,
            chunker=self.chunker,
        )
        self.results: List[ReviewResult] = []
        # 프롬프트 주입 지원
        self.custom_system_prompt = system_prompt
        self.custom_build_user_prompt = build_user_prompt_fn

    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("MODEL_NAME") or os.getenv("GAR_MODEL") or "qwen3-coder:30b"
        logger.info(f"Starting security review with model: {model}")
        
        iterations = 0
        successful = 0
        failed = 0
        
        while iterations < max_iterations:
            chunk = self.planner.select_next()
            if not chunk:
                logger.info("No more chunks to process")
                break
            
            stored = self.offload.get(chunk.chunk_id)
            if not stored:
                iterations += 1
                continue
            
            chunk_content, meta = stored
            logger.info(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.info(f"Security 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
        
        # 커스텀 프롬프트 함수가 있으면 사용, 없으면 보안 모드 기본 사용
        if self.custom_build_user_prompt:
            user_prompt = self.custom_build_user_prompt(context, chunk_id, meta)
        else:
            user_prompt = build_security_user_prompt(context, chunk_id, meta)
        
        # 커스텀 system_prompt가 있으면 사용, 없으면 보안 모드 기본 사용
        if self.custom_system_prompt:
            system_prompt = self.custom_system_prompt
        else:
            system_prompt = SYSTEM_PROMPT_JSON_MULTI if meta and meta.get("file_paths") else SYSTEM_PROMPT_JSON
        
        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_from_entries(
    file_entries: List[Dict[str, Any]],
    max_iterations: int = 50,
    chunking_mode: str = "fast",
    system_prompt: Optional[str] = None,
    build_user_prompt_fn: Optional[Any] = None,
    planner: Optional[Planner] = None,
) -> Tuple[MemoryStore, List[ReviewResult]]:
    """Run context-aware security 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")
        system_prompt: Optional custom system prompt
        build_user_prompt_fn: Optional custom function to build user prompt
        planner: Optional custom planner (defaults to SecurityPlanner)
        
    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)
    # 기본값을 SecurityPlanner로 사용
    security_planner = planner if planner is not None else SecurityPlanner()
    runner = IterativeLoopRunner(
        chunker=chunker,
        planner=security_planner,
        system_prompt=system_prompt,
        build_user_prompt_fn=build_user_prompt_fn,
    )

    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.info(
        f"Prepared {num_chunks} chunks from {len(file_entries)} files (mode: {chunking_mode})"
    )
    
    memory = runner.run(max_iterations=max_iterations)
    return memory, runner.results
