import os
import json
import logging
from datetime import datetime
from flask import Flask, render_template, request, jsonify, send_file
from flask_cors import CORS
import threading
from pathlib import Path
import numpy as np
import networkx as nx
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from io import BytesIO
import base64

from sentence_transformers import SentenceTransformer
from groq import Groq
import PyPDF2
import pandas as pd
from langchain.text_splitter import RecursiveCharacterTextSplitter

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = Flask(__name__)
CORS(app)

UPLOAD_FOLDER = Path('./data/uploads')
GRAPH_DATA_FOLDER = Path('./data/graph_data')
UPLOAD_FOLDER.mkdir(parents=True, exist_ok=True)
GRAPH_DATA_FOLDER.mkdir(parents=True, exist_ok=True)

app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
app.config['MAX_CONTENT_LENGTH'] = 50 * 1024 * 1024

embedding_model = SentenceTransformer('all-MiniLM-L6-v2')

groq_api_key = os.getenv('GROQ_API_KEY', '')
if groq_api_key:
    groq_client = Groq(api_key=groq_api_key)
else:
    groq_client = None
    logger.warning("GROQ_API_KEY not set - chat queries will fail")

documents_state = {}
graph_data = {}

class DocumentProcessor:
    @staticmethod
    def extract_text(file_path):
        file_ext = Path(file_path).suffix.lower()

        if file_ext == '.pdf':
            text = DocumentProcessor._extract_pdf(file_path)
        elif file_ext == '.csv':
            text = DocumentProcessor._extract_csv(file_path)
        elif file_ext == '.txt':
            with open(file_path, 'r', encoding='utf-8') as f:
                text = f.read()
        else:
            raise ValueError(f"Unsupported file type: {file_ext}")

        return text

    @staticmethod
    def _extract_pdf(file_path):
        text = []
        with open(file_path, 'rb') as f:
            reader = PyPDF2.PdfReader(f)
            for page in reader.pages:
                text.append(page.extract_text())
        return '\n'.join(text)

    @staticmethod
    def _extract_csv(file_path):
        df = pd.read_csv(file_path)
        return df.to_string()

    @staticmethod
    def chunk_text(text, chunk_size=500, overlap=100):
        splitter = RecursiveCharacterTextSplitter(
            chunk_size=chunk_size,
            chunk_overlap=overlap,
            separators=["\n\n", "\n", " ", ""]
        )
        chunks = splitter.split_text(text)
        return chunks

class GraphBuilder:
    @staticmethod
    def build_knowledge_graph(chunks, doc_name):
        graph = nx.DiGraph()
        entities = set()

        for chunk in chunks:
            words = chunk.split()[:10]
            main_entity = f"{doc_name}_chunk_{chunks.index(chunk)}"
            graph.add_node(main_entity, type='chunk', content=chunk[:200])

            for word in words:
                if len(word) > 3:
                    word_node = word.lower()
                    graph.add_node(word_node, type='entity')
                    graph.add_edge(main_entity, word_node, weight=1.0)
                    entities.add(word_node)

        return graph, list(entities)

    @staticmethod
    def visualize_graph(graph, output_path):
        plt.figure(figsize=(14, 10))

        if len(graph.nodes()) == 0:
            plt.text(0.5, 0.5, 'Empty Graph', ha='center', va='center')
        else:
            pos = nx.spring_layout(graph, k=2, iterations=50, seed=42)

            node_colors = []
            for node in graph.nodes():
                if graph.nodes[node].get('type') == 'chunk':
                    node_colors.append('lightblue')
                else:
                    node_colors.append('lightgreen')

            nx.draw_networkx_nodes(graph, pos, node_color=node_colors,
                                  node_size=500, alpha=0.9)
            nx.draw_networkx_edges(graph, pos, edge_color='gray',
                                  arrows=True, alpha=0.5, width=1.5)

            labels = {node: node[:15] for node in graph.nodes()}
            nx.draw_networkx_labels(graph, pos, labels, font_size=8)

        plt.title('Knowledge Graph Visualization', fontsize=16, fontweight='bold')
        plt.axis('off')
        plt.tight_layout()
        plt.savefig(output_path, dpi=150, bbox_inches='tight', facecolor='white')
        plt.close()

        logger.info(f"Graph visualization saved to {output_path}")

