import dash
from dash import dcc, html, Input, Output, State, callback
import dash_bootstrap_components as dbc
import base64
import io
import pandas as pd
import openai
import os
import time
from dash.exceptions import PreventUpdate
import PyPDF2
import docx
import chardet
import json
from dash.exceptions import PreventUpdate

# Initialize the Dash app
app = dash.Dash(__name__, external_stylesheets=[dbc.themes.BOOTSTRAP])

# Get OpenAI API key from Hugging Face Spaces environment variable
openai.api_key = os.environ.get('OPENAI_API_KEY')

# Global variables
uploaded_files = {}
current_matrix = None
matrix_type = None

# Matrix types and their descriptions
matrix_types = {
    "Project Mangement Matrix": "Generate a project management matrix outlining all aspects of project management that summarizes components such as scope, schedule, budget, quality, resources, communications, risk, procurement, and stakeholder management.",
    "Communications Plan Matrix": "Create a matrix showing stakeholders, communication methods, frequency, and responsibilities.",
    "Project Kick-off Matrix": "Generate a matrix outlining key project details, goals, team roles, and initial timelines.",
    "Decision Matrix": "Develop a matrix for evaluating options against criteria, with weighted scores analysis of alternatives style.",
    "Lessons Learned Matrix": "Create a matrix capturing project experiences, challenges, solutions, and recommendations.",
    "Key Performance Indicator Matrix": "Generate a matrix of KPIs, their measurable targets, actual performance, and status.",
    "Prioritization Matrix": "Develop a matrix for ranking tasks or features based on importance and urgency.",
    "Risk Matrix": "Create a matrix identifying tasks with potential risks, their likelihood, impact, and mitigation strategies.",
    "RACI Matrix": "Generate a matrix showing team members and their roles (Responsible, Accountable, Consulted, Informed) for each task.",
    "Project Schedule Matrix": "Develop a matrix showing project phases, tasks, durations, and dependencies.",
    "Quality Control Matrix": "Create a matrix outlining measurable quality standards, testing methods, and acceptance criteria.",
    "Requirements Traceability Matrix": "Generate a matrix linking requirements to their sources, test cases, and status.",
    "Sprint Planning Matrix": "Develop a matrix for sprint nubmer, sprint tasks in that sprint number, story points, assignees, and status.",
    "Test Traceability Matrix": "Create a matrix linking test cases to requirements, execution status, and results.",
    "Sprint Backlog": "Generate a matrix of user stories, tasks, estimates, and priorities for the sprint.",
    "Sprint Retrospective": "Develop a matrix capturing what went well, what didn't, and action items from the sprint.",
    "SWOT Matrix": "Create a matrix analyzing Strengths, Weaknesses, Opportunities, and Threats."
}

app.layout = dbc.Container([
    dbc.Row([
        dbc.Col([
            html.H4("Project Artifacts", className="mt-3 mb-4"),
            dcc.Upload(
                id='upload-files',
                children=html.Div([
                    'Drag and Drop or ',
                    html.A('Select Files')
                ]),
                style={
                    'width': '100%',
                    'height': '60px',
                    'lineHeight': '60px',
                    'borderWidth': '1px',
                    'borderStyle': 'dashed',
                    'borderRadius': '5px',
                    'textAlign': 'center',
                    'margin': '10px 0'
                },
                multiple=True
            ),
            html.Div([
                html.H5("Uploaded Files"),
                dbc.ListGroup(id='file-list')
            ], className="mt-3"),
            html.Hr(),
            html.Div([
                dbc.Button(
                    matrix_type,
                    id={'type': 'matrix-button', 'index': matrix_type},
                    color="link",
                    className="mb-2 w-100 text-left",
                    style={
                        'textAlign': 'left',
                        'textDecoration': 'none',
                        'padding': '0.375rem 0.75rem',
                        'transition': 'background-color 0.3s'
                    }
                ) for matrix_type in matrix_types.keys()
            ], className="left-buttons-container")
        ], width=3),
        dbc.Col([
            html.Div(style={"height": "20px"}),  # Small gap
            html.Div(id="loading-output"),  # Text status box
            dcc.Loading(
                id="loading-indicator",
                type="dot",
                children=[html.Div()]  # Empty div for the blinking dots
            ),
            html.Div(id='matrix-preview', className="border p-3 mb-3"),
            dbc.Button("Download Matrix", id="btn-download", color="success", className="mt-3"),
            dcc.Download(id="download-matrix"),
            html.Hr(),
            html.Div(style={"height": "20px"}),  # Small gap
            dcc.Loading(
                id="chat-loading",
                type="dot",
                children=[
                    dbc.Input(id="chat-input", type="text", placeholder="Chat with GPT to update matrix...", className="mb-2"),
                    dbc.Button("Send", id="btn-send-chat", color="primary", className="mb-3"),
                    html.Div(id="chat-output")
                ]
            )
        ], width=9)
    ]),
    html.Div(id='uploaded-files-state', style={'display': 'none'}),
    dcc.Store(id='matrix-type-store'),
    html.Div(id='matrix-button-clicks', style={'display': 'none'})
], fluid=True)

