import os
import json
import argparse
import pyperclip
from rich.console import Console
from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn, TaskProgressColumn
from rich.panel import Panel
from rich.text import Text
import hashlib

# Initialize rich console
console = Console()

# Constants
EXCLUDED_DIRS = {
    "migrations",
    "static",
    "media",
    "__pycache__",
    ".git",
    "venv",
    "node_modules",
}
EXCLUDED_EXTENSIONS = {
    ".pyc", ".css", ".scss", ".png", ".jpg", ".jpeg", ".svg", ".sqlite3"
}
EXCLUDED_FILES = {"manage.py", "wsgi.py", "asgi.py", "package-lock.json"}
VALID_EXTENSIONS = {".py", ".js", ".html", ".json"}

HASH_CACHE_FILE = "hash_cache.json"   # Stores file-level hashes and old summaries
SUMMARIES_OUTPUT = "summaries.json"   # Final combined summaries

def valid_source_file(file_path):
    """
    Check if the file is a valid source file based on its extension and name.
    """
    _, ext = os.path.splitext(file_path)
    filename = os.path.basename(file_path)
    return (
        ext.lower() in VALID_EXTENSIONS
        and ext.lower() not in EXCLUDED_EXTENSIONS
        and filename not in EXCLUDED_FILES
    )

def scan_repo(repo_path):
    """
    Recursively walk the repo and return a list of valid source file paths.
    Skips excluded directories and files.
    """
    file_paths = []
    for root, dirs, files in os.walk(repo_path):
        dirs[:] = [d for d in dirs if d not in EXCLUDED_DIRS]  # Skip excluded directories
        for f in files:
            fp = os.path.join(root, f)
            if valid_source_file(fp):
                file_paths.append(fp)
    return file_paths

def chunk_file(file_path, lines_per_chunk=50):
    """
    Splits the file content into chunks of N lines.
    """
    with open(file_path, "r", encoding="utf-8", errors="ignore") as f:
        content = f.readlines()

    chunks = []
    for i in range(0, len(content), lines_per_chunk):
        chunk_lines = content[i:i+lines_per_chunk]
        chunks.append("".join(chunk_lines))
    return chunks

def calculate_file_hash(file_path):
    """
    Compute a hash (SHA-256) for the file content.
    """
    hasher = hashlib.sha256()
    try:
        with open(file_path, "rb") as f:
            for chunk in iter(lambda: f.read(8192), b""):
                hasher.update(chunk)
        return hasher.hexdigest()
    except Exception as e:
        console.print(f"[red]Error reading file for hashing: {file_path} - {e}[/red]")
        return None

def load_hash_cache():
    """
    Load existing hash/summaries from HASH_CACHE_FILE, if exists.
    Returns a dict: { file_path: { 'hash': <str>, 'summaries': [ {chunk_id, summary} ... ] } }
    """
    if not os.path.exists(HASH_CACHE_FILE):
        return {}
    try:
        with open(HASH_CACHE_FILE, "r", encoding="utf-8") as f:
            return json.load(f)
    except:
        return {}

def save_hash_cache(cache_data):
    """
    Save hash cache to file.
    """
    with open(HASH_CACHE_FILE, "w", encoding="utf-8") as f:
        json.dump(cache_data, f, indent=2)

def summarize_chunk(chunk, progress=None):
    """
    Summarize a code chunk using OpenAI's GPT-4o model.
    Assumes OPENAI_API_KEY is set in environment variables.
    """
    from openai import OpenAI
    import os
    
    client = OpenAI()  # Uses OPENAI_API_KEY from environment by default
    
    try:
        if progress:
            progress.update(progress.task_ids[-1], description="[yellow]🤔 Analyzing with AI...[/yellow]")
        
        completion = client.chat.completions.create(
            model="gpt-4o",  
            store=True,  # Store the completion
            messages=[
                {
                    "role": "user",
                    "content": f"Please summarize this code chunk concisely:\n\n{chunk}"
                }
            ]
        )
        
        return completion.choices[0].message.content.strip()
    except Exception as e:
        error_msg = f"Error summarizing chunk: {str(e)}"
        console.print(f"[red]❌ {error_msg}[/red]")
        return error_msg

