import os
import subprocess
import io
import base64
import threading
import logging
import tempfile
import shutil
import uuid
import json
import re
from pathlib import Path
from flask import request
from dash import Dash, dcc, html, Input, Output, State, ctx, no_update
import dash_bootstrap_components as dbc
import openai
import anthropic
from google import generativeai as genai
import requests
import graphviz

logging.basicConfig(level=logging.INFO)

os.environ['PATH'] += ':/usr/bin:/usr/local/bin:/home/user/.local/bin'

try:
    result = subprocess.run(['which', 'dot'], check=True, capture_output=True, text=True)
    print(f"dot location: {result.stdout}")
    result = subprocess.run(['dot', '-V'], check=True, capture_output=True, text=True)
    print(f"Graphviz version: {result.stdout}")
except Exception as e:
    print(f"Error checking Graphviz: {e}")

print(f"Current PATH: {os.environ['PATH']}")

openai.api_key = os.getenv("OPENAI_API_KEY")
if not openai.api_key:
    logging.warning("OPENAI_API_KEY not set. GPT-3.5 model will not be available.")

google_api_key = os.getenv("GOOGLE_API_KEY")
if google_api_key:
    try:
        genai.configure(api_key=google_api_key)
    except Exception as e:
        logging.error(f"Failed to configure Google Gemini: {e}")
        genai = None
else:
    genai = None
    logging.warning("GOOGLE_API_KEY not set. Gemini model will not be available.")

anthropic_api_key = os.getenv("ANTHROPIC_API_KEY")
if not anthropic_api_key:
    logging.warning("ANTHROPIC_API_KEY not set. Claude model will not be available.")

grok_api_key = os.getenv("GROK_API_KEY")
if not grok_api_key:
    logging.warning("GROK_API_KEY not set. Groq model will not be available.")

SESSION_DIR = Path(tempfile.gettempdir()) / "arch_diagram_sessions"
SESSION_DIR.mkdir(parents=True, exist_ok=True)
session_locks = {}
session_data = {}

def get_or_create_session_id():
    sid = request.cookies.get("sid")
    if not sid or not isinstance(sid, str) or len(sid) < 8:
        sid = str(uuid.uuid4())
    return sid

def get_session_dir(sid):
    p = SESSION_DIR / sid
    p.mkdir(parents=True, exist_ok=True)
    return p

def get_session_lock(sid):
    if sid not in session_locks:
        session_locks[sid] = threading.Lock()
    return session_locks[sid]

def get_session_data(sid):
    if sid not in session_data:
        session_data[sid] = {"diagram_base64": None, "description": "", "model": "gpt-3.5-turbo"}
    return session_data[sid]

def save_session_state(sid, data):
    session_data[sid] = data
    sdir = get_session_dir(sid)
    with open(sdir / "session.json", "w") as f:
        json.dump(data, f)

def load_session_state(sid):
    sdir = get_session_dir(sid)
    sf = sdir / "session.json"
    if sf.exists():
        try:
            with open(sf) as f:
                session_data[sid] = json.load(f)
        except Exception as e:
            logging.error(f"Session load error: {e}")

def clear_session(sid):
    sdir = get_session_dir(sid)
    try:
        if sdir.exists():
            shutil.rmtree(sdir)
        session_data.pop(sid, None)
        session_locks.pop(sid, None)
        logging.info(f"Session {sid} cleared.")
    except Exception as e:
        logging.error(f"Error clearing session {sid}: {e}")

def strip_markdown_codeblock(text):
    if text is None:
        return ""
    pattern = r"^```[a-zA-Z0-9]*\s*([\s\S]*?)```$"
    match = re.search(pattern, text.strip(), re.MULTILINE)
    if match:
        return match.group(1).strip()
    return text.strip()

def generate_diagram(dot_code, sid):
    dot_code = dot_code.strip()
    try:
        src = graphviz.Source(dot_code)
        out_img = src.pipe(format="png")
        diagram_base64 = base64.b64encode(out_img).decode('utf-8')
        return diagram_base64
    except Exception as e:
        logging.error(f"Graphviz error: {e}")
        raise Exception("Graphviz rendering failed. Please check your description or try again.")

