import gradio as gr
import subprocess
import os
import requests
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import Optional

# ==========================================
# 1. FastAPI Setup & API Endpoints
# ==========================================

# Initialize FastAPI with custom docs URL
app = FastAPI(
    title="Qwen 2.5 Coder API",
    description="Automated API for Ollama on Hugging Face Spaces",
    version="1.0.0",
    docs_url="/api/docs", # Exposes Swagger UI here
    redoc_url=None
)

# Define the expected JSON payload for the /ask endpoint
class AskRequest(BaseModel):
    prompt: str
    system: Optional[str] = "You are an expert programming assistant."
    stream: Optional[bool] = False

@app.post("/ask", summary="Generate code or text using Qwen2.5-Coder:3b")
def ask_model(request: AskRequest):
    """
    Sends a prompt to the local Ollama instance and returns the generated text.
    """
    try:
        # Forward the request internally to the local Ollama server running on 11434
        response = requests.post("http://localhost:11434/api/generate", json={
            "model": "qwen2.5-coder:3b",
            "prompt": request.prompt,
            "system": request.system,
            "stream": request.stream
        })
        response.raise_for_status()
        return response.json()
    except requests.exceptions.RequestException as e:
        raise HTTPException(status_code=500, detail=f"Error connecting to Ollama: {str(e)}")

@app.get("/api/health")
def health_check():
    return {"status": "active", "model": "qwen2.5-coder:3b is ready"}


# ==========================================
# 2. Gradio Terminal Setup
# ==========================================

class Terminal:
    def __init__(self):
        self.current_dir = "/app"
        self.output_history = []
    
    def execute_stream(self, command):
        if not command.strip():
            yield self.get_full_output()
            return

        cmd = command.strip()
        
        if cmd.startswith("cd "):
            try:
                target = cmd[3:].strip()
                if not target: target = os.path.expanduser("~")
                if not os.path.isabs(target): target = os.path.join(self.current_dir, target)
                target = os.path.normpath(target)
                
                if os.path.isdir(target):
                    os.chdir(target)
                    self.current_dir = os.getcwd()
                    self.output_history.append(f"$ {cmd}\nChanged directory to: {self.current_dir}")
                else:
                    self.output_history.append(f"$ {cmd}\ncd: {target}: No such directory")
            except Exception as e:
                self.output_history.append(f"$ {cmd}\ncd error: {str(e)}")
            yield self.get_full_output()
            return

        if cmd in ["clear", "cls"]:
            self.output_history = []
            yield self.get_full_output()
            return

        try:
            process = subprocess.Popen(
                ["bash", "-c", cmd],
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                text=True,
                cwd=self.current_dir,
                bufsize=1,
                universal_newlines=True
            )
            
            live_output = ""
            for line in iter(process.stdout.readline, ''):
                live_output += line
                recent_history = self.output_history[-15:] 
                current_view = "\n\n".join(recent_history)
                if current_view:
                    current_view += "\n\n"
                    
                yield current_view + f"$ {cmd}\n{live_output}\n{self.current_dir}$ [Running...]"
            
            process.wait()
            self.output_history.append(f"$ {cmd}\n{live_output}".strip())
            yield self.get_full_output()
            
        except Exception as e:
            self.output_history.append(f"$ {cmd}\nError: {str(e)}")
            yield self.get_full_output()

    def get_full_output(self):
        if not self.output_history:
            return f"{self.current_dir}$ "
        
        recent_history = self.output_history[-15:]
        full_output = "\n\n".join(recent_history)
        full_output += f"\n\n{self.current_dir}$ "
        return full_output

terminal = Terminal()

custom_css = """
#terminal-output textarea {
    font-family: 'Monaco', 'Menlo', 'Ubuntu Mono', monospace !important;
    font-size: 14px !important;
    background-color: #1e1e1e !important;
    color: #e0e0e0 !important;
    border: 1px solid #444 !important;
}
#terminal-output {
    max-height: 600px;
    overflow-y: auto;
}
"""

with gr.Blocks(title="Ubuntu AI Terminal", css=custom_css) as demo:
    gr.Markdown("# 🐧 Root Terminal (API & Ollama Running)")
    gr.Markdown("**API Docs:** Available at `/api/docs` | **API Endpoint:** Available at `/ask`")
    
    terminal_output = gr.Textbox(
        value=terminal.get_full_output(), lines=25, elem_id="terminal-output", interactive=False, show_label=False
    )
    
    with gr.Row():
        cmd_input = gr.Textbox(placeholder="Enter command (press Enter to execute)...", scale=8, show_label=False)
        execute_btn = gr.Button("Run", variant="primary", scale=1)
        clear_btn = gr.Button("Clear", variant="secondary", scale=1)
    
    gr.Markdown("### 🚀 Quick AI Commands")
    with gr.Row():
        btn_check_model = gr.Button("1. Check Downloaded Models", size="sm")
        btn_test_api = gr.Button("2. Test Local /ask API", size="sm")
        btn_expose_api = gr.Button("3. Expose via Localtunnel", size="sm")
        
        btn_check_model.click(lambda: "ollama list", outputs=[cmd_input])
        btn_test_api.click(lambda: 'curl -X POST http://localhost:7860/ask -H "Content-Type: application/json" -d \'{"prompt":"Write a Python hello world"}\'', outputs=[cmd_input])
        btn_expose_api.click(lambda: "lt --port 7860 --subdomain my-custom-ai-agent", outputs=[cmd_input])

    cmd_input.submit(terminal.execute_stream, inputs=[cmd_input], outputs=[terminal_output]).then(lambda: "", outputs=[cmd_input])
    execute_btn.click(terminal.execute_stream, inputs=[cmd_input], outputs=[terminal_output]).then(lambda: "", outputs=[cmd_input])
    clear_btn.click(lambda: next(terminal.execute_stream("clear")), outputs=[terminal_output])

# ==========================================
# 3. Mount Gradio to FastAPI
# ==========================================
app = gr.mount_gradio_app(app, demo, path="/")

# Note: The app is now started by Uvicorn via the Dockerfile CMD, not by demo.launch()