#!/usr/bin/env python3
"""
app.py — Project File Aggregator for LLM Input (with Gradio UI)

اضافه شده: پشتیبانی از فایل‌های ZIP به عنوان ورودی
"""

import hashlib
import tempfile
import zipfile
from pathlib import Path
from datetime import datetime
from concurrent.futures import ThreadPoolExecutor
import os
import shutil

import gradio as gr

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

IGNORE_DIRS_DEFAULT = (
    ".git,.idea,.vscode,__pycache__,.pytest_cache,node_modules,venv,.venv,"
    "env,build,dist,.next,.cache,site-packages,.mypy_cache,target,.gradle,bin,obj"
)

IGNORE_EXTENSIONS = {
    ".exe", ".dll", ".so", ".dylib", ".png", ".jpg", ".jpeg", ".gif", ".ico",
    ".svg", ".webp", ".zip", ".rar", ".7z", ".tar", ".gz", ".pdf", ".mp4",
    ".mp3", ".wav", ".mov", ".woff", ".woff2", ".ttf", ".eot", ".pyc",
    ".class", ".jar", ".db", ".sqlite3", ".lock", ".bin", ".dat",
}

LANGUAGE_MAP = {
    ".py": "Python", ".js": "JavaScript", ".jsx": "JavaScript (JSX)",
    ".ts": "TypeScript", ".tsx": "TypeScript (TSX)", ".html": "HTML",
    ".css": "CSS", ".scss": "SCSS", ".json": "JSON", ".yaml": "YAML",
    ".yml": "YAML", ".toml": "TOML", ".xml": "XML", ".md": "Markdown",
    ".sql": "SQL", ".java": "Java", ".go": "Go", ".rb": "Ruby",
    ".php": "PHP", ".c": "C", ".cpp": "C++", ".cs": "C#", ".swift": "Swift",
    ".kt": "Kotlin", ".sh": "Shell", ".txt": "Text", ".env": "Env",
    ".ini": "INI", ".cfg": "Config", ".rst": "reStructuredText",
}

EXTENSION_CATEGORIES = {
    "Python": {".py"},
    "JS / TypeScript": {".js", ".jsx", ".ts", ".tsx"},
    "Web (HTML/CSS)": {".html", ".css", ".scss"},
    "Config / Data (JSON, YAML, ...)": {".json", ".yaml", ".yml", ".toml", ".ini", ".cfg", ".xml", ".env"},
    "Documentation (Markdown/Text)": {".md", ".txt", ".rst"},
    "Other Languages (Java, Go, C++, ...)": {".java", ".go", ".rb", ".php", ".c", ".cpp", ".cs", ".swift", ".kt", ".sql", ".sh"},
}

CHARS_PER_TOKEN = 4.0


# ---------------------------------------------------------------------------
# Filtering logic
# ---------------------------------------------------------------------------

def should_ignore_dir(dir_name: str, extra_ignore_dirs: set) -> bool:
    return dir_name in extra_ignore_dirs or dir_name.startswith(".")


# ---------------------------------------------------------------------------
# ZIP handling functions (جدید)
# ---------------------------------------------------------------------------

def is_zip_file(file_path: str) -> bool:
    """بررسی می‌کند که آیا فایل یک ZIP معتبر است"""
    try:
        return zipfile.is_zipfile(file_path)
    except:
        return False