def main(repo_path="./my_repo"):
    console.print(Panel.fit(
        Text("🚀 Repository Summarizer", justify="center", style="bold cyan"),
        subtitle="Let's make your code speak!"
    ))

    # Load the existing hash cache to skip unchanged files
    hash_cache = load_hash_cache()
    
    with Progress(
        SpinnerColumn(),
        TextColumn("[bold blue]{task.description}"),
        BarColumn(),
        TaskProgressColumn(),
        console=console
    ) as progress:
        # Scan repository task
        scan_task = progress.add_task("[cyan]🔍 Scanning repository...", total=None)
        files = scan_repo(repo_path)
        progress.update(scan_task, total=1, completed=1)
        console.print(f"[green]✨ Found {len(files)} source files to analyze![/green]")
        
        overall_task = progress.add_task("[cyan]📝 Processing files...", total=len(files))
        combined_summaries = []
        
        for f in files:
            file_task = progress.add_task(
                f"[yellow]Processing {os.path.basename(f)}...",
                total=1
            )

            # Calculate current file hash
            current_hash = calculate_file_hash(f)
            old_hash_data = hash_cache.get(f, {})
            old_hash = old_hash_data.get("hash")
            old_summaries = old_hash_data.get("summaries", [])

            if current_hash and current_hash == old_hash:
                # File unchanged => reuse old summaries
                progress.update(file_task, description="[green]Skipping unchanged file[/green]")
                for s in old_summaries:
                    combined_summaries.append({
                        "file": f,
                        "chunk_id": s["chunk_id"],
                        "summary": s["summary"]
                    })
            else:
                # File changed or not in cache => re-summarize
                chunks = chunk_file(f, lines_per_chunk=30)
                new_summaries = []
                for idx, chunk_content in enumerate(chunks):
                    chunk_task = progress.add_task(
                        f"[blue]Analyzing chunk {idx+1}/{len(chunks)}...",
                        total=1
                    )
                    summary = summarize_chunk(chunk_content, progress)
                    new_summaries.append({"chunk_id": idx, "summary": summary})
                    combined_summaries.append({
                        "file": f,
                        "chunk_id": idx,
                        "summary": summary
                    })
                    progress.update(chunk_task, completed=1)
                    progress.remove_task(chunk_task)

                # Update hash_cache after summarizing
                hash_cache[f] = {
                    "hash": current_hash,
                    "summaries": new_summaries
                }
            
            progress.update(file_task, completed=1)
            progress.remove_task(file_task)
            progress.update(overall_task, advance=1)

        # Save combined summaries to JSON
        save_task = progress.add_task("[yellow]💾 Saving summaries to JSON...[/yellow]", total=1)
        with open(SUMMARIES_OUTPUT, "w", encoding="utf-8") as jf:
            json.dump(combined_summaries, jf, indent=2)
        progress.update(save_task, completed=1)
        console.print("[green]✅ Summaries saved successfully![/green]")

        # Save the updated hash cache
        save_hash_cache(hash_cache)

        # Copy the latest summary to the clipboard
        if combined_summaries:
            latest_summary = combined_summaries[-1]["summary"]
            pyperclip.copy(latest_summary)
            console.print("\n[bold green]📋 Latest summary has been copied to your clipboard![/bold green]")
            preview = latest_summary[:200] + "..." if len(latest_summary) > 200 else latest_summary
            console.print(Panel(preview))
        else:
            console.print("[red]⚠️ No summaries generated or found.[/red]")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Summarize a local repo's code in chunked form."
    )
    parser.add_argument(
        "--repo_path",
        type=str,
        default="./my_repo",
        help="Path to the local repository you want to summarize."
    )
    args = parser.parse_args()
    main(args.repo_path)
