"""
ChunkHive CLI - Embedding Commands
===================================

New commands for semantic embeddings and vector search.
"""

import typer
from pathlib import Path
from typing import Optional
import json

embed_app = typer.Typer(help="Semantic embedding commands for code search")


# =============================================================================
# HELPER: Load chunks from JSONL file
# =============================================================================

def _load_chunks(chunks_file: str, require_embeddings: bool = False):
    """
    Load chunks from a JSONL file.
    
    Args:
        chunks_file: Path to the JSONL file
        require_embeddings: If True, only load chunks that have embeddings
        
    Returns:
        Tuple of (chunks, embeddings_list)
        embeddings_list is empty if require_embeddings=False
    """
    from chunkhive.chunkers.schema import SemanticChunk
    
    chunks = []
    embeddings_list = []
    skipped = 0
    
    with open(chunks_file, 'r', encoding='utf-8') as f:
        for line in f:
            if not line.strip():
                continue
            chunk_data = json.loads(line)
            chunk = SemanticChunk.from_dict(chunk_data)
            
            if require_embeddings:
                if chunk.embedding:
                    chunks.append(chunk)
                    embeddings_list.append(chunk.embedding)
                else:
                    skipped += 1
            else:
                chunks.append(chunk)
    
    return chunks, embeddings_list, skipped


# =============================================================================
# COMMAND: generate
# =============================================================================

@embed_app.command("generate")
def generate_embeddings(
    chunks_file: str = typer.Argument(..., help="Path to chunks.jsonl file"),
    output: Optional[str] = typer.Option(None, "-o", "--output", help="Output file path"),
    model: str = typer.Option("minilm", "--model", help="Model to use (minilm, codebert, unixcoder)"),
    batch_size: int = typer.Option(32, "--batch-size", help="Batch size for processing"),
    overwrite: bool = typer.Option(False, "--overwrite", help="Re-generate even if embeddings exist"),
):
    """
    Generate semantic embeddings for code chunks.

    Example:
        chunkhive embed generate chunks.jsonl -o chunks_embedded.jsonl
        chunkhive embed generate chunks.jsonl --model codebert
    """
    from rich.console import Console
    console = Console()

    try:
        from chunkhive.embeddings import SemanticEmbeddingEngine
        import numpy as np

        console.print(f"[cyan]Loading chunks from {chunks_file}...[/cyan]")
        chunks, existing_embeddings, _ = _load_chunks(chunks_file, require_embeddings=False)
        console.print(f"[green]✓[/green] Loaded {len(chunks)} chunks")

        # Check if embeddings already exist
        already_embedded = sum(1 for c in chunks if c.embedding)
        
        if already_embedded > 0 and not overwrite:
            console.print(
                f"[yellow]⚠[/yellow] {already_embedded}/{len(chunks)} chunks already have embeddings.\n"
                f"  Using existing embeddings. Use --overwrite to re-generate."
            )
            # Only embed chunks that don't have embeddings yet
            chunks_to_embed = [c for c in chunks if not c.embedding]
            if chunks_to_embed:
                console.print(f"[cyan]Generating embeddings for {len(chunks_to_embed)} new chunks...[/cyan]")
                engine = SemanticEmbeddingEngine(model_name=model)
                new_embeddings = engine.generate_embeddings(chunks_to_embed, batch_size=batch_size)
                engine.attach_embeddings_to_chunks(chunks_to_embed, new_embeddings)
                console.print(f"[green]✓[/green] Generated {len(new_embeddings)} new embeddings")
            else:
                console.print(f"[green]✓[/green] All chunks already embedded, nothing to do!")
        else:
            # Generate for all chunks
            console.print(f"[cyan]Generating embeddings using {model}...[/cyan]")
            engine = SemanticEmbeddingEngine(model_name=model)
            embeddings = engine.generate_embeddings(chunks, batch_size=batch_size)
            engine.attach_embeddings_to_chunks(chunks, embeddings)
            console.print(f"[green]✓[/green] Generated {len(embeddings)} embeddings")

        # Save output
        output_file = output or chunks_file.replace('.jsonl', '_embedded.jsonl')
        output_path = Path(output_file)
        output_path.parent.mkdir(parents=True, exist_ok=True)

        with open(output_file, 'w', encoding='utf-8') as f:
            for chunk in chunks:
                f.write(json.dumps(chunk.to_dict()) + '\n')

        console.print(f"[green]✓[/green] Saved to {output_file}")
        
        # Show which model was used
        model_used = chunks[0].embedding_model if chunks and chunks[0].embedding_model else model
        console.print(f"[cyan]Embedding model: {model_used}[/cyan]")

    except ImportError as e:
        console.print(f"[red]✗[/red] Missing dependencies: {e}")
        console.print("[yellow]Install with: pip install sentence-transformers faiss-cpu[/yellow]")
    except Exception as e:
        console.print(f"[red]✗[/red] Error: {e}")
        import traceback
        traceback.print_exc()