def get_ai_response(model, user_prompt):
    instruction = (
        "You are an expert architecture-to-graphviz converter. "
        "Given any architecture description, output ONLY the Graphviz DOT code to represent the architecture, "
        "using the best layout for architectural diagrams (using digraph and edge/node attributes as needed). "
        "Output ONLY the code, in a markdown ```dot block, and nothing else. "
        "No explanations, no intro or outro, just the code. "
    )
    prompt = f"{instruction}\n\n{user_prompt}"
    if model == 'gpt-3.5-turbo':
        response = openai.ChatCompletion.create(
            model="gpt-3.5-turbo",
            messages=[{"role": "user", "content": prompt}]
        )
        return response.choices[0].message.content
    elif model == 'gemini-1.5-flash-latest':
        if genai is None:
            raise Exception("Gemini API not configured.")
        model_obj = genai.GenerativeModel('gemini-1.5-pro-latest')
        response = model_obj.generate_content(prompt)
        return response.text
    elif model == 'claude-3-5-haiku-20241022':
        if not anthropic_api_key:
            raise Exception("Anthropic API not configured.")
        client = anthropic.Anthropic(api_key=anthropic_api_key)
        response = client.messages.create(
            model="claude-3-5-haiku-20241022",
            system=None,
            max_tokens=2048,
            messages=[{"role": "user", "content": prompt}]
        )
        return response.content[0].text
    elif model == 'grok-3-mini-fast-beta':
        if not grok_api_key:
            raise Exception("Groq API not configured.")
        groq_url = "https://api.groq.com/openai/v1/chat/completions"
        headers = {
            "Authorization": f"Bearer {grok_api_key}",
            "Content-Type": "application/json"
        }
        data = {
            "model": "grok-3-mini-fast-beta",
            "messages": [{"role": "user", "content": prompt}]
        }
        response = requests.post(groq_url, json=data, headers=headers)
        return response.json()['choices'][0]['message']['content']
    else:
        raise Exception("Error: Invalid model selected.")

app = Dash(__name__, external_stylesheets=[dbc.themes.BOOTSTRAP])

app.layout = dbc.Container(fluid=True, children=[
    dbc.Row([
        dbc.Col([
            html.H1("Architecture Diagram Generator", className="mb-4"),
        ], width=12)
    ], className="mt-2"),
    dbc.Row([
        dbc.Col([
            dbc.Card([
                dbc.CardHeader("Controls"),
                dbc.CardBody([
                    html.Div([
                        dcc.Dropdown(
                            id='model-dropdown',
                            options=[
                                {'label': 'GPT-3.5 Turbo', 'value': 'gpt-3.5-turbo'},
                                {'label': 'Gemini 1.5 Flash', 'value': 'gemini-1.5-flash-latest'},
                                {'label': 'Claude 3.5 Haiku', 'value': 'claude-3-5-haiku-20241022'},
                                {'label': 'Grok 3 Mini Fast', 'value': 'grok-3-mini-fast-beta'}
                            ],
                            value="gpt-3.5-turbo",
                            className="mb-3"
                        ),
                        dbc.Textarea(
                            id='description-input', 
                            placeholder="Enter architecture description here...",
                            style={'height': '200px', 'whiteSpace': 'pre-wrap', 'wordWrap': 'break-word'},
                            className="mb-3"
                        ),
                        dbc.Button("Generate Diagram", 
                                   id='generate-button', 
                                   color="primary", 
                                   className="me-2 mb-2"),
                        dbc.Button("Download Diagram", 
                                   id='download-button', 
                                   color="secondary", 
                                   className="mb-2", 
                                   disabled=True),
                        dbc.Button("Clear Session", 
                                   id='clear-session-button',
                                   color="danger",
                                   className="mb-2",
                                   style={'marginTop': '0.5em'}),
                        dcc.Download(id="download-diagram"),
                        html.Div(id='status-message', className="mb-3"),
                    ]),
                ])
            ])
        ], width=4, style={'backgroundColor': '#f8f9fa', 'minHeight': '100vh'}),
        dbc.Col([
            dbc.Card([
                dbc.CardHeader("Diagram"),
                dbc.CardBody([
                    dcc.Loading(
                        id="loading",
                        type="default",
                        fullscreen=False,
                        children=[
                            html.Img(
                                id='diagram-output',
                                src="",
                                style={'width': '100%', 'minHeight': '300px'}
                            ),
                        ]
                    )
                ])
            ]),
        ], width=8, style={'backgroundColor': '#fff', 'minHeight': '100vh'}),
    ])
])