def process_zip_input(zip_path: str, output_dir: str = None) -> Path:
    """
    فایل ZIP را استخراج کرده و مسیر پوشه استخراج شده را برمی‌گرداند
    اگر output_dir مشخص نشده باشد، از tempfile استفاده می‌کند
    """
    zip_path = Path(zip_path)
    
    # اگر دایرکتوری خروجی مشخص نشده، از دایرکتوری موقت استفاده کن
    if output_dir is None:
        temp_dir = tempfile.mkdtemp(prefix="context_llm_zip_")
        output_dir = Path(temp_dir)
    else:
        output_dir = Path(output_dir)
        output_dir.mkdir(parents=True, exist_ok=True)
    
    # استخراج فایل‌های زیپ
    try:
        with zipfile.ZipFile(zip_path, 'r') as zip_ref:
            # بررسی وجود فایل‌های مخرب (مسیرهای نسبی)
            for member in zip_ref.namelist():
                # جلوگیری از path traversal attacks
                member_path = (output_dir / member).resolve()
                if not str(member_path).startswith(str(output_dir.resolve())):
                    raise Exception(f"فایل مخرب در ZIP: {member}")
            
            zip_ref.extractall(output_dir)
            return output_dir
    except Exception as e:
        # در صورت خطا، دایرکتوری موقت را پاک کن
        if output_dir is not None and str(output_dir).startswith(tempfile.gettempdir()):
            shutil.rmtree(output_dir, ignore_errors=True)
        raise Exception(f"خطا در استخراج ZIP: {str(e)}")


def cleanup_temp_dir(temp_dir: Path):
    """پاک کردن دایرکتوری موقت (اگر از tempfile ایجاد شده باشد)"""
    if temp_dir is not None and str(temp_dir).startswith(tempfile.gettempdir()):
        try:
            shutil.rmtree(temp_dir, ignore_errors=True)
        except:
            pass


# ---------------------------------------------------------------------------
# Core scanning logic (تغییر داده شده برای پشتیبانی از ورودی‌های مختلف)
# ---------------------------------------------------------------------------

def scan_directory(
    root_path: Path,
    ignore_dirs: set = None,
    max_file_size_mb: int = 5,
    include_hidden: bool = False,
    use_gitignore: bool = False,
    is_temp_dir: bool = False  # جدید: مشخص می‌کند که دایرکتوری موقت است
) -> tuple[list, list, int, list]:
    """
    اسکن دایرکتوری و جمع‌آوری فایل‌ها
    برمی‌گرداند: (file_infos, root_name, total_files, category_counts)
    """
    if ignore_dirs is None:
        ignore_dirs = set(IGNORE_DIRS_DEFAULT.split(","))

    root_path = Path(root_path)
    if not root_path.exists():
        raise FileNotFoundError(f"مسیر وجود ندارد: {root_path}")

    # برای پشتیبانی از فایل‌های زیپ، اگر ورودی فایل است، آن را استخراج کن
    if root_path.is_file() and is_zip_file(str(root_path)):
        # این یک فایل ZIP است
        temp_dir = process_zip_input(str(root_path))
        is_temp_dir = True
        root_path = temp_dir
        root_name = root_path.name
    elif root_path.is_file():
        # اگر فایل معمولی است، فقط همان را پردازش کن
        root_name = root_path.name
        # ایجاد یک دایرکتوری موقت برای پردازش فایل تنها
        temp_dir = tempfile.mkdtemp(prefix="context_llm_single_")
        shutil.copy2(root_path, Path(temp_dir) / root_name)
        root_path = Path(temp_dir) / root_name
        is_temp_dir = True
    else:
        root_name = root_path.name
        is_temp_dir = False

    files = []
    total_size = 0
    total_files = 0
    category_counts = {cat: 0 for cat in EXTENSION_CATEGORIES}

    for item in root_path.rglob("*"):
        if item.is_dir():
            # Check if directory should be ignored
            if should_ignore_dir(item.name, ignore_dirs) or item.name.startswith("."):
                continue
            continue

        # Check extension
        ext = item.suffix.lower()
        if ext in IGNORE_EXTENSIONS:
            continue

        # Skip if hidden file (unless requested)
        if not include_hidden and item.name.startswith("."):
            continue

        # Skip if file too large
        if item.stat().st_size > max_file_size_mb * 1024 * 1024:
            continue

        # Determine category
        category = "Other"
        for cat, exts in EXTENSION_CATEGORIES.items():
            if ext in exts:
                category = cat
                break

        category_counts[category] = category_counts.get(category, 0) + 1

        # Read file content
        try:
            content = item.read_text(encoding="utf-8", errors="ignore")
        except Exception:
            # Skip files that can't be read as text
            continue

        # Relative path from root
        rel_path = str(item.relative_to(root_path))

        files.append({
            "path": rel_path,
            "content": content,
            "ext": ext,
            "language": LANGUAGE_MAP.get(ext, "Unknown"),
            "category": category,
            "size": item.stat().st_size,
        })
        total_files += 1

    return files, root_name, total_files, category_counts, is_temp_dir


