import gradio as gr
import requests
import logging
from typing import Optional, Dict, Any
from pydantic import BaseModel

logging.basicConfig(filename="ui.log", level=logging.INFO)
logger = logging.getLogger(__name__)

class ResourceLimits(BaseModel):
    cpu_limit: Optional[int] = None
    memory_limit: Optional[int] = None
    gpu_power_limit: Optional[int] = None

class APIClient:
    def __init__(self, base_url: str, api_key: str = None):
        self.base_url = base_url.rstrip("/")
        self.api_key = api_key

    def _headers(self):
        return {"Authorization": f"Bearer {self.api_key}"} if self.api_key else {}

    def get_metrics(self) -> Dict[str, Any]:
        try:
            resp = requests.get(f"{self.base_url}/metrics", headers=self._headers(), timeout=10)
            resp.raise_for_status()
            return resp.json()
        except Exception as e:
            logger.error(f"Error fetching metrics: {e}")
            return {}

    def set_limits(self, limits: ResourceLimits) -> bool:
        try:
            resp = requests.post(
                f"{self.base_url}/limits",
                json=limits.dict(exclude_none=True),
                headers=self._headers(),
                timeout=10,
            )
            resp.raise_for_status()
            return True
        except Exception as e:
            logger.error(f"Error setting limits: {e}")
            return False

def build_dashboard(api_client: APIClient):
    def fetch_metrics():
        return api_client.get_metrics()

    def update_limits(cpu, memory, gpu):
        limits = ResourceLimits(cpu_limit=cpu, memory_limit=memory, gpu_power_limit=gpu)
        success = api_client.set_limits(limits)
        return "Limits updated!" if success else "Failed to update limits."

    with gr.Row():
        with gr.Column():
            metrics_output = gr.JSON(label="Current Metrics")
            refresh_btn = gr.Button("Refresh Metrics")
        with gr.Column():
            cpu = gr.Number(label="CPU Limit")
            memory = gr.Number(label="Memory Limit")
            gpu = gr.Number(label="GPU Power Limit")
            set_btn = gr.Button("Set Limits")
            status = gr.Textbox(label="Status", interactive=False)

    refresh_btn.click(fetch_metrics, outputs=metrics_output)
    set_btn.click(update_limits, inputs=[cpu, memory, gpu], outputs=status)

def main():
    api_client = APIClient(base_url="http://localhost:8000")
    with gr.Blocks(theme=gr.themes.Default()) as app:
        gr.Markdown("# Resource Dashboard")
        build_dashboard(api_client)
    app.launch(server_name="0.0.0.0", server_port=7860, share=False)

if __name__ == "__main__":
    main()