import gradio as gr
from huggingface_hub import hf_hub_download
from llama_cpp import Llama
import re
from datasets import load_dataset
import random
import logging
import os
import autopep8
import textwrap

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Define the model options
gguf_models = {
    "Q8_0 (8-bit)": "leetmonkey_peft__q8_0.gguf",
    "Exact Copy": "leetmonkey_peft_exact_copy.gguf",
    "F16": "leetmonkey_peft_f16.gguf",
    "Super Block Q6": "leetmonkey_peft_super_block_q6.gguf"
}

def download_model(model_name):
    logger.info(f"Downloading model: {model_name}")
    model_path = hf_hub_download(
        repo_id="sugiv/leetmonkey-peft-gguf",
        filename=model_name,
        cache_dir="./models",
        force_download=True,
        resume_download=True
    )
    logger.info(f"Model downloaded: {model_path}")
    return model_path

# Download and load the 8-bit model at startup
q8_model_path = download_model(gguf_models["Q8_0 (8-bit)"])
llm = Llama(
    model_path=q8_model_path,
    n_ctx=2048,
    n_threads=4,
    n_gpu_layers=0,
    verbose=False
)
logger.info("8-bit model loaded successfully")

# Load the dataset
dataset = load_dataset("sugiv/leetmonkey_python_dataset")
train_dataset = dataset["train"]

# Generation parameters
generation_kwargs = {
    "max_tokens": 2048,
    "stop": ["```", "### Instruction:", "### Response:"],
    "echo": False,
    "temperature": 0.2,
    "top_k": 50,
    "top_p": 0.95,
    "repeat_penalty": 1.1
}

def generate_solution(instruction, model):
    system_prompt = "You are a Python coding assistant specialized in solving LeetCode problems. Provide only the complete implementation of the given function. Ensure proper indentation and formatting. Do not include any explanations or multiple solutions."
    full_prompt = f"""### Instruction:
{system_prompt}

Implement the following function for the LeetCode problem:

{instruction}

### Response:
Here's the complete Python function implementation:

```python
"""
    
    response = model(full_prompt, **generation_kwargs)
    return response["choices"][0]["text"]

def extract_and_format_code(text):
    # Extract code between triple backticks
    code_match = re.search(r'```python\s*(.*?)\s*```', text, re.DOTALL)
    if code_match:
        code = code_match.group(1)
    else:
        code = text

    # Remove any text before the function definition
    code = re.sub(r'^.*?(?=def\s+\w+\s*\()', '', code, flags=re.DOTALL)

    # Dedent the code to remove any common leading whitespace
    code = textwrap.dedent(code)

    # Split the code into lines
    lines = code.split('\n')

    # Find the function definition line
    func_def_index = next((i for i, line in enumerate(lines) if line.strip().startswith('def ')), 0)

    # Ensure proper indentation
    indented_lines = [lines[func_def_index]]  # Keep the function definition as is
    for line in lines[func_def_index + 1:]:
        if line.strip():  # If the line is not empty
            indented_lines.append('    ' + line)  # Add 4 spaces of indentation
        else:
            indented_lines.append(line)  # Keep empty lines as is

    formatted_code = '\n'.join(indented_lines)

    try:
        return autopep8.fix_code(formatted_code)
    except:
        return formatted_code

def select_random_problem():
    return random.choice(train_dataset)['instruction']

def update_solution(problem, model_name):
    if model_name == "Q8_0 (8-bit)":
        model = llm
    else:
        model_path = download_model(gguf_models[model_name])
        model = Llama(model_path=model_path, n_ctx=2048, n_threads=4, n_gpu_layers=0, verbose=False)
    
    logger.info(f"Generating solution using {model_name} model")
    generated_output = generate_solution(problem, model)
    formatted_code = extract_and_format_code(generated_output)
    logger.info("Solution generated successfully")
    return formatted_code

def stream_solution(problem, model_name):
    if model_name == "Q8_0 (8-bit)":
        model = llm
    else:
        model_path = download_model(gguf_models[model_name])
        model = Llama(model_path=model_path, n_ctx=2048, n_threads=4, n_gpu_layers=0, verbose=False)
    
    logger.info(f"Generating solution using {model_name} model")
    system_prompt = "You are a Python coding assistant specialized in solving LeetCode problems. Provide only the complete implementation of the given function. Ensure proper indentation and formatting. Do not include any explanations or multiple solutions."
    full_prompt = f"""### Instruction:
{system_prompt}

Implement the following function for the LeetCode problem:

{problem}

### Response:
Here's the complete Python function implementation:

```python
"""
    
    generated_text = ""
    for chunk in model(full_prompt, stream=True, **generation_kwargs):
        token = chunk["choices"][0]["text"]
        generated_text += token
        yield generated_text
    
    formatted_code = extract_and_format_code(generated_text)
    logger.info("Solution generated successfully")
    yield formatted_code

with gr.Blocks() as demo:
    gr.Markdown("# LeetCode Problem Solver")
    
    with gr.Row():
        with gr.Column():
            problem_display = gr.Textbox(label="LeetCode Problem", lines=10)
            select_problem_btn = gr.Button("Select Random Problem")
        
        with gr.Column():
            model_dropdown = gr.Dropdown(choices=list(gguf_models.keys()), label="Select GGUF Model", value="Q8_0 (8-bit)")
            solution_display = gr.Code(label="Generated Solution", language="python", lines=25)
            generate_btn = gr.Button("Generate Solution")
    
    select_problem_btn.click(select_random_problem, outputs=problem_display)
    generate_btn.click(stream_solution, inputs=[problem_display, model_dropdown], outputs=solution_display)

if __name__ == "__main__":
    logger.info("Starting Gradio interface")
    demo.launch(share=True)