# ---------------------------------------------------------------------------
# Chunking and tokenization logic (بدون تغییر)
# ---------------------------------------------------------------------------

def chunk_file_content(file_info: dict, chunk_size_tokens: int = 4000, overlap_tokens: int = 200) -> list:
    """Split file content into chunks based on token estimation."""
    content = file_info["content"]
    # Estimate token count
    estimated_tokens = len(content) / CHARS_PER_TOKEN
    if estimated_tokens <= chunk_size_tokens:
        return [content]

    # Rough chunking based on character count (not exact token count)
    chunk_chars = int(chunk_size_tokens * CHARS_PER_TOKEN)
    overlap_chars = int(overlap_tokens * CHARS_PER_TOKEN)

    chunks = []
    start = 0
    while start < len(content):
        end = min(start + chunk_chars, len(content))
        # Try to cut at newline or space
        if end < len(content):
            # Look for newline
            newline_pos = content.rfind("\n", start, end)
            if newline_pos != -1 and newline_pos > start:
                end = newline_pos + 1
            else:
                # Look for space
                space_pos = content.rfind(" ", start, end)
                if space_pos != -1 and space_pos > start:
                    end = space_pos + 1

        chunks.append(content[start:end])
        start = end - overlap_chars
        if start < 0:
            start = 0

    return chunks


def estimate_tokens(text: str) -> int:
    """Estimate token count based on character count."""
    return int(len(text) / CHARS_PER_TOKEN)


# ---------------------------------------------------------------------------
# Main aggregation function (تغییر داده شده)
# ---------------------------------------------------------------------------