# =============================================================================
# COMMAND: search
# =============================================================================

@embed_app.command("search")
def search_chunks(
    query: str = typer.Argument(..., help="Search query"),
    chunks_file: str = typer.Option(..., "--chunks", help="Path to chunks.jsonl with embeddings"),
    model: str = typer.Option(None, "--model", help="Override model for query (auto-detected from file)"),
    k: int = typer.Option(5, "-k", "--top-k", help="Number of results"),
    show_code: bool = typer.Option(True, "--show-code/--no-code", help="Show code preview"),
):
    """
    Search code chunks semantically.

    Example:
        chunkhive embed search "authentication function" --chunks embedded.jsonl
        chunkhive embed search "hash password" --chunks embedded.jsonl -k 10
    """
    from rich.console import Console
    from rich.table import Table
    from rich.syntax import Syntax
    console = Console()

    try:
        from chunkhive.embeddings import SemanticEmbeddingEngine, VectorSearchIndex
        import numpy as np

        console.print(f"[cyan]Loading chunks from {chunks_file}...[/cyan]")
        chunks, embeddings_list, skipped = _load_chunks(chunks_file, require_embeddings=True)

        if not chunks:
            console.print("[red]✗[/red] No chunks with embeddings found!")
            console.print("[yellow]Tip: Run 'chunkhive embed generate' first[/yellow]")
            return

        if skipped > 0:
            console.print(f"[yellow]⚠[/yellow] Skipped {skipped} chunks without embeddings")

        # Auto-detect model from file if not specified
        detected_model = None
        for chunk in chunks:
            if chunk.embedding_model:
                # Map model_id back to model name
                model_id = chunk.embedding_model
                if "codebert" in model_id:
                    detected_model = "codebert"
                elif "unixcoder" in model_id:
                    detected_model = "unixcoder"
                elif "codesage" in model_id:
                    detected_model = "codesage"
                else:
                    detected_model = "minilm"
                break

        search_model = model or detected_model or "minilm"
        
        if detected_model and model and detected_model != model:
            console.print(
                f"[yellow]⚠ Warning:[/yellow] Embeddings were created with '{detected_model}' "
                f"but you specified '{model}'. Using '{detected_model}' for consistency."
            )
            search_model = detected_model

        console.print(f"[green]✓[/green] Loaded {len(chunks)} chunks with embeddings")
        console.print(f"[dim]Using model: {search_model}[/dim]")

        # Build FAISS index from existing embeddings
        console.print(f"[cyan]Building search index...[/cyan]")
        embeddings = np.array(embeddings_list, dtype='float32')
        engine = SemanticEmbeddingEngine(model_name=search_model)
        
        # Use existing embeddings, don't re-generate!
        index = VectorSearchIndex(dimension=embeddings.shape[1])
        index.add_embeddings(embeddings, chunks)

        # Search
        console.print(f"[cyan]Searching for: '{query}'[/cyan]\n")
        results = index.search(query, engine, k=k)

        if not results:
            console.print("[yellow]No results found![/yellow]")
            return

        # Display results table
        table = Table(title=f"Top {len(results)} Results for: '{query}'")
        table.add_column("#", style="cyan", width=3)
        table.add_column("File", style="green")
        table.add_column("Name", style="yellow")
        table.add_column("Type", style="blue")
        table.add_column("Distance", style="magenta")

        for i, (chunk, distance) in enumerate(results, 1):
            # Similarity score: lower distance = more similar
            similarity = f"{distance:.4f}"
            table.add_row(
                str(i),
                Path(chunk.file_path).name,
                chunk.ast.name or "N/A",
                chunk.chunk_type,
                similarity
            )

        console.print(table)
        console.print("[dim]Lower distance = more similar[/dim]\n")

        # Show code preview for top result
        if show_code and results:
            top_chunk, top_distance = results[0]
            console.print(f"[bold]📄 Top Result: {top_chunk.ast.name or 'N/A'}[/bold]")
            console.print(f"[dim]File: {top_chunk.file_path}[/dim]")
            preview = top_chunk.get_code_preview(max_lines=15)
            console.print(f"[dim]{preview}[/dim]")

    except ImportError as e:
        console.print(f"[red]✗[/red] Missing dependencies: {e}")
        console.print("[yellow]Install with: pip install sentence-transformers faiss-cpu[/yellow]")
    except Exception as e:
        console.print(f"[red]✗[/red] Error: {e}")
        import traceback
        traceback.print_exc()


