import threading
import asyncio
import json
import os
import ollama
from stream import Stream
from tool import format_tools_for_engine, build_tool_prompt_prefix, ToolSpec

from llama_cpp import Llama
from openrouter import OpenRouter

VALID_ROLES = ['user', 'assistant', 'system', 'tool']
CONTEXT_WINDOW = 10_000

# TODO: Rename to ChatMessage. Add methods for constructing from instruct/openai/etc formatted objects.
class LlamaMessage:
    def __init__(self, role, content):
        if role not in VALID_ROLES:
            raise Exception("message role must be one of: " + ",".join(VALID_ROLES))
        self.role = role
        self.content = content

    def to_dict(self):
        """Serialize to the dict format expected by llama-cpp / ollama."""
        d = {"role": self.role, "content": self.content}
        return d

    def __repr__(self):
        return f"LlamaMessage(role={self.role!r}, content={self.content!r})"

class Models:
    llama_8b = "meta/meta-llama-3-8b-instruct"
    llama_70b = "meta/meta-llama-3-70b-instruct"
    qwen3_8b = "lucataco/qwen3-vl-8b-instruct:39e893666996acf464cff75688ad49ac95ef54e9f1c688fbc677330acc478e11"
    qwen3_235b = "qwen/qwen3-235b-a22b-instruct-2507"

# llama cpp models
model_cache = {}
last_model_file = None
last_model_instance = None

def get_llama_cpp_model(app, model_file):
    global last_model_file, last_model_instance
    if last_model_file == model_file:
        return last_model_instance

    if last_model_instance is not None:
        last_model_instance = None
        #del last_model_instance

    last_model_file = model_file
    last_model_instance = Llama(
        model_path=f"{app.data_directory}/models/{model_file}",
        n_gpu_layers=-1, # Uncomment to use GPU acceleration
        # seed=1337, # Uncomment to set a specific seed
        n_ctx=CONTEXT_WINDOW,
    )
    return last_model_instance



# Generate a conversation based on agent inputs
def build_messages(system_prompt="", instruction=None, input=None, messages=None):
    def normalize(message):
        if isinstance(message, LlamaMessage):
            return message
        elif isinstance(message, dict):
            return LlamaMessage(message['role'], message['content'])
        else:
            raise Exception("Unsupported message type: " + repr(type(message)))

    # shouldn't be necessary. maybe an assert instead?
    messages = [normalize(msg) for msg in messages]
    # default user message
    prompt = input

    # Instruction extends system_prompt if there's input or messages
    # Otherwise instruction becomes the user prompt
    # Input is always added as a user message at the end, if present

    if input is None and not messages:
        prompt = instruction
    elif instruction:
        system_prompt = f"{system_prompt}\n{instruction}" if system_prompt else instruction

    # Add default messages to list
    if system_prompt:
        messages.insert(0, LlamaMessage("system", system))
    if prompt:
        messages.append(LlamaMessage("user", prompt))

    return messages


async def chat_completion(app,
    model_id=None,
    instruction=None,
    input=None,
    messages=None,
    system_prompt="",
    tools=[],
    opts={},
    agent_run=None
):
    db_model = (await app.get_model(model_id)).value

    # Integrate arguments into message list
    messages = build_messages(
        system_prompt=system_prompt,
        instruction=instruction,
        input=input,
        messages=messages,
    )

    model_engine = db_model['engine']
    model_name = db_model['model']

    # Warn when tools are passed to engines that do not support native tool calling
    if tools:
        if model_engine not in ("llama-cpp",):
            print(f"WARNING: Tool calls not supported for engine '{model_engine}'. Tools will be ignored.")

    if model_engine == "llama-cpp":
        result = await local_chat_completion(app, messages, model=model_name, tools=tools, opts=opts, agent_run=agent_run)
    elif model_engine == "replicate":
        # Clone opts to avoid mutating the caller's dict
        opts = dict(opts)
        if 'parameters' in db_model:
            opts['parameters'] = db_model['parameters']
        if 'format' in db_model:
            opts['format'] = db_model['format']
        result = await replicate_completion(app, messages, model=model_name, tools=tools, opts=opts, agent_run=agent_run)
    elif model_engine == "ollama":
        result = await ollama_chat_completion(app, messages, model=model_name, tools=tools, opts=opts, agent_run=agent_run)
    elif model_engine == "openrouter":
        result = await openrouter_chat_completion(app, messages, model=model_name, tools=tools, opts=opts, agent_run=agent_run)
    else:
        raise Exception("Unsupported model: " + repr(model))

    if agent_run:
        return None
    return result


