import os
import gradio as gr
import ezdxf
import tempfile
import json
import groq

# Function to parse the DXF and extract simple components and connections
def analyze_dxf(file_path):
    try:
        doc = ezdxf.readfile(file_path)
        msp = doc.modelspace()

        components = []
        connections = []

        for entity in msp:
            if entity.dxftype() == "INSERT":
                components.append({
                    "type": "BlockReference",
                    "name": entity.dxf.name,
                    "location": f"({entity.dxf.insert[0]}, {entity.dxf.insert[1]})"
                })
            elif entity.dxftype() == "TEXT":
                components.append({
                    "type": "Text",
                    "content": entity.dxf.text,
                    "location": f"({entity.dxf.insert[0]}, {entity.dxf.insert[1]})"
                })
            elif entity.dxftype() == "LINE":
                connections.append({
                    "type": "Line",
                    "start": f"({entity.dxf.start[0]}, {entity.dxf.start[1]})",
                    "end": f"({entity.dxf.end[0]}, {entity.dxf.end[1]})"
                })

        return components, connections
    except Exception as e:
        return [], []

# Function to ask Groq to analyze components

def ask_groq_question(context_dict, query):
    try:
        client = groq.Client(api_key=os.getenv("GROQ_API_KEY"))

        component_text = "\n".join(
            [f"{comp['type']} - {comp.get('name', comp.get('content', ''))} at {comp['location']}" for comp in context_dict["components"]]
        )
        connection_text = "\n".join(
            [f"{conn['type']} from {conn.get('start', '')} to {conn.get('end', '')}" for conn in context_dict["connections"]]
        )
        context = f"COMPONENTS:\n{component_text}\n\nCONNECTIONS:\n{connection_text}\n\nQUESTION:\n{query}"

        response = client.chat.completions.create(
            model="llama3-8b-8192",
            messages=[{"role": "user", "content": context}]
        )
        return response.choices[0].message.content
    except Exception as e:
        return f"Groq Error: {str(e)}"

# Gradio interface
def process_file(file):
    if file is None:
        return "⚠️ Please upload a valid AutoCAD DXF file."
    components, connections = analyze_dxf(file.name)
    context = {
        "components": components,
        "connections": connections
    }
    # Store context in a temp file to retrieve later
    temp_path = os.path.join(tempfile.gettempdir(), "context.json")
    with open(temp_path, "w") as f:
        json.dump(context, f)
    return "✅ File processed and understood. Now you can ask questions."

def query_groq(user_query):
    temp_path = os.path.join(tempfile.gettempdir(), "context.json")
    if not os.path.exists(temp_path):
        return "⚠️ No DXF context available. Please upload a file first."
    with open(temp_path, "r") as f:
        context = json.load(f)
    return ask_groq_question(context, user_query)

with gr.Blocks() as demo:
    gr.Markdown("# 🧠 RAG-based AutoCAD Analyzer with Groq")

    with gr.Row():
        file_input = gr.File(label="📁 Upload AutoCAD DXF File", file_types=['.dxf'])
        file_output = gr.Textbox(label="Status", interactive=False)

    process_btn = gr.Button("🔍 Analyze DXF")
    process_btn.click(process_file, inputs=[file_input], outputs=[file_output])

    with gr.Row():
        user_query = gr.Textbox(label="🔎 Ask a Question about the Circuit")
        response = gr.Textbox(label="💬 Groq Answer")

    ask_btn = gr.Button("🤖 Ask Groq")
    ask_btn.click(query_groq, inputs=[user_query], outputs=[response])

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