from __future__ import annotations

import os
import subprocess
import tempfile
from functools import partial
from textwrap import dedent
from typing import Literal

from ramalama.arg_types import RagArgsType
from ramalama.chat import ChatOperationalArgs
from ramalama.common import ensure_image, perror, set_accel_env_vars
from ramalama.compat import StrEnum
from ramalama.config import ActiveConfig, Config
from ramalama.engine import BuildEngine, Engine, is_healthy, stop_container, wait_for_healthy
from ramalama.path_utils import get_container_mount_path
from ramalama.transports.base import Transport
from ramalama.transports.oci.oci import OCI

INPUT_DIR = "/docs"


class VectorDBEngine(Engine):
    """Generate a RAG vector database from source content."""

    def __init__(self, args):
        super().__init__(args)
        self.add_user()

    def add_input(self, path: str) -> None:
        if os.path.exists(path):
            fpath = os.path.realpath(path)
            input_name = os.path.basename(fpath)
            self.add_volume(fpath, f"{INPUT_DIR}/{input_name}")
        else:
            raise ValueError(f"{path} does not exist")

    def add_user(self):
        if self.use_docker:
            # Run the command in the container with the same uid as the current user.
            # Otherwise the files written to the tempdir will be owned by root and
            # they will not be readable by the subsequent "docker build".
            # Note: geteuid() doesn't exist on Windows, but this is only needed on Unix
            if hasattr(os, 'geteuid'):
                self.add_args("--user", str(os.geteuid()))


class Rag:
    target: str = ""
    oci: bool = False

    def __init__(self, target: str):
        self.target = target
        self.oci = self.target.startswith("oci://") or (
            not self.target.startswith("file://") and self.target[0] not in {".", "/"}
        )
        if self.oci and not self.target.islower():
            raise ValueError(f"invalid reference format: repository name '{self.target}' must be lowercase")
        set_accel_env_vars()

    def build(self, source: str, args):
        perror(f"Building {self.target} ...")
        contextdir = os.path.dirname(source)
        src = os.path.basename(source)
        engine = BuildEngine(args)
        return engine.build_containerfile(
            dedent(f"""
                FROM scratch
                COPY {src} /vector.db
                """),
            contextdir,
            tag=self.target,
        )

    def generate(self, args, cmd):
        args.nocapdrop = True
        if not args.container:
            raise KeyError("rag command requires a container. Can not be run with --nocontainer option.")
        if not args.engine or args.engine == "":
            raise KeyError("rag command requires a container. Can not be run without a container engine.")
        engine = VectorDBEngine(args)

        for path in args.PATHS:
            engine.add_input(path)

        if self.oci:
            ragdb = tempfile.TemporaryDirectory(prefix="RamaLama_rag_")
            dbdir = os.path.join(ragdb.name, "vectordb")
            os.makedirs(dbdir)
        else:
            dbdir = self.target

        engine.add_volume(dbdir, "/output", opts="rw")
        if not args.dryrun:
            config = ActiveConfig()
            should_pull = config.pull in ["always", "missing", "newer"]
            args.image = ensure_image(
                args.engine, args.image, should_pull=should_pull, quiet=getattr(args, "quiet", False)
            )
        engine.add_args(args.image)
        engine.add_args(*cmd)
        try:
            if args.dryrun:
                engine.dryrun()
            else:
                engine.run()
            if self.oci:
                print(self.build(dbdir, args))
        except subprocess.CalledProcessError as e:
            raise e
        finally:
            if self.oci:
                ragdb.cleanup()


def rag_image(config: Config) -> str:
    return config.default_rag_image


class RagSource(StrEnum):
    DB = "db"
    IMAGE = "image"