def process_document_async(filename, file_path):
    try:
        logger.info(f"Processing document: {filename}")
        documents_state[filename] = {
            'status': 'processing',
            'progress': 10,
            'error': None,
            'chunks': 0,
            'entities': 0
        }

        text = DocumentProcessor.extract_text(file_path)
        documents_state[filename]['progress'] = 40

        chunks = DocumentProcessor.chunk_text(text)
        documents_state[filename]['progress'] = 60
        logger.info(f"Created {len(chunks)} chunks from {filename}")

        graph, entities = GraphBuilder.build_knowledge_graph(chunks, filename)
        documents_state[filename]['progress'] = 80
        logger.info(f"Built graph with {len(graph.nodes())} nodes for {filename}")

        graph_path = GRAPH_DATA_FOLDER / f"{filename}_graph.png"
        GraphBuilder.visualize_graph(graph, graph_path)
        documents_state[filename]['progress'] = 95

        embeddings = embedding_model.encode(chunks, show_progress_bar=False)

        graph_data[filename] = {
            'chunks': chunks,
            'embeddings': embeddings.tolist(),
            'graph': nx.node_link_data(graph),
            'entities': entities,
            'graph_image': str(graph_path)
        }

        documents_state[filename] = {
            'status': 'ready',
            'progress': 100,
            'error': None,
            'chunks': len(chunks),
            'entities': len(entities),
            'graph_image': f"/graph-image/{filename}",
            'timestamp': datetime.now().isoformat()
        }

        logger.info(f"Successfully processed {filename}")

    except Exception as e:
        logger.error(f"Error processing {filename}: {str(e)}")
        documents_state[filename] = {
            'status': 'error',
            'progress': 0,
            'error': str(e),
            'chunks': 0,
            'entities': 0
        }

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/api/documents', methods=['GET'])
def get_documents():
    return jsonify({
        'documents': documents_state,
        'api_key_set': bool(os.getenv('GROQ_API_KEY')),
        'timestamp': datetime.now().isoformat()
    })

@app.route('/api/upload', methods=['POST'])
def upload_document():
    if 'files' not in request.files:
        return jsonify({'error': 'No files provided'}), 400

    files = request.files.getlist('files')
    results = {'successful': 0, 'failed': 0, 'files': []}

    for file in files:
        if not file or file.filename == '':
            results['failed'] += 1
            continue

        filename = file.filename
        file_path = UPLOAD_FOLDER / filename
        file.save(file_path)

        documents_state[filename] = {
            'status': 'queued',
            'progress': 0,
            'error': None,
            'chunks': 0,
            'entities': 0
        }

        thread = threading.Thread(target=process_document_async, args=(filename, file_path))
        thread.daemon = True
        thread.start()

        results['successful'] += 1
        results['files'].append(filename)

    return jsonify({
        'success': True,
        'message': f"✅ {results['successful']} file(s) queued for processing",
        **results
    })

@app.route('/api/query', methods=['POST'])
def query():
    data = request.json
    query_text = data.get('query', '').strip()
    doc_name = data.get('document', '')

    if not query_text or not doc_name:
        return jsonify({'error': 'Missing query or document'}), 400

    if doc_name not in graph_data:
        return jsonify({'error': 'Document not found or not ready'}), 404

    if not groq_client:
        return jsonify({'error': 'GROQ_API_KEY not configured. Chat is unavailable.'}), 500

    try:
        doc_info = graph_data[doc_name]
        query_embedding = embedding_model.encode(query_text, show_progress_bar=False)

        embeddings = np.array(doc_info['embeddings'])
        similarities = np.dot(embeddings, query_embedding) / (
            np.linalg.norm(embeddings, axis=1) * np.linalg.norm(query_embedding) + 1e-10
        )

        top_k_indices = np.argsort(similarities)[-3:][::-1]
        relevant_chunks = [doc_info['chunks'][i] for i in top_k_indices]

        # Filter by lower threshold but always include at least top 1
        high_sim_indices = [i for i in top_k_indices if similarities[i] > 0.1]
        if high_sim_indices:
            top_k_indices = high_sim_indices
            relevant_chunks = [doc_info['chunks'][i] for i in top_k_indices]

        # Always use top chunks even if similarity is low
        if not relevant_chunks:
            return jsonify({
                'answer': 'Unable to process this query. Please try with a different document or question.',
                'sources': [],
                'confidence': 0.0
            })

        context = "\n".join(relevant_chunks)

        message = groq_client.chat.completions.create(
            model="llama-3.3-70b-versatile",
            max_tokens=500,
            messages=[
                {"role": "user", "content": f"""Based on this context, answer the question concisely.

Context: {context}

Question: {query_text}

Answer:"""}
            ]
        )

        answer = message.choices[0].message.content.strip()

        return jsonify({
            'answer': answer,
            'sources': [f"Chunk {idx+1}" for idx in top_k_indices],
            'confidence': float(max(similarities[top_k_indices]))
        })

    except Exception as e:
        logger.error(f"Query error: {str(e)}")
        return jsonify({'error': str(e)}), 500

@app.route('/graph-image/<filename>')
def get_graph_image(filename):
    graph_path = GRAPH_DATA_FOLDER / f"{filename}_graph.png"
    if graph_path.exists():
        return send_file(graph_path, mimetype='image/png')
    return jsonify({'error': 'Graph not found'}), 404

@app.route('/api/delete/<filename>', methods=['DELETE'])
def delete_document(filename):
    if filename in documents_state:
        del documents_state[filename]
    if filename in graph_data:
        del graph_data[filename]

    file_path = UPLOAD_FOLDER / filename
    if file_path.exists():
        file_path.unlink()

    graph_path = GRAPH_DATA_FOLDER / f"{filename}_graph.png"
    if graph_path.exists():
        graph_path.unlink()

    return jsonify({'success': True})

if __name__ == '__main__':
    port = int(os.getenv('PORT', 7860))
    app.run(host='0.0.0.0', port=port, debug=False, threaded=True)