"""Local AI autocomplete engine using llama-cpp-python.

Downloads and runs Qwen2.5-Coder-1.5B (Q8_0 GGUF) for fill-in-the-middle
code completions. Used by the bridge to serve ghost-text suggestions to Studio.
"""

import os
import re
import sys
import tempfile
import threading

from ._env import autocomplete_model_path, models_dir

# Qwen2.5-Coder FIM special tokens
_FIM_PREFIX = "<|fim_prefix|>"
_FIM_SUFFIX = "<|fim_suffix|>"
_FIM_MIDDLE = "<|fim_middle|>"

# Tokens that should stop generation (FIM boundaries + Qwen special tokens)
_FIM_STOP_TOKENS = [
    "<|endoftext|>",
    "<|fim_prefix|>",
    "<|fim_middle|>",
    "<|fim_suffix|>",
    "<|fim_pad|>",
    "<|repo_name|>",
    "<|file_sep|>",
    "<|im_start|>",
    "<|im_end|>",
]

# Regex to strip leaked special tokens from output
_SPECIAL_TOKEN_RE = re.compile(
    r"<\|(?:endoftext|fim_prefix|fim_middle|fim_suffix|fim_pad"
    r"|repo_name|file_sep|im_start|im_end)\|>"
)

MODEL_URL = (
    "https://huggingface.co/ggml-org/Qwen2.5-Coder-1.5B-Q8_0-GGUF"
    "/resolve/main/qwen2.5-coder-1.5b-q8_0.gguf"
)


def model_exists():
    """Check if the autocomplete model has been downloaded."""
    return os.path.isfile(autocomplete_model_path())


def model_path():
    """Return the autocomplete model path."""
    return autocomplete_model_path()


def download_model():
    """Download the autocomplete model from Hugging Face.

    Streams to a temp file first, then renames on success (atomic).
    Returns the path on success, None on failure.
    """
    import requests

    dest = autocomplete_model_path()
    if os.path.isfile(dest):
        return dest

    os.makedirs(models_dir(), exist_ok=True)

    print(f"  Downloading Dash Assist autocomplete model...")
    print(f"  {MODEL_URL}")

    try:
        resp = requests.get(MODEL_URL, stream=True, timeout=30)
        resp.raise_for_status()

        total = int(resp.headers.get("content-length", 0))
        downloaded = 0

        # Write to temp file in the same directory (for atomic rename)
        fd, tmp_path = tempfile.mkstemp(dir=models_dir(), suffix=".tmp")
        try:
            with os.fdopen(fd, "wb") as f:
                for chunk in resp.iter_content(chunk_size=1024 * 1024):  # 1MB chunks
                    f.write(chunk)
                    downloaded += len(chunk)
                    if total > 0:
                        pct = downloaded * 100 // total
                        mb_done = downloaded / (1024 * 1024)
                        mb_total = total / (1024 * 1024)
                        bar = "#" * (pct // 5) + "-" * (20 - pct // 5)
                        sys.stdout.write(
                            f"\r  [{bar}] {pct}% ({mb_done:.0f}/{mb_total:.0f} MB)"
                        )
                        sys.stdout.flush()

            print()  # newline after progress bar
            os.rename(tmp_path, dest)
            print(f"  Saved to {dest}")
            return dest

        except Exception:
            # Clean up temp file on failure
            try:
                os.unlink(tmp_path)
            except OSError:
                pass
            raise

    except Exception as e:
        print(f"\n  Warning: Failed to download autocomplete model: {e}")
        return None


class AutocompleteEngine:
    """Manages a local LLM for code completions.

    Thread-safe: the model is guarded by a lock and generation can be
    cancelled by bumping the request ID.
    """

    def __init__(self):
        self._model = None
        self._lock = threading.Lock()
        self._loading = False
        self._loaded = False
        self._request_id = 0

    def load_in_background(self):
        """Start loading the model in a daemon thread."""
        if self._loaded or self._loading:
            return
        self._loading = True
        t = threading.Thread(target=self._load, daemon=True)
        t.start()

    def _load(self):
        """Load the GGUF model. Called from background thread."""
        try:
            from llama_cpp import Llama

            path = autocomplete_model_path()
            if not os.path.isfile(path):
                print("[Dash Assist] Autocomplete model not found, skipping.")
                return

            print("[Dash Assist] Loading autocomplete model...")
            with self._lock:
                self._model = Llama(
                    model_path=path,
                    n_ctx=4096,
                    n_gpu_layers=-1,
                    verbose=False,
                )
            self._loaded = True
            print("[Dash Assist] Autocomplete ready.")
        except Exception as e:
            print(f"[Dash Assist] Failed to load autocomplete model: {e}")
        finally:
            self._loading = False

    @property
    def ready(self):
        return self._loaded and self._model is not None

    async def complete(self, prefix, suffix, request_id):
        """Generate a completion asynchronously.

        Returns the completion text, or None if cancelled/unavailable.
        """
        if not self.ready:
            return None

        # Accept this request — any older in-flight request is now stale
        self._request_id = request_id

        import asyncio
        return await asyncio.to_thread(self._generate, prefix, suffix, request_id)

    def _generate(self, prefix, suffix, request_id):
        """Synchronous generation. Checks request_id for cancellation."""
        if self._request_id != request_id:
            return None  # A newer request arrived

        # Don't block waiting for a stale generation to finish
        acquired = self._lock.acquire(timeout=0.05)
        if not acquired:
            return None

        try:
            if self._request_id != request_id:
                return None

            # Manually construct the FIM prompt — llama-cpp-python's suffix=
            # parameter is broken (FIM token IDs hardcoded to 0).
            fim_prompt = f"{_FIM_PREFIX}{prefix}{_FIM_SUFFIX}{suffix}{_FIM_MIDDLE}"

            result = self._model.create_completion(
                prompt=fim_prompt,
                max_tokens=128,
                temperature=0.01,
                top_k=40,
                top_p=0.90,
                min_p=0.05,
                stop=_FIM_STOP_TOKENS + ["\n\n", "\ndef ", "\nclass "],
            )
            # Check cancellation after generation
            if self._request_id != request_id:
                return None

            text = result["choices"][0]["text"] if result.get("choices") else None
            if not text:
                return None

            return self._postprocess(text, prefix, suffix)
        except Exception:
            return None
        finally:
            self._lock.release()

    @staticmethod
    def _postprocess(text, prefix, suffix):
        """Clean up the raw completion text."""
        # Strip any leaked FIM / special tokens
        text = _SPECIAL_TOKEN_RE.sub("", text)

        # Strip trailing whitespace
        text = text.rstrip()

        # Reject whitespace-only completions
        if not text.strip():
            return None

        # Detect prefix echo: if the completion starts with the tail of the
        # prefix, strip the overlapping portion.
        if prefix:
            tail = prefix[-len(text):] if len(text) <= len(prefix) else prefix
            for i in range(min(len(tail), len(text)), 0, -1):
                if text[:i] == tail[-i:]:
                    text = text[i:]
                    break

        # Detect suffix overlap: if the completion ends with the start of the
        # suffix, trim it.
        if suffix:
            head = suffix[:len(text)] if len(text) <= len(suffix) else suffix
            for i in range(min(len(head), len(text)), 0, -1):
                if text[-i:] == head[:i]:
                    text = text[:-i]
                    break

        # Final empty check after trimming
        if not text.strip():
            return None

        return text