@app.server.after_request
def set_cookie(response):
    try:
        sid = request.cookies.get("sid")
        if not sid or not isinstance(sid, str) or len(sid) < 8:
            sid = str(uuid.uuid4())
        response.set_cookie("sid", sid, max_age=60*60*24*30, httponly=False, samesite="Lax")
    except Exception as e:
        logging.error(f"Error setting cookie: {e}")
    return response

@app.callback(
    Output('diagram-output', 'src'),
    Output('status-message', 'children'),
    Output('download-button', 'disabled'),
    Output('description-input', 'value'),
    Output('model-dropdown', 'value'),
    Output('download-diagram', 'data'),
    Input('generate-button', 'n_clicks'),
    Input('download-button', 'n_clicks'),
    Input('clear-session-button', 'n_clicks'),
    State('description-input', 'value'),
    State('model-dropdown', 'value'),
    prevent_initial_call=True
)
def main_callback(generate_clicks, download_clicks, clear_clicks, description, model):
    trigger = ctx.triggered_id
    try:
        sid = get_or_create_session_id()
    except Exception as e:
        logging.error(f"Session id error: {e}")
        return "", "Session error. Try refreshing the page.", True, "", "gpt-3.5-turbo", None

    lock = get_session_lock(sid)

    if trigger == "generate-button":
        if not description:
            return no_update, "Please enter a description.", True, no_update, no_update, None
        with lock:
            try:
                ai_response = get_ai_response(model, description)
                dot_code = strip_markdown_codeblock(ai_response)
                diagram_base64 = generate_diagram(dot_code, sid)
                sdata = get_session_data(sid)
                sdata["diagram_base64"] = diagram_base64
                sdata["description"] = description
                sdata["model"] = model
                save_session_state(sid, sdata)
                logging.info(f"Session {sid}: Diagram generated successfully.")
                return f'data:image/png;base64,{diagram_base64}', "Diagram generated successfully!", False, description, model, None
            except Exception as e:
                logging.error(f"Session {sid}: Error generating diagram: {str(e)}")
                return no_update, f"Error: {str(e)}", True, no_update, no_update, None

    elif trigger == "download-button":
        with lock:
            load_session_state(sid)
            sdata = get_session_data(sid)
            diagram_base64 = sdata.get("diagram_base64", None)
            if not diagram_base64:
                return no_update, "No diagram to download.", True, sdata.get("description", ""), sdata.get("model", "gpt-3.5-turbo"), None
            try:
                diagram_bytes = base64.b64decode(diagram_base64)
                logging.info(f"Session {sid}: Diagram downloaded.")
                return no_update, "Download started.", False, sdata.get("description", ""), sdata.get("model", "gpt-3.5-turbo"), dcc.send_bytes(diagram_bytes, "architecture_diagram.png")
            except Exception as e:
                logging.error(f"Session {sid}: Download error: {e}")
                return no_update, f"Download error: {e}", False, sdata.get("description", ""), sdata.get("model", "gpt-3.5-turbo"), None

    elif trigger == "clear-session-button":
        with lock:
            clear_session(sid)
            logging.info(f"Session {sid}: Session cleared by user.")
        return "", "Session cleared. All data deleted.", True, "", "gpt-3.5-turbo", None

    else:
        with lock:
            load_session_state(sid)
            sdata = get_session_data(sid)
            diagram_base64 = sdata.get("diagram_base64", None)
            description_val = sdata.get("description", "")
            model_val = sdata.get("model", "gpt-3.5-turbo")
            if diagram_base64:
                return f'data:image/png;base64,{diagram_base64}', "", False, description_val, model_val, None
            else:
                return "", "", True, description_val, model_val, None

if __name__ == '__main__':
    print("Starting the Dash application...")
    app.run(debug=True, host='0.0.0.0', port=7860, threaded=True)
    print("Dash application has finished running.")