import os
import hashlib
import json
from pathlib import Path
from datetime import datetime
from typing import Generator
from dotenv import load_dotenv
load_dotenv()
from openai import OpenAI
from rich.console import Console
from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn, TaskProgressColumn

console = Console()
client = OpenAI()

HASH_FILE = Path(".vectormind/file_hashes.json")


# ---------------------------------------------------------------------------
# Chunking
# ---------------------------------------------------------------------------

def chunk_text(text: str, chunk_size: int, overlap: int) -> list[str]:
    """Split text into overlapping chunks by line boundaries."""
    lines = text.splitlines(keepends=True)
    chunks = []
    current = []
    current_len = 0

    for line in lines:
        current.append(line)
        current_len += len(line.split())

        if current_len >= chunk_size:
            chunks.append("".join(current))
            # Keep overlap lines for next chunk
            overlap_lines = current[-overlap:] if overlap < len(current) else current
            current = list(overlap_lines)
            current_len = sum(len(l.split()) for l in current)

    if current:
        chunks.append("".join(current))

    return [c for c in chunks if c.strip()]


def get_file_hash(path: Path) -> str:
    """MD5 hash of file contents for change detection."""
    return hashlib.md5(path.read_bytes()).hexdigest()


def load_hashes() -> dict:
    if HASH_FILE.exists():
        with open(HASH_FILE) as f:
            return json.load(f)
    return {}


def save_hashes(hashes: dict):
    with open(HASH_FILE, "w") as f:
        json.dump(hashes, f, indent=2)


# ---------------------------------------------------------------------------
# File crawling
# ---------------------------------------------------------------------------

def crawl_files(root: str, config: dict) -> Generator[Path, None, None]:
    """Walk project directory and yield files matching configured extensions."""
    extensions = set(config["extensions"])
    exclude_dirs = set(config["exclude_dirs"])
    root_path = Path(root).resolve()

    for dirpath, dirnames, filenames in os.walk(root_path):
        # Prune excluded directories in-place
        dirnames[:] = [
            d for d in dirnames
            if d not in exclude_dirs and not d.startswith(".")
        ]

        for filename in filenames:
            filepath = Path(dirpath) / filename
            if filepath.suffix in extensions:
                yield filepath


# ---------------------------------------------------------------------------
# Embeddings
# ---------------------------------------------------------------------------

def embed_chunks(chunks: list[str], model: str) -> list[list[float]]:
    """Embed a list of text chunks using OpenAI embeddings API."""
    if not chunks:
        return []

    # OpenAI supports batching up to 2048 inputs
    batch_size = 100
    all_embeddings = []

    for i in range(0, len(chunks), batch_size):
        batch = chunks[i:i + batch_size]
        response = client.embeddings.create(
            input=batch,
            model=model
        )
        all_embeddings.extend([e.embedding for e in response.data])

    return all_embeddings


# ---------------------------------------------------------------------------
# Main indexer
# ---------------------------------------------------------------------------

def index_project(root: str = ".", config: dict = None, incremental: bool = False) -> dict:
    """
    Crawl project, chunk files, embed, and store in ChromaDB.
    Returns stats dict with files, chunks, updated, skipped counts.
    """
    from direction.db import get_collection, upsert_chunks, delete_file_chunks

    chunk_size = config["chunk_size"]
    chunk_overlap = config["chunk_overlap"]
    embedding_model = config["embedding_model"]

    collection = get_collection()
    hashes = load_hashes() if incremental else {}

    files_processed = 0
    files_skipped = 0
    total_chunks = 0
    new_hashes = {}

    files = list(crawl_files(root, config))

    with Progress(
        SpinnerColumn(),
        TextColumn("[progress.description]{task.description}"),
        BarColumn(),
        TaskProgressColumn(),
        console=console
    ) as progress:
        task = progress.add_task("Indexing files...", total=len(files))

        for filepath in files:
            rel_path = str(filepath.relative_to(Path(root).resolve()))
            current_hash = get_file_hash(filepath)
            new_hashes[rel_path] = current_hash

            # Skip unchanged files in incremental mode
            if incremental and hashes.get(rel_path) == current_hash:
                files_skipped += 1
                progress.advance(task)
                continue

            try:
                text = filepath.read_text(encoding="utf-8", errors="ignore")
            except Exception:
                progress.advance(task)
                continue

            if not text.strip():
                progress.advance(task)
                continue

            chunks = chunk_text(text, chunk_size, chunk_overlap)

            if not chunks:
                progress.advance(task)
                continue

            # Embed all chunks for this file
            embeddings = embed_chunks(chunks, embedding_model)

            # Build metadata for each chunk
            metadatas = [
                {
                    "file": rel_path,
                    "chunk_index": i,
                    "extension": filepath.suffix,
                    "indexed_at": datetime.utcnow().isoformat()
                }
                for i in range(len(chunks))
            ]

            # Build unique IDs
            ids = [
                f"{rel_path}::chunk_{i}::{current_hash[:8]}"
                for i in range(len(chunks))
            ]

            # Remove old chunks for this file before upserting
            if incremental:
                delete_file_chunks(collection, rel_path)

            upsert_chunks(
                collection=collection,
                ids=ids,
                embeddings=embeddings,
                documents=chunks,
                metadatas=metadatas
            )

            files_processed += 1
            total_chunks += len(chunks)
            progress.advance(task)

    # Persist hashes
    if incremental:
        hashes.update(new_hashes)
        save_hashes(hashes)
    else:
        save_hashes(new_hashes)

    # Update index metadata
    meta_file = Path(".vectormind/index_meta.json")
    meta = {
        "last_indexed": datetime.utcnow().isoformat(),
        "total_files": files_processed + files_skipped,
        "total_chunks": total_chunks
    }
    with open(meta_file, "w") as f:
        json.dump(meta, f, indent=2)

    return {
        "files": files_processed,
        "chunks": total_chunks,
        "updated": files_processed,
        "skipped": files_skipped
    }