class RagEngine(Engine):
    """Engine for executing a RAG proxy"""

    sourcetype: RagSource

    def __init__(self, args, /, *, sourcetype: RagSource):
        self.sourcetype = sourcetype
        super().__init__(args)
        self.add_rag()

    def add_labels(self):
        super().add_labels()
        rag = getattr(self.args, "rag", None)
        if rag is not None:
            self.add_label(f"ai.ramalama.rag.{self.sourcetype}={rag}")

    def add_rag(self):
        rag = getattr(self.args, "rag", None)
        if rag is None:
            return
        if self.sourcetype is RagSource.DB:
            # Convert to container-friendly path format (handles Windows path conversion)
            rag_path = get_container_mount_path(rag)
            # Read-write is the default behaviour for bind mounts in both Docker and Podman
            self.add_args(f"--mount=type=bind,source={rag_path},destination=/rag/vector.db{self.relabel()}")
        else:
            # Image mounts default to read-only in Podman, so we need rw=true for write access
            # Docker does not support type=image mounts
            if self.use_podman:
                self.add_args(f"--mount=type=image,source={rag},destination=/rag,rw=true")
            else:
                # Docker falls back to using volumes or other mechanisms
                self.add_args(f"--mount=type=image,source={rag},destination=/rag")


class RagTransport(OCI):
    """Run a RAG proxy, dispatching to a backend LLM"""

    type: str = "Model+RAG"

    def format_model(self, model: str) -> str:
        if getattr(self, "kind", None) is RagSource.DB:
            return model
        if "/" not in model:
            model = f"localhost/{model}"
        return super().format_model(model)

    def __init__(self, imodel: Transport, cmd: list[str], args: RagArgsType):
        self.kind = RagSource.DB if os.path.exists(args.rag) else RagSource.IMAGE
        super().__init__(args.rag, args.store, args.engine)
        args.rag = self.model
        self.imodel = imodel
        self.model_cmd = cmd

    def exists(self) -> bool:
        if self.kind is RagSource.DB:
            return os.path.exists(self.model)
        return super().exists()

    def new_engine(self, args: RagArgsType) -> "RagEngine":
        return RagEngine(args, sourcetype=self.kind)

    def setup_mounts(self, args: RagArgsType) -> None:
        pass

    def chat_operational_args(self, args: RagArgsType) -> ChatOperationalArgs:
        return ChatOperationalArgs(name=args.model_args.name)

    def _handle_container_chat(self, args: RagArgsType, server_process: int) -> Literal[0]:
        # Clear args.rag so RamaLamaShell doesn't treat it as local data for RAG context
        args.rag = None  # type: ignore[assignment]
        return super()._handle_container_chat(args, server_process)

    def serve(self, args: RagArgsType, cmd: list[str]):
        args.model_args.name = self.imodel.get_container_name(args.model_args)
        process = self.imodel.serve_nonblocking(args.model_args, self.model_cmd)
        if not args.dryrun:
            if process and process.wait() != 0:
                raise subprocess.CalledProcessError(
                    process.returncode,
                    " ".join(self.model_cmd),
                )
        try:
            self.execute_command(cmd, args)
        finally:
            if getattr(args.model_args, "name", None):
                args.model_args.ignore = True
                stop_container(args.model_args, args.model_args.name, remove=True)

    def run(self, args: RagArgsType, cmd: list[str]):
        args.model_args.name = self.imodel.get_container_name(args.model_args)
        process = self.imodel.serve_nonblocking(args.model_args, self.model_cmd)
        rag_process = self.serve_nonblocking(args, cmd)

        if args.dryrun:
            return

        if process and process.wait() != 0:
            raise subprocess.CalledProcessError(
                process.returncode,
                " ".join(self.model_cmd),
            )
        if rag_process and rag_process.wait() != 0:
            raise subprocess.CalledProcessError(
                rag_process.returncode,
                " ".join(cmd),
            )
        return self._connect_and_chat(args, rag_process)

    def wait_for_healthy(self, args: RagArgsType) -> None:
        self.imodel.wait_for_healthy(args.model_args)
        wait_for_healthy(args, partial(is_healthy, model_name=f"{self.imodel.model_alias}+rag"))
