import gradio as gr
import os
import requests
from urllib.parse import urlparse

try:
    from groq import Groq
except ImportError:
    print("Error: The Groq Python client is required. Install with 'pip install groq'")

def extract_repo_id(repo_input):
    if repo_input.startswith(("http://", "https://")):
        parsed = urlparse(repo_input)
        path = parsed.path.strip("/")
        parts = path.split("/")
        if len(parts) >= 2:
            return f"{parts[0]}/{parts[1]}"
        return None
    else:
        if "/" in repo_input:
            return repo_input
        return None

def get_repo_info(repo_id):
    try:
        # Try model endpoint
        model_response = requests.get(f"https://huggingface.co/api/models/{repo_id}")
        if model_response.status_code == 200:
            return model_response.json(), "model"
        
        # Try dataset endpoint
        dataset_response = requests.get(f"https://huggingface.co/api/datasets/{repo_id}")
        if dataset_response.status_code == 200:
            return dataset_response.json(), "dataset"
        # Try Spaces endpoint
        spaces_response = requests.get(f"https://huggingface.co/api/spaces/{repo_id}")
        if spaces_response.status_code == 200:
            return spaces_response.json(), "spaces"
        
        return None, None
    except Exception as e:
        print(f"Error fetching repo info: {e}")
        return None, None

def generate_readme(repo_input):
    try:
        # Validate and extract repo ID
        repo_id = extract_repo_id(repo_input)
        if not repo_id:
            return "Invalid repository format. Please use 'user/repo' or a Hugging Face URL"
        
        # Get repository information
        repo_info, repo_type = get_repo_info(repo_id)
        if not repo_info:
            return "Repository not found or inaccessible"
        
        # Prepare prompt for Groq
        prompt = f"""Generate a professional README.md for the Hugging Face {repo_type} repository {repo_id}.
        Use the following information:
        - Author: {repo_info.get('author', 'Unknown')}
        - Description: {repo_info.get('description', 'No description')}
        - Tags: {', '.join(repo_info.get('tags', []))}
        - License: {repo_info.get('license', 'Unknown')}
        
        Include these sections:
        1. Overview
        2. Installation
        3. Usage
        4. Examples
        5. License
        6. Citation (if applicable)
        
        Format the README in proper Markdown with code blocks where appropriate."""
        
        # Initialize Groq client
        client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
        
        # Generate completion
        completion = client.chat.completions.create(
            messages=[{"role": "user", "content": prompt}],
            model="mixtral-8x7b-32768",
            temperature=0.3,
            max_tokens=1024
        )
        
        return completion.choices[0].message.content
        
    except Exception as e:
        return f"Error generating README: {str(e)}"

with gr.Blocks(
    theme=os.environ.get("gradio_theme"), 
    title="HF Repo README Generator",
) as demo:
    gr.Markdown("# 🚀 Hugging Face Repository README Generator")
    
    with gr.Row():
        with gr.Column():
            repo_input = gr.Textbox(
                label="Hugging Face Repository URL or ID",
                placeholder="Enter 'user/repo' or full URL...",
                max_lines=1
            )
            submit_btn = gr.Button("Generate README", variant="primary")
    
    with gr.Row():
        output = gr.Textbox(label="Generated README", lines=20, interactive=True)
    
    submit_btn.click(
        fn=generate_readme,
        inputs=repo_input,
        outputs=output
    )

if __name__ == "__main__":
    demo.launch()