# =============================================================================
# COMMAND: index
# =============================================================================

@embed_app.command("index")
def build_search_index(
    chunks_file: str = typer.Argument(..., help="Path to chunks.jsonl WITH embeddings"),
    output: str = typer.Option("search_index", "-o", "--output", help="Output index path"),
):
    """
    Build a persistent FAISS search index from pre-embedded chunks.
    
    IMPORTANT: Input file must already have embeddings!
    Run 'chunkhive embed generate' first.

    Example:
        chunkhive embed generate chunks.jsonl -o embedded.jsonl
        chunkhive embed index embedded.jsonl -o my_index
    """
    from rich.console import Console
    console = Console()

    try:
        from chunkhive.embeddings import VectorSearchIndex
        import numpy as np

        console.print(f"[cyan]Loading chunks with embeddings from {chunks_file}...[/cyan]")
        chunks, embeddings_list, skipped = _load_chunks(chunks_file, require_embeddings=True)

        if not chunks:
            console.print("[red]✗[/red] No chunks with embeddings found!")
            console.print("[yellow]Tip: Run 'chunkhive embed generate' first to create embeddings[/yellow]")
            return

        if skipped > 0:
            console.print(f"[yellow]⚠[/yellow] Skipped {skipped} chunks without embeddings")

        console.print(f"[green]✓[/green] Loaded {len(chunks)} chunks with embeddings")

        # Detect dimension from existing embeddings
        embeddings = np.array(embeddings_list, dtype='float32')
        dimension = embeddings.shape[1]
        
        # Detect model
        model_used = "unknown"
        for chunk in chunks:
            if chunk.embedding_model:
                model_used = chunk.embedding_model
                break
        
        console.print(f"[dim]Embedding dimension: {dimension}[/dim]")
        console.print(f"[dim]Model: {model_used}[/dim]")

        # Build FAISS index from EXISTING embeddings (NO re-generation!)
        console.print(f"[cyan]Building FAISS index...[/cyan]")
        index = VectorSearchIndex(dimension=dimension)
        index.add_embeddings(embeddings, chunks)

        # Save index to disk
        output_path = Path(output)
        output_path.parent.mkdir(parents=True, exist_ok=True)
        index.save(output_path)

        console.print(f"[green]✓[/green] Index saved!")
        console.print(f"[dim]  {output_path}.faiss  ← FAISS vector index[/dim]")
        console.print(f"[dim]  {output_path}.chunks ← Chunk metadata[/dim]")
        console.print(f"[cyan]Total vectors: {index.index.ntotal}[/cyan]")
        console.print(f"[cyan]Dimension: {dimension}[/cyan]")
        console.print(f"\n[green]To search using this index:[/green]")
        console.print(f"[dim]chunkhive embed search-index 'your query' --index {output}[/dim]")

    except ImportError as e:
        console.print(f"[red]✗[/red] Missing dependencies: {e}")
    except Exception as e:
        console.print(f"[red]✗[/red] Error: {e}")
        import traceback
        traceback.print_exc()


# =============================================================================
# COMMAND: search-index (search using persistent index)
# =============================================================================