def build_context(
    input_path: str,
    ignore_dirs: str = None,
    max_file_size_mb: int = 5,
    chunk_size_tokens: int = 4000,
    overlap_tokens: int = 200,
    max_total_tokens: int = 100000,
    include_hidden: bool = False,
    use_gitignore: bool = False,
    include_hash: bool = True,
    include_tree: bool = True,
) -> tuple[str, str, dict]:
    """
    Build aggregated context for LLM input.
    برمی‌گرداند: (context_text, context_hash, stats)
    """
    # Parse ignore dirs
    if ignore_dirs:
        ignore_set = set(ignore_dirs.split(","))
    else:
        ignore_set = set(IGNORE_DIRS_DEFAULT.split(","))

    # Scan directory
    root_path = Path(input_path)
    is_temp = False
    
    try:
        files, root_name, total_files, category_counts, is_temp = scan_directory(
            root_path=root_path,
            ignore_dirs=ignore_set,
            max_file_size_mb=max_file_size_mb,
            include_hidden=include_hidden,
            use_gitignore=use_gitignore,
        )
    except Exception as e:
        return f"❌ خطا: {str(e)}", "", {}

    if not files:
        return "⚠️ هیچ فایل متنی قابل خواندنی در مسیر مشخص شده یافت نشد.", "", {}

    # Sort files by category then path
    files.sort(key=lambda x: (x["category"], x["path"]))

    # Build context
    context_parts = []
    context_parts.append(f"# Project Context: {root_name}")
    context_parts.append(f"# Generated: {datetime.now().isoformat()}")
    context_parts.append(f"# Total files: {len(files)}")
    context_parts.append("")

    # Add tree structure
    if include_tree:
        tree = build_tree_structure(files)
        context_parts.append("## Directory Structure")
        context_parts.append("```")
        context_parts.append(tree)
        context_parts.append("```")
        context_parts.append("")

    total_est_tokens = 0
    file_chunks = []

    for file_info in files:
        chunks = chunk_file_content(
            file_info,
            chunk_size_tokens=chunk_size_tokens,
            overlap_tokens=overlap_tokens,
        )

        for i, chunk in enumerate(chunks):
            # Estimate tokens in this chunk
            chunk_tokens = estimate_tokens(chunk)

            # Check if adding this chunk would exceed max tokens
            if total_est_tokens + chunk_tokens > max_total_tokens:
                # Add truncation message and stop
                remaining = max_total_tokens - total_est_tokens
                if remaining > 0:
                    # Truncate chunk to fit remaining tokens
                    char_limit = int(remaining * CHARS_PER_TOKEN)
                    truncated = chunk[:char_limit]
                    header = f"### {file_info['path']} [{file_info['language']}] (chunk {i+1}/{len(chunks)}) — TRUNCATED\n\n"
                    context_parts.append(header + truncated)
                context_parts.append("\n\n⚠️ **Context truncated due to token limit**")
                context_parts.append(f"Limit: {max_total_tokens:,} tokens")
                break

            # Add file header
            chunk_header = f"### {file_info['path']} [{file_info['language']}] (chunk {i+1}/{len(chunks)})\n\n"
            context_parts.append(chunk_header + chunk)
            total_est_tokens += chunk_tokens + estimate_tokens(chunk_header)

        else:
            # If no break, continue to next file
            continue
        # If break was hit, exit outer loop too
        break

    context_text = "\n\n".join(context_parts)

    # Generate hash
    if include_hash:
        context_hash = hashlib.md5(context_text.encode()).hexdigest()
    else:
        context_hash = ""

    # Stats
    stats = {
        "files": len(files),
        "total_tokens": total_est_tokens,
        "categories": category_counts,
        "chunk_size": chunk_size_tokens,
        "overlap": overlap_tokens,
        "max_tokens": max_total_tokens,
        "truncated": total_est_tokens >= max_total_tokens,
    }

    # Cleanup temp directory if needed
    if is_temp and root_path:
        cleanup_temp_dir(root_path.parent if root_path.is_file() else root_path)

    return context_text, context_hash, stats


def build_tree_structure(files: list) -> str:
    """Build a simple tree structure from file paths."""
    tree = {}
    for file_info in files:
        parts = Path(file_info["path"]).parts
        current = tree
        for part in parts[:-1]:
            current = current.setdefault(part, {})
        current[parts[-1]] = None

    def render_tree(node, prefix=""):
        lines = []
        items = sorted(node.items())
        for i, (key, value) in enumerate(items):
            is_last = (i == len(items) - 1)
            if value is None:
                lines.append(f"{prefix}{'└── ' if is_last else '├── '}{key}")
            else:
                lines.append(f"{prefix}{'└── ' if is_last else '├── '}{key}/")
                lines.extend(render_tree(
                    value,
                    prefix + ("    " if is_last else "│   ")
                ))
        return lines

    return "\n".join(render_tree(tree)) if tree else "(empty)"


# ---------------------------------------------------------------------------
# Gradio UI (تغییر داده شده)
# ---------------------------------------------------------------------------