# llama.cpp
async def local_chat_completion(app, messages, model, tools=[], opts={}, agent_run=None):
    if agent_run:
        content_stream = agent_run.stream
    else:
        content_stream = Stream(app)

    def worker():
        llm = get_llama_cpp_model(app, model)

        kwargs = {}

        temp = opts.get("temperature")
        if temp is not None:
            kwargs['temperature'] = temp

        repeat = opts.get("repetition_penalty")
        if repeat is not None:
            kwargs['repeat_penalty'] = repeat

        # Format tools for llama-cpp (OpenAI-compatible function format)
        formatted_tools = format_tools_for_engine(tools, "llama-cpp") if tools else None
        if formatted_tools:
            kwargs['tools'] = formatted_tools
            kwargs['tool_choice'] = 'auto'

        # With stream=True, returns an iterator that infers tokens on demand (openai compatible)
        llm_output = llm.create_chat_completion(
            messages=[msg.to_dict() for msg in messages],
            stream=True,
            **kwargs
        )

        for chunk in llm_output:
            delta = chunk['choices'][0]['delta']
            if 'content' in delta:
                chunk_text = delta['content']
                content_stream.feed(chunk_text)
            # Handle tool_call deltas if the model emits them
            if 'tool_calls' in delta:
                # Forward tool call chunks as JSON strings for the consumer
                for tc in delta['tool_calls']:
                    content_stream.feed(json.dumps({'tool_call': tc}))
                if agent_run:
                    for tc in delta['tool_calls']:
                        agent_run.emit_event("tool_call", {
                            "id": tc.get("id"),
                            "function": tc.get("function"),
                        })

        content_stream.close()

    thread = threading.Thread(target=worker)
    thread.start()

    if agent_run:
        return None
    return content_stream

async def ollama_list(app, messages, model, opts={}):
    # TODO: test the output, and use this
    return ollama.list()

async def ollama_chat_completion(app, messages, model, tools=[], opts={}, agent_run=None):
    if agent_run:
        content_stream = agent_run.stream
    else:
        content_stream = Stream(app)

    async def worker():
        print("START OLLAMA INFERENCE")
        kwargs = {
            'model': model,
            'messages': [msg.to_dict() for msg in messages],
            'stream': True,
        }

        # Pass tools in Ollama format if provided
        formatted_tools = format_tools_for_engine(tools, "ollama") if tools else None
        if formatted_tools:
            kwargs['tools'] = formatted_tools

        stream_response = await ollama.AsyncClient().chat(**kwargs)

        async for part in stream_response:
            if part.message.content:
                content_stream.feed(part.message.content)

        content_stream.close()
        print("END OLLAMA INFERENCE")

    asyncio.create_task(worker())

    if agent_run:
        return None
    return content_stream


async def openrouter_chat_completion(app, messages, model, tools=[], opts={}, agent_run=None):
    if agent_run:
        content_stream = agent_run.stream
    else:
        content_stream = Stream(app)

    async def worker():
        print("START OPENROUTER INFERENCE")
        if not app.openrouter_client:
            raise Exception("OPENROUTER_API_KEY environment variable is not set")

        kwargs = {
            'model': model,
            'messages': [msg.to_dict() for msg in messages],
            'stream': True,
        }

        # Pass tools in OpenAI-compatible format if provided
        formatted_tools = format_tools_for_engine(tools, "openrouter") if tools else None
        if formatted_tools:
            kwargs['tools'] = formatted_tools

        stream_response = await app.openrouter_client.chat.send_async(**kwargs)

        async for event in stream_response:
            if event.choices and len(event.choices) > 0:
                delta = event.choices[0].delta
                if delta.content:
                    content_stream.feed(delta.content)

        content_stream.close()
        print("END OPENROUTER INFERENCE")

    asyncio.create_task(worker())

    if agent_run:
        return None
    return content_stream


async def replicate_completion(app, messages, model=None, tools=[], opts={}, agent_run=None):
    if model is None:
        model = Models.llama_8b

    parameter_names = opts.get('parameters', {})

    arguments = {}

    # Build the prompt, injecting tool descriptions for models that rely on
    # prompt-based tool usage rather than native function calling.
    prompt_text = format_instruct(messages)
    if tools:
        tool_prefix = build_tool_prompt_prefix(tools)
        if tool_prefix:
            prompt_text = tool_prefix + "\n\n" + prompt_text

    arguments[parameter_names.get('prompt', 'prompt')] = prompt_text
    arguments[parameter_names.get('temperature', 'temperature')] = opts.get('temperature') or 0.75
    arguments[parameter_names.get('max_tokens', 'max_tokens')] = opts.get('max_tokens', 2048)

    arguments[parameter_names.get('system_prompt', 'system_prompt')] = ""

    result = await replicate_base_completion(app, model, arguments, agent_run=agent_run)
    if agent_run:
        return None
    return result

async def replicate_base_completion(app, model, input, agent_run=None):
    if agent_run:
        content_stream = agent_run.stream
    else:
        content_stream = Stream(app)
    thread = threading.Thread(
        target=run_replicate_inference_loop,
        kwargs={
            'replicate_client': app.replicate_client,
            'model': model,
            'content_stream': content_stream,
            'input': input
        }
    )
    thread.start()
    #response = await content_stream.content_wait()
    #return response.lstrip()
    if agent_run:
        return None
    return content_stream


def run_replicate_inference_loop(replicate_client, model, input, content_stream):
    print("START INFERENCE")
    completion = replicate_client.run(model, input=input, stream=True)
    for event in completion:
        if event:
            content_stream.feed(str(event))
    print("END INFERENCE")
    content_stream.close()

def format_instruct(messages):
    prompt = "<|begin_of_text|>"
    for message in messages:
        prompt += f"<|start_header_id|>{message.role}<|end_header_id|>\n\n{message.content}<|eot_id|>"
    prompt += f"<|start_header_id|>assistant<|end_header_id|>"
    return prompt