__all__ = [
    "CompilationDatabase",
    "CompilationUnit",
]


from collections import defaultdict
from contextlib import contextmanager
from dataclasses import dataclass
from heapq import heappop, heappush
from typing import Any, DefaultDict, Dict, Iterator, List, Optional, Set, Tuple

from beet import Container, DataPack, Generator, TextFile, TextFileBase
from beet.core.utils import extra_field

from .ast import AstRoot
from .diagnostic import DiagnosticCollection


@dataclass
class CompilationUnit:
    """Information associated with the compilation of a specific function file."""

    ast: Optional[AstRoot] = None
    source: Optional[str] = None
    filename: Optional[str] = None
    resource_location: Optional[str] = None
    pack: Optional[DataPack] = None
    priority: int = 0

    diagnostics: DiagnosticCollection = extra_field(init=False)
    perf: Dict[int, float] = extra_field(init=False)

    def __post_init__(self):
        self.diagnostics = DiagnosticCollection(
            filename=self.filename,
            hint=self.resource_location,
        )
        self.perf = {}


class CompilationDatabase(Container[TextFileBase[Any], CompilationUnit]):
    """Compilation database."""

    indices: DefaultDict[Optional[DataPack], Dict[str, TextFileBase[Any]]]
    session: Set[TextFileBase[Any]]
    queue: List[Tuple[int, int, str, int, TextFileBase[Any]]]
    step: int
    current: TextFileBase[Any]
    count: int

    generate: Optional[Generator] = None

    def __init__(self):
        super().__init__()
        self.indices = defaultdict(dict)
        self.session = set()
        self.queue = []
        self.step = -1
        self.current = TextFile()
        self[self.current] = CompilationUnit()
        self.count = 0

        self.generate = None

    @property
    def index(self) -> Dict[str, TextFileBase[Any]]:
        return self.indices[self[self.current].pack]

    def process(
        self,
        key: TextFileBase[Any],
        value: CompilationUnit,
    ) -> CompilationUnit:
        pack_index = self.indices[value.pack]
        for index in [value.filename, value.resource_location]:
            if index:
                pack_index[index] = key
        value.diagnostics.file = key
        return value

    def __delitem__(self, key: TextFileBase[Any]):
        pack_index = self.indices[self[key].pack]
        for index in [self[key].filename, self[key].resource_location]:
            if index:
                del pack_index[index]
        super().__delitem__(key)

    def setup_compilation(self):
        """Setup database for compilation."""
        self.session.clear()
        self.queue.clear()

    def enqueue(self, key: TextFileBase[Any], step: int = -1, priority: int = 0):
        """Enqueue a file and schedule it to be processed with the given step."""
        self.session.add(key)
        self.count += 1
        heappush(
            self.queue,
            (
                step,
                priority,
                self[key].resource_location or "<unknown>",
                self.count,
                key,
            ),
        )

    def dequeue(self) -> Tuple[int, TextFileBase[Any]]:
        """Dequeue the next file that needs to be processed."""
        step, _, _, _, key = heappop(self.queue)
        return step, key

    def process_queue(self) -> Iterator[Tuple[int, TextFileBase[Any]]]:
        """Yield database entries from the queue."""
        while self.queue:
            self.step, self.current = self.dequeue()
            with self.process_context(self.current):
                yield self.step, self.current

    @contextmanager
    def process_context(self, file_instance: TextFileBase[Any]):
        if self.generate is not None:
            compilation_unit = self[file_instance]
            if compilation_unit.pack is not None:
                with self.generate.overlays[compilation_unit.pack.overlay_name].push():
                    yield
                    return
        yield