@embed_app.command("search-index")
def search_with_index(
    query: str = typer.Argument(..., help="Search query"),
    index_path: str = typer.Option(..., "--index", help="Path to saved index (without extension)"),
    model: str = typer.Option(None, "--model", help="Model to use for query embedding"),
    k: int = typer.Option(5, "-k", "--top-k", help="Number of results"),
):
    """
    Search using a pre-built persistent FAISS index (FAST!).
    
    This is faster than 'search' for repeated searches because
    the index is pre-built and loaded from disk.

    Example:
        chunkhive embed search-index "auth function" --index data/embeddings/my_index
    """
    from rich.console import Console
    from rich.table import Table
    console = Console()

    try:
        from chunkhive.embeddings import VectorSearchIndex, SemanticEmbeddingEngine

        console.print(f"[cyan]Loading index from {index_path}...[/cyan]")
        
        # Check index files exist
        faiss_file = Path(index_path).with_suffix('.faiss')
        chunks_file = Path(index_path).with_suffix('.chunks')
        
        if not faiss_file.exists():
            console.print(f"[red]✗[/red] Index file not found: {faiss_file}")
            console.print("[yellow]Tip: Run 'chunkhive embed index' first to build the index[/yellow]")
            return
        
        # Load index
        index_path_obj = Path(index_path)
        
        # We need to detect dimension from the loaded index
        import faiss as faiss_lib  # type: ignore
        faiss_index = faiss_lib.read_index(str(faiss_file))
        dimension = faiss_index.d
        
        # Load using our class
        loaded_index = VectorSearchIndex.load(index_path_obj, dimension=dimension)
        
        console.print(f"[green]✓[/green] Loaded index with {loaded_index.index.ntotal} vectors")
        
        # Detect model from chunks
        search_model = model or "minilm"
        for chunk in loaded_index.chunks[:5]:
            if chunk.embedding_model:
                model_id = chunk.embedding_model
                if "codebert" in model_id:
                    search_model = "codebert"
                elif "unixcoder" in model_id:
                    search_model = "unixcoder"
                elif "codesage" in model_id:
                    search_model = "codesage"
                else:
                    search_model = "minilm"
                break
        
        console.print(f"[dim]Using model: {search_model}[/dim]")
        
        # Load engine for query embedding
        engine = SemanticEmbeddingEngine(model_name=search_model)
        
        # Search
        console.print(f"[cyan]Searching for: '{query}'[/cyan]\n")
        results = loaded_index.search(query, engine, k=k)

        if not results:
            console.print("[yellow]No results found![/yellow]")
            return

        # Display results
        table = Table(title=f"Top {len(results)} Results")
        table.add_column("#", style="cyan", width=3)
        table.add_column("File", style="green")
        table.add_column("Name", style="yellow")
        table.add_column("Type", style="blue")
        table.add_column("Distance", style="magenta")

        for i, (chunk, distance) in enumerate(results, 1):
            table.add_row(
                str(i),
                Path(chunk.file_path).name,
                chunk.ast.name or "N/A",
                chunk.chunk_type,
                f"{distance:.4f}"
            )

        console.print(table)
        
        if results:
            top_chunk, _ = results[0]
            console.print(f"\n[bold]Top Result:[/bold] {top_chunk.ast.name or 'N/A'}")
            console.print(f"[dim]{top_chunk.get_code_preview(max_lines=10)}[/dim]")

    except ImportError as e:
        console.print(f"[red]✗[/red] Missing dependencies: {e}")
    except Exception as e:
        console.print(f"[red]✗[/red] Error: {e}")
        import traceback
        traceback.print_exc()


# =============================================================================
# COMMAND: models
# =============================================================================

@embed_app.command("models")
def list_models():
    """List available embedding models."""
    from rich.console import Console
    from rich.table import Table
    console = Console()

    try:
        from chunkhive.embeddings import AVAILABLE_MODELS

        table = Table(title="Available Embedding Models")
        table.add_column("Name", style="cyan")
        table.add_column("Model ID", style="green")
        table.add_column("Dimension", style="yellow")
        table.add_column("Best For", style="blue")
        table.add_column("Speed", style="magenta")

        speeds = {
            "minilm": "⚡⚡⚡ Fast",
            "codebert": "⚡⚡ Medium",
            "unixcoder": "⚡⚡ Medium",
            "codesage": "⚡ Slow",
        }

        for name, config in AVAILABLE_MODELS.items():
            table.add_row(
                name,
                config.model_id,
                str(config.dimension),
                config.best_for,
                speeds.get(name, "⚡⚡ Medium")
            )

        console.print(table)
        console.print("\n[dim]Usage: chunkhive embed generate chunks.jsonl --model <name>[/dim]")
        console.print("[dim]Tip: Use 'minilm' for speed, 'codebert' for best quality[/dim]")

    except ImportError:
        console.print("[red]✗[/red] Embeddings module not available")
        console.print("[yellow]Install with: pip install sentence-transformers[/yellow]")