"""
Main entry point for the deep agent search application.

Demonstrates:
- Creating a virtual file system.
- Adding a virtual file with sample data.
- Performing a search query using a simple deep agent.
"""

from __future__ import annotations

import sys
from typing import List

import numpy as np
import torch
from sklearn.feature_extraction.text import TfidfVectorizer

from virtual_file_system import VirtualFileSystem, VirtualFile


class SearchAgent:
    """
    Simple search agent that ranks lines from a virtual file based on
    cosine similarity between TF-IDF vectors of the query and the lines.
    """

    def __init__(self, vfs: VirtualFileSystem):
        self.vfs = vfs

    def search(self, file_name: str, query: str, top_k: int = 5) -> List[str]:
        """
        Search for the most relevant lines in the specified virtual file.

        Parameters
        ----------
        file_name : str
            Name of the virtual file to search.
        query : str
            Search query string.
        top_k : int, optional
            Number of top results to return.

        Returns
        -------
        List[str]
            List of the most relevant lines.
        """
        vf = self.vfs.get_file(file_name)
        data = vf.read().decode("utf-8")
        lines = [line.strip() for line in data.splitlines() if line.strip()]
        if not lines:
            return []

        # Vectorize lines and query
        vectorizer = TfidfVectorizer()
        doc_vectors = vectorizer.fit_transform(lines).toarray()
        query_vec = vectorizer.transform([query]).toarray()

        # Convert to torch tensors for similarity calculation
        doc_tensors = torch.tensor(doc_vectors, dtype=torch.float32)
        query_tensor = torch.tensor(query_vec, dtype=torch.float32)

        # Normalize vectors
        doc_norm = doc_tensors / doc_tensors.norm(dim=1, keepdim=True)
        query_norm = query_tensor / query_tensor.norm()

        # Cosine similarity
        similarities = torch.matmul(doc_norm, query_norm.t()).squeeze()

        # Get top_k indices
        top_indices = similarities.topk(top_k).indices.tolist()
        return [lines[i] for i in top_indices]


def main(argv: List[str]) -> None:
    """
    Example usage of the virtual file system and search agent.

    Creates a virtual file with sample text and performs a search query.
    """
    vfs = VirtualFileSystem()

    # Sample data: a small collection of sentences
    sample_text = """\
Deep learning has revolutionized many fields.
Neural networks can approximate complex functions.
PyTorch provides dynamic computation graphs.
Scikit-learn offers a wide range of machine learning tools.
Numpy is essential for numerical operations.
"""

    # Create a virtual file
    vf = vfs.create_file("sample.txt", sample_text.encode("utf-8"))

    # Instantiate the search agent
    agent = SearchAgent(vfs)

    # Perform a search query
    query = "neural networks"
    results = agent.search("sample.txt", query, top_k=3)

    print(f"Search results for query: '{query}'")
    for idx, line in enumerate(results, 1):
        print(f"{idx}. {line}")

    # Demonstrate unload
    vf.unload()
    try:
        vf.read()
    except RuntimeError as e:
        print(f"After unload: {e}")


if __name__ == "__main__":
    main(sys.argv[1:])