def parse_file_content(contents, filename):
    content_type, content_string = contents.split(',')
    decoded = base64.b64decode(content_string)
    try:
        if filename.endswith('.pdf'):
            with io.BytesIO(decoded) as pdf_file:
                reader = PyPDF2.PdfReader(pdf_file)
                return ' '.join([page.extract_text() for page in reader.pages])
        elif filename.endswith('.docx'):
            with io.BytesIO(decoded) as docx_file:
                doc = docx.Document(docx_file)
                return ' '.join([para.text for para in doc.paragraphs])
        elif filename.endswith('.txt') or filename.endswith('.rtf'):
            encoding = chardet.detect(decoded)['encoding']
            return decoded.decode(encoding)
        else:
            return "Unsupported file format"
    except Exception as e:
        print(f"Error processing file {filename}: {str(e)}")
        return "Error processing file"

@app.callback(
    Output('file-list', 'children'),
    Input('upload-files', 'contents'),
    State('upload-files', 'filename'),
    State('file-list', 'children')
)
def update_output(list_of_contents, list_of_names, existing_files):
    global uploaded_files
    if list_of_contents is not None:
        new_files = []
        for content, name in zip(list_of_contents, list_of_names):
            file_content = parse_file_content(content, name)
            uploaded_files[name] = file_content
            new_files.append(dbc.ListGroupItem([
                dbc.Button("×", id={'type': 'remove-file', 'index': name}, 
                           color="danger", size="sm", className="me-2"),
                html.Span(name)
            ]))
        existing_files = existing_files or []
        return existing_files + new_files
    return existing_files or []

@app.callback(
    Output('file-list', 'children', allow_duplicate=True),
    Input({'type': 'remove-file', 'index': dash.ALL}, 'n_clicks'),
    State('file-list', 'children'),
    prevent_initial_call=True
)
def remove_file(n_clicks, existing_files):
    global uploaded_files
    ctx = dash.callback_context
    if not ctx.triggered:
        raise PreventUpdate
    
    triggered_id = ctx.triggered[0]['prop_id']
    removed_file = json.loads(triggered_id.split('.')[0])['index']
    
    uploaded_files.pop(removed_file, None)
    return [file for file in existing_files if file['props']['children'][1]['children'] != removed_file]

def generate_matrix_with_gpt(matrix_type, file_contents):
    prompt = f"""Generate a {matrix_type} based on the following project artifacts:
{' '.join(file_contents)}
Instructions:
1. Create the {matrix_type} as a table.
2. Use ONLY pipe symbols (|) to separate columns.
3. Do NOT include any introductory text, descriptions, or explanations.
4. Do NOT use any dashes (-) or other formatting characters.
5. The first row should be the column headers.
6. Start the output immediately with the column headers.
7. Each subsequent row should represent a single item in the matrix.
Example format:
Header1|Header2|Header3
Item1A|Item1B|Item1C
Item2A|Item2B|Item2C
Now, generate the {matrix_type}:
"""

    response = openai.ChatCompletion.create(
        model="gpt-4-turbo",
        messages=[
            {"role": "system", "content": "You are a precise matrix generator that outputs only the requested matrix without any additional text."},
            {"role": "user", "content": prompt}
        ]
    )

    matrix_text = response.choices[0].message.content.strip()
    print("Raw matrix text from GPT:", matrix_text)  # For debugging

    lines = [line.strip() for line in matrix_text.split('\n') if '|' in line]
    data = [line.split('|') for line in lines]
    data = [[cell.strip() for cell in row] for row in data]

    headers = data[0]
    data = data[1:]

    return pd.DataFrame(data, columns=headers)

@app.callback(
    Output('uploaded-files-state', 'children'),
    Input('file-list', 'children')
)
def update_uploaded_files_state(file_list):
    global uploaded_files
    return json.dumps(list(uploaded_files.keys()))

@app.callback(
    Output('matrix-preview', 'children'),
    Output('loading-output', 'children'),
    Input({'type': 'matrix-button', 'index': dash.ALL}, 'n_clicks'),
    State('uploaded-files-state', 'children'),
    prevent_initial_call=True
)
def generate_matrix(n_clicks, uploaded_files_json):
    global current_matrix, matrix_type
    ctx = dash.callback_context
    if not ctx.triggered:
        raise PreventUpdate
    
    button_id = ctx.triggered[0]['prop_id']
    matrix_type = json.loads(button_id.split('.')[0])['index']
    
    if not uploaded_files_json:
        return html.Div("Please upload project artifacts before generating a matrix."), ""

    uploaded_files_list = json.loads(uploaded_files_json)
    
    if not uploaded_files_list:
        return html.Div("Please upload project artifacts before generating a matrix."), ""
    
    try:
        current_matrix = generate_matrix_with_gpt(matrix_type, [uploaded_files[file] for file in uploaded_files_list])
        return dbc.Table.from_dataframe(current_matrix, striped=True, bordered=True, hover=True), f"{matrix_type} generated"
    except Exception as e:
        print(f"Error generating matrix: {str(e)}")
        return html.Div(f"Error generating matrix: {str(e)}"), "Error"

@app.callback(
    Output('chat-output', 'children'),
    Output('matrix-preview', 'children', allow_duplicate=True),
    Input('btn-send-chat', 'n_clicks'),
    State('chat-input', 'value'),
    prevent_initial_call=True
)
def update_matrix_via_chat(n_clicks, chat_input):
    global current_matrix, matrix_type
    if not chat_input or current_matrix is None:
        raise PreventUpdate
    
    prompt = f"""Update the following {matrix_type} based on this instruction: {chat_input}
Current matrix:
{current_matrix.to_string(index=False)}
Instructions:
1. Provide ONLY the updated matrix as a table.
2. Use ONLY pipe symbols (|) to separate columns.
3. Do NOT include any introductory text, descriptions, or explanations.
4. Do NOT use any dashes (-) or other formatting characters.
5. The first row should be the column headers.
6. Start the output immediately with the column headers.
7. Each subsequent row should represent a single item in the matrix.
Now, provide the updated {matrix_type}:
"""
    
    response = openai.ChatCompletion.create(
        model="gpt-3.5-turbo",
        messages=[
            {"role": "system", "content": "You are a precise matrix updater that outputs only the requested matrix without any additional text. You will make assumptions as a project manager to produce the matrix based on the limited informaton provided"},
            {"role": "user", "content": prompt}
        ]
    )

    updated_matrix_text = response.choices[0].message.content.strip()
    print("Raw updated matrix text from GPT:", updated_matrix_text)  # For debugging

    lines = [line.strip() for line in updated_matrix_text.split('\n') if '|' in line]
    data = [line.split('|') for line in lines]
    data = [[cell.strip() for cell in row] for row in data]

    headers = data[0]
    data = data[1:]

    current_matrix = pd.DataFrame(data, columns=headers)
    
    return f"Matrix updated based on: {chat_input}", dbc.Table.from_dataframe(current_matrix, striped=True, bordered=True, hover=True)
    
@app.callback(
    Output("download-matrix", "data"),
    Input("btn-download", "n_clicks"),
    prevent_initial_call=True
)
def download_matrix(n_clicks):
    global current_matrix, matrix_type
    if current_matrix is None:
        raise PreventUpdate
    
    # Create an in-memory Excel file
    output = io.BytesIO()
    with pd.ExcelWriter(output, engine='xlsxwriter') as writer:
        current_matrix.to_excel(writer, sheet_name='Sheet1', index=False)
    
    return dcc.send_bytes(output.getvalue(), f"{matrix_type}.xlsx")

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