import gradio as gr
import requests
import json
from faster_whisper import WhisperModel

OLLAMA_URL = "http://localhost:11434"

MODELS = {
    "⭐ Qwen3 30B-A3B (Best)": "hf.co/bartowski/Qwen_Qwen3-30B-A3B-GGUF:Q4_K_M",
    "Qwen2.5 Coder 7B": "qwen2.5-coder:7b",
    "Qwen2.5 Coder 3B": "qwen2.5-coder:3b",
    "Qwen2.5 Coder 1.5B (Fast)": "qwen2.5-coder:1.5b",
    "DeepSeek Coder 6.7B": "deepseek-coder:6.7b",
    "DeepSeek Coder 1.3B (Fast)": "deepseek-coder:1.3b",
    "StarCoder2 7B": "starcoder2:7b",
    "StarCoder2 3B": "starcoder2:3b",
    "CodeGemma 7B": "codegemma:7b",
    "CodeGemma 2B (Fast)": "codegemma:2b",
}

LANGUAGES = [
    "Python", "JavaScript", "TypeScript", "Go", "Rust", 
    "Java", "C++", "C#", "PHP", "Ruby", "Swift", "Kotlin",
    "HTML/CSS", "SQL", "Bash", "PowerShell"
]

print("Loading Whisper...")
whisper_model = WhisperModel("tiny", device="cpu", compute_type="int8")
print("Whisper ready!")

def get_status():
    try:
        r = requests.get(f"{OLLAMA_URL}/api/tags", timeout=2)
        if r.status_code == 200:
            models = r.json().get("models", [])
            return f"✅ Online | {len(models)} models loaded"
    except:
        pass
    return "⏳ Starting..."

def transcribe_audio(audio):
    if audio is None:
        return ""
    try:
        segments, _ = whisper_model.transcribe(audio)
        return " ".join([seg.text for seg in segments]).strip()
    except Exception as e:
        return f"[STT Error: {e}]"

def chat_stream(message, history, model_name, temperature, max_tokens):
    if not message.strip():
        yield history
        return
    
    model = MODELS.get(model_name, "qwen2.5-coder:3b")
    messages = [{"role": "system", "content": "You are an expert coding assistant. Provide clear, well-commented code. Always use markdown code blocks with language tags."}]

    for user_msg, assistant_msg in history:
        messages.append({"role": "user", "content": user_msg})
        if assistant_msg:
            messages.append({"role": "assistant", "content": assistant_msg})

    messages.append({"role": "user", "content": message})

    try:
        response = requests.post(
            f"{OLLAMA_URL}/api/chat",
            json={"model": model, "messages": messages, "stream": True, "options": {"temperature": temperature, "num_predict": max_tokens}},
            stream=True, timeout=300
        )

        full = ""
        for line in response.iter_lines():
            if line:
                try:
                    data = json.loads(line)
                    if "message" in data:
                        full += data["message"].get("content", "")
                        yield history + [[message, full]]
                except:
                    continue
    except Exception as e:
        yield history + [[message, f"❌ Error: {e}"]]

def generate_code(prompt, language, model_name, temperature, max_tokens):
    if not prompt.strip():
        return "⚠️ Please describe what you want to build."
    
    model = MODELS.get(model_name, "qwen2.5-coder:3b")
    full_prompt = f"""Write {language} code for the following task:

{prompt}

Requirements:
- Clean, production-ready code
- Add helpful comments
- Handle edge cases
- Output ONLY the code in a markdown code block"""

    try:
        r = requests.post(
            f"{OLLAMA_URL}/api/generate",
            json={"model": model, "prompt": full_prompt, "stream": False, "options": {"temperature": temperature, "num_predict": max_tokens}},
            timeout=300
        )
        if r.status_code == 200:
            result = r.json().get("response", "")
            if "```" in result:
                parts = result.split("```")
                if len(parts) >= 2:
                    code = parts[1]
                    if "\n" in code:
                        code = code.split("\n", 1)[-1]
                    return code.strip()
            return result
        return f"❌ Error: {r.text}"
    except Exception as e:
        return f"❌ Error: {e}"

def explain_code(code, model_name, detail_level, max_tokens):
    if not code.strip():
        return "⚠️ Paste code to explain."
    
    model = MODELS.get(model_name, "qwen2.5-coder:3b")
    
    detail_prompts = {
        "Brief": "Give a brief 2-3 sentence explanation of what this code does.",
        "Normal": "Explain what this code does, including the main logic and any important details.",
        "Detailed": "Give a detailed explanation including: purpose, how it works step-by-step, time/space complexity, and potential improvements."
    }
    
    prompt = f"""{detail_prompts[detail_level]}
{code}

    try:
        r = requests.post(
            f"{OLLAMA_URL}/api/generate",
            json={"model": model, "prompt": prompt, "stream": False, "options": {"num_predict": max_tokens}},
            timeout=300
        )
        return r.json().get("response", "") if r.status_code == 200 else f"❌ Error: {r.text}"
    except Exception as e:
        return f"❌ Error: {e}"

def fix_code(code, error, model_name, max_tokens):
    if not code.strip():
        return "⚠️ Paste code to fix."
    
    model = MODELS.get(model_name, "qwen2.5-coder:3b")
    prompt = f"""Fix the following code and explain what was wrong.

Code:
```
{code}
```

Error/Problem: {error or 'Code is not working as expected'}

Provide:
1. The fixed code
2. Brief explanation of what was wrong
3. Any suggestions to prevent similar issues"""

    try:
        r = requests.post(
            f"{OLLAMA_URL}/api/generate",
            json={"model": model, "prompt": prompt, "stream": False, "options": {"temperature": 0.3, "num_predict": max_tokens}},
            timeout=300
        )
        return r.json().get("response", "") if r.status_code == 200 else f"❌ Error: {r.text}"
    except Exception as e:
        return f"❌ Error: {e}"

def review_code(code, model_name, max_tokens):
    if not code.strip():
        return "⚠️ Paste code to review."
    
    model = MODELS.get(model_name, "qwen2.5-coder:3b")
    prompt = f"""Review this code and provide feedback on:

1. **Code Quality** - Is it clean, readable, well-structured?
2. **Bugs/Issues** - Any potential bugs or problems?
3. **Performance** - Any performance concerns?
4. **Security** - Any security issues?
5. **Suggestions** - How could it be improved?
{code}

    try:
        r = requests.post(
            f"{OLLAMA_URL}/api/generate",
            json={"model": model, "prompt": prompt, "stream": False, "options": {"temperature": 0.4, "num_predict": max_tokens}},
            timeout=300
        )
        return r.json().get("response", "") if r.status_code == 200 else f"❌ Error: {r.text}"
    except Exception as e:
        return f"❌ Error: {e}"

css = """
.container { max-width: 1200px; margin: auto; }
footer { display: none !important; }
"""

with gr.Blocks(title="Axon v6", css=css, theme=gr.themes.Soft(primary_hue="violet", secondary_hue="blue")) as demo:
    
    gr.Markdown("""
    # 🔥 Axon v6
    ### AI Coding Assistant • 10 Models • 100% Local • No Rate Limits
    """)
    
    status = gr.Markdown(value=get_status, every=5)
    
    with gr.Row():
        model_dropdown = gr.Dropdown(choices=list(MODELS.keys()), value="Qwen2.5 Coder 3B", label="🤖 Model", scale=3)
        temperature = gr.Slider(0, 1, value=0.7, step=0.1, label="🌡️ Creativity", scale=2)
        max_tokens = gr.Slider(256, 8192, value=2048, step=256, label="📏 Max Length", scale=2)

    with gr.Tabs():
        
        with gr.TabItem("💬 Chat"):
            chatbot = gr.Chatbot(height=450, show_copy_button=True)
            with gr.Row():
                msg = gr.Textbox(placeholder="Ask a coding question...", show_label=False, scale=8, lines=2)
                send = gr.Button("Send", variant="primary", scale=1)
            with gr.Row():
                audio_input = gr.Audio(sources=["microphone"], type="filepath", label="🎤 Voice")
                transcribe_btn = gr.Button("🎤 Transcribe")
                clear = gr.Button("🗑️ Clear")
            with gr.Accordion("💡 Examples", open=False):
                gr.Examples(["Write a Python quicksort function", "Explain async/await in JavaScript"], inputs=msg)
        
        with gr.TabItem("⚡ Generate"):
            with gr.Row():
                with gr.Column(scale=1):
                    gen_prompt = gr.Textbox(label="📝 Describe what you want", lines=4)
                    gen_lang = gr.Dropdown(LANGUAGES, value="Python", label="🔤 Language")
                    gen_temp = gr.Slider(0, 1, value=0.3, step=0.1, label="🌡️ Creativity")
                    gen_btn = gr.Button("⚡ Generate", variant="primary", size="lg")
                with gr.Column(scale=2):
                    gen_output = gr.Code(label="Generated Code", language="python", lines=20)
        
        with gr.TabItem("🔍 Explain"):
            with gr.Row():
                with gr.Column(scale=1):
                    explain_input = gr.Code(label="📋 Paste code", lines=12)
                    explain_detail = gr.Radio(["Brief", "Normal", "Detailed"], value="Normal", label="📊 Detail")
                    explain_btn = gr.Button("🔍 Explain", variant="primary", size="lg")
                with gr.Column(scale=1):
                    explain_output = gr.Markdown(label="Explanation")
        
        with gr.TabItem("🔧 Fix"):
            with gr.Row():
                with gr.Column(scale=1):
                    fix_input = gr.Code(label="🐛 Buggy code", lines=10)
                    fix_error = gr.Textbox(label="❌ Error message", lines=3)
                    fix_btn = gr.Button("🔧 Fix", variant="primary", size="lg")
                with gr.Column(scale=1):
                    fix_output = gr.Markdown(label="Fixed Code")
        
        with gr.TabItem("📋 Review"):
            with gr.Row():
                with gr.Column(scale=1):
                    review_input = gr.Code(label="📋 Code to review", lines=15)
                    review_btn = gr.Button("📋 Review", variant="primary", size="lg")
                with gr.Column(scale=1):
                    review_output = gr.Markdown(label="Review")

    def respond(message, history, model, temp, tokens):
        history = history or []
        for updated_history in chat_stream(message, history, model, temp, tokens):
            yield updated_history, ""
    
    msg.submit(respond, [msg, chatbot, model_dropdown, temperature, max_tokens], [chatbot, msg])
    send.click(respond, [msg, chatbot, model_dropdown, temperature, max_tokens], [chatbot, msg])
    clear.click(lambda: [], None, chatbot)
    transcribe_btn.click(transcribe_audio, audio_input, msg)
    gen_btn.click(generate_code, [gen_prompt, gen_lang, model_dropdown, gen_temp, max_tokens], gen_output)
    explain_btn.click(explain_code, [explain_input, model_dropdown, explain_detail, max_tokens], explain_output)
    fix_btn.click(fix_code, [fix_input, fix_error, model_dropdown, max_tokens], fix_output)
    review_btn.click(review_code, [review_input, model_dropdown, max_tokens], review_output)

demo.launch(server_name="0.0.0.0", server_port=7860)