import os
import hashlib
import json
from datetime import datetime
import time


class DocMetaData:
    """
    Class that stores document meta data. Using this so we know that the MetaData is always
    initialized and can avoid a lot of checking whether keys exist or not
    """

    def __init__(self, mod_time, hash, last_updated, chunk_ids):
        self.mod_time = mod_time
        self.hash = hash
        self.last_updated = last_updated
        self.chunk_ids = chunk_ids

    def __eq__(self, other):
        if isinstance(other, DocMetaData):
            return (
                self.mod_time == other.mod_time
                and self.hash == other.hash
                and self.last_updated == other.last_updated
                and self.chunk_ids == other.chunk_ids
            )
        return NotImplemented

    def update_chunks(self, chunks):
        self.chunk_ids = chunks

    def to_dict(self):
        return {
            "mod_time": self.mod_time,
            "hash": self.hash,
            "last_updated": self.last_updated,
            "chunk_ids": self.chunk_ids,
        }

    @classmethod
    def from_file(cls, file_path):
        return cls(
            os.path.getmtime(file_path), calculate_file_hash(file_path), time.time(), []
        )

    @classmethod
    def from_dict(cls, input_dict):
        return cls(
            input_dict["mod_time"],
            input_dict["hash"],
            input_dict["last_updated"],
            input_dict["chunk_ids"],
        )


class DocumentTracker:
    """
    Tracks document changes using file hashes, modification times, and chunk IDs
    """

    def __init__(self, tracking_file):
        self.tracking_file = tracking_file
        self.doc_info = self._load_tracking_data()

    def _load_tracking_data(self):
        """Load existing tracking data if available"""
        doc_info = dict()
        if os.path.exists(self.tracking_file):
            with open(self.tracking_file, "r") as f:
                serialized = json.load(f)
                for k, v in serialized.items():
                    doc_info[k] = DocMetaData.from_dict(v)

        return doc_info

    def _save_tracking_data(self):
        """Save tracking data to file"""
        output = dict()
        for k, v in self.doc_info.items():
            output[k] = v.to_dict()
        with open(self.tracking_file, "w") as f:
            json.dump(output, f, indent=2)

    def get_changed_files(self, directory, file_extension=".txt"):
        """
        Detect new, modified, and deleted files
        Returns: dict with 'new', 'modified', and 'deleted' lists
        """
        current_file_mod_times = {}
        for root, _, files in os.walk(directory):
            for file in files:
                if file.endswith(file_extension):
                    file_path = os.path.join(root, file)
                    mod_time = os.path.getmtime(file_path)
                    current_file_mod_times[file_path] = mod_time

        new_files = []
        modified_files = []

        # Check for new or modified files
        for file_path, mod_time in current_file_mod_times.items():
            if file_path not in self.doc_info:
                new_files.append(file_path)
            elif mod_time > self.doc_info[file_path].mod_time:
                # Check if content actually changed using hash
                current_hash = calculate_file_hash(file_path)
                if current_hash != self.doc_info[file_path].hash:
                    modified_files.append(file_path)
                    self.doc_info[file_path].hash = current_hash
                self.doc_info[file_path].mod_time = mod_time

        # Check for deleted files
        deleted_files = [f for f in self.doc_info if f not in current_file_mod_times]

        # Update tracking information for new files
        for file_path in new_files:
            self.doc_info[file_path] = DocMetaData(
                current_file_mod_times[file_path],
                calculate_file_hash(file_path),
                time.time(),
                [],  # Will store IDs of chunks derived from this document
            )

        # Update last_updated for modified files
        for file_path in modified_files:
            self.doc_info[file_path].last_updated = datetime.now().isoformat()

        self._save_tracking_data()

        return {"new": new_files, "modified": modified_files, "deleted": deleted_files}

    def update_chunk_mappings(self, file_path, chunk_ids):
        """Store the chunk IDs associated with a document"""
        if file_path not in self.doc_info:
            self.doc_info[file_path] = DocMetaData.from_file(file_path)

        self.doc_info[file_path].chunk_ids = chunk_ids
        self._save_tracking_data()

    def get_chunks_to_delete(self, deleted_files):
        """Get all chunk IDs associated with deleted files"""
        chunks_to_delete = []
        for file_path in deleted_files:
            if file_path in self.doc_info and self.doc_info[file_path].chunk_ids:
                chunks_to_delete.extend(self.doc_info[file_path].chunk_ids)
                # Remove the file from tracking after processing
                del self.doc_info[file_path]

        self._save_tracking_data()
        return chunks_to_delete

    def get_chunks_for_modified_files(self, modified_files):
        """Get chunk IDs for modified files that need to be deleted before re-indexing"""
        chunks_to_delete = []
        for file_path in modified_files:
            if file_path in self.doc_info and self.doc_info[file_path].chunk_ids:
                chunks_to_delete.extend(self.doc_info[file_path].chunk_ids)
                # Clear the chunk IDs (will be updated with new ones)
                self.doc_info[file_path].chunk_ids = []

        self._save_tracking_data()
        return chunks_to_delete


def calculate_file_hash(file_path):
    """Calculate MD5 hash of file contents"""
    hasher = hashlib.md5()
    with open(file_path, "rb") as f:
        buf = f.read(65536)
        while len(buf) > 0:
            hasher.update(buf)
            buf = f.read(65536)
    return hasher.hexdigest()