def create_ui():
    with gr.Blocks(title="Context LLM Builder", theme=gr.themes.Soft()) as demo:
        gr.Markdown("""
        # 📚 Context LLM Builder

        این ابزار فایل‌های یک پروژه را اسکن کرده و آن‌ها را به‌صورت یک متن یکپارچه
        برای ارسال به مدل‌های زبانی بزرگ (LLM) آماده می‌کند.

        **✅ پشتیبانی از ورودی‌ها:**
        - **پوشه** (کل پروژه)
        - **فایل ZIP** (استخراج خودکار در حافظه)
        - **فایل تکی** (مانند یک فایل پایتون)
        """)

        with gr.Row():
            with gr.Column(scale=1):
                input_path = gr.Textbox(
                    label="مسیر پوشه / فایل ZIP / فایل تکی",
                    placeholder="/path/to/your/project یا /path/to/file.zip",
                    value=".",
                )

                with gr.Accordion("⚙️ تنظیمات پیشرفته", open=False):
                    ignore_dirs = gr.Textbox(
                        label="پوشه‌های نادیده گرفته شده (با کاما جدا کنید)",
                        value=IGNORE_DIRS_DEFAULT,
                    )
                    max_file_size_mb = gr.Slider(
                        label="حداکثر حجم فایل (MB)",
                        minimum=1,
                        maximum=50,
                        value=5,
                        step=1,
                    )
                    chunk_size_tokens = gr.Slider(
                        label="حجم هر تکه (توکن)",
                        minimum=500,
                        maximum=16000,
                        value=4000,
                        step=100,
                    )
                    overlap_tokens = gr.Slider(
                        label="همپوشانی بین تکه‌ها (توکن)",
                        minimum=0,
                        maximum=1000,
                        value=200,
                        step=50,
                    )
                    max_total_tokens = gr.Slider(
                        label="حداکثر توکن کل خروجی",
                        minimum=10000,
                        maximum=200000,
                        value=100000,
                        step=1000,
                    )
                    include_hidden = gr.Checkbox(
                        label="شامل فایل‌های پنهان (.bashrc و ...)",
                        value=False,
                    )
                    include_hash = gr.Checkbox(
                        label="شامل هش (MD5) متن خروجی",
                        value=True,
                    )
                    include_tree = gr.Checkbox(
                        label="شامل ساختار دایرکتوری",
                        value=True,
                    )

                build_btn = gr.Button("🚀 ساخت کانتکست", variant="primary", size="lg")

            with gr.Column(scale=2):
                status = gr.Markdown("⬅️ تنظیمات را وارد کرده و دکمه را بزنید")
                with gr.Row():
                    with gr.Column(scale=3):
                        output_text = gr.Textbox(
                            label="📄 متن کانتکست",
                            lines=30,
                            max_lines=50,
                            interactive=False,
                        )
                    with gr.Column(scale=1):
                        hash_text = gr.Textbox(
                            label="🔑 هش (MD5)",
                            lines=1,
                            interactive=False,
                            max_lines=1,
                        )
                        stats_json = gr.JSON(
                            label="📊 آمار",
                            value={},
                        )

                with gr.Row():
                    copy_btn = gr.Button("📋 کپی متن", variant="secondary", size="sm")
                    download_btn = gr.DownloadButton("💾 دانلود فایل", variant="secondary", size="sm")

        # Events
        build_btn.click(
            fn=build_context,
            inputs=[
                input_path,
                ignore_dirs,
                max_file_size_mb,
                chunk_size_tokens,
                overlap_tokens,
                max_total_tokens,
                include_hidden,
                include_hash,
                include_tree,
            ],
            outputs=[output_text, hash_text, stats_json],
        )

        # Copy button - اصلاح شده
        copy_btn.click(
            fn=lambda x: x,
            inputs=[output_text],
            outputs=[],
            js="(text) => { navigator.clipboard.writeText(text); return text; }",
        )

        # Download button - اصلاح شده با روش ساده‌تر
        def download_file(text):
            if not text:
                return None
            temp = tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False)
            temp.write(text)
            temp.close()
            return temp.name

        download_btn.click(
            fn=download_file,
            inputs=[output_text],
            outputs=[gr.File(label="📥 دانلود فایل")],
        )

    return demo

if __name__ == "__main__":
    demo = create_ui()
    demo.launch(server_name="0.0.0.0", server_port=7860, share=False)