import pandas as pd
import json
import sys
import os
import fitz  # PyMuPDF
from docx import Document
import requests
import argparse

def extract_text_from_txt(file_path):
    with open(file_path, 'r', encoding='utf-8') as file:
        return file.read()

def extract_text_from_xlsx(file_path):
    xlsx_data = pd.read_excel(file_path)
    text = '\n'.join(xlsx_data.astype(str).agg(' '.join, axis=1))
    return text

def extract_text_from_docx(file_path):
    doc = Document(file_path)
    return '\n'.join([para.text for para in doc.paragraphs])

def extract_text_from_pdf(file_path):
    text = ""
    with fitz.open(file_path) as doc:
        for page in doc:
            text += page.get_text()
    return text

def summarize_text(text, model_name, prompt, context):
    url = "http://localhost:11434/api/generate"

    #if context.startswith('[') and context.endswith(']'):
    #    context = context[1:-1].strip()
    
    payload = {
        "model": model_name, 
        #"context": f"{context}",
        #"prompt": f"{prompt}\n\n\n\n{text}",
        "prompt": f"Document: {text}\n\nQuestion: {prompt}",
        "system": "You are an AI assistant that answers questions based on the given document.",
        "stream": True,
        "options": {
            "temperature": 0.8
        },
        #"keep_alive": 600  # Keep model in memory for 10 minutes
        "keep_alive": -1  # Keep model in memory indefinitely
    }

    if context:
        try:
            payload["context"] = json.loads(context)
        except json.JSONDecodeError:
            print("Warning: Invalid JSON in context. Proceeding without context.")

    payload = json.dumps(payload)
    
    headers = {
        'Content-Type': 'application/json'
    }
    
    with requests.post(url, headers=headers, data=payload, stream=True) as response:
        for line in response.iter_lines():
            if line:
                decoded_line = line.decode('utf-8')
                try:
                    json_response = json.loads(decoded_line)
                    if 'response' in json_response:
                        yield json_response['response']
                    if 'context' in json_response:
                        yield f"\n\n<Context>{json_response['context']}"
                except json.JSONDecodeError:
                    continue

def process_document(file_path, model_name, prompt, context):
    if file_path.lower().endswith('.pdf'):
        extracted_text = extract_text_from_pdf(file_path)
    elif file_path.lower().endswith('.docx'):
        extracted_text = extract_text_from_docx(file_path)
    elif file_path.lower().endswith('.xlsx'):
        extracted_text = extract_text_from_xlsx(file_path)
    elif file_path.lower().endswith('.txt'):
        extracted_text = extract_text_from_txt(file_path)
    else:
        raise ValueError('Unsupported file format')
    
    #json_text = json.dumps({"text": extracted_text})
    json_text = extracted_text
    for chunk in summarize_text(json_text, model_name, prompt, context):
        yield chunk

def validate_document_path(path):
    if not os.path.exists(path):
        raise argparse.ArgumentTypeError(f"No such file or directory: '{path}'")
    valid_extensions = {'.pdf', '.docx', '.xlsx', '.txt'}
    _, ext = os.path.splitext(path)
    if ext.lower() not in valid_extensions:
        raise argparse.ArgumentTypeError("File must be one of the following types: pdf, docx, xlsx, txt.")
    return path

# The rest of the functions (validate_model_name, get_user_input, read_context_file, main) remain unchanged

def validate_model_name(model):
    url = "http://localhost:11434/api/tags"
    response = requests.get(url)
    if response.status_code == 200:
        models = [m['name'] for m in response.json()['models']]
        if model not in models:
            raise argparse.ArgumentTypeError(f"The model {model} is not available. Available models: {', '.join(models)}")
    else:
        print("Warning: Unable to validate model name. Proceeding with the provided model.")
    return model

def get_user_input(prompt, default):
    user_input = input(f"{prompt} (default: {default}): ").strip()
    return user_input if user_input else default

def read_context_file(file_path):
    try:
        with open(file_path, 'r') as file:
            return file.read().strip()
    except FileNotFoundError:
        print(f"Warning: Context file {file_path} not found. Using empty context.")
        return ""

def main():
    parser = argparse.ArgumentParser(description="Summarize a document using Ollama")
    parser.add_argument("model_name", nargs='?', help="Name of the Ollama model to use")
    parser.add_argument("prompt_file", nargs='?', help="Path to the file containing the prompt")
    parser.add_argument("documentpath", nargs='?', help="Path to the document")
    parser.add_argument("context_file", nargs='?', default="", help="Path to the context file (optional)")

    args = parser.parse_args()

    if not args.model_name:
        print("No command-line arguments provided. Using interactive mode.")
        args.model_name = get_user_input("Enter the model name", "llama3.2")
        args.prompt_file = get_user_input("Enter the path to the prompt file", "c:\\tmp\\prompt.txt")
        args.documentpath = get_user_input("Enter the document path", "c:\\tmp\\disclosure.docx")
        args.context_file = get_user_input("Enter the path to the context file (optional)", "")

    args.documentpath = validate_document_path(args.documentpath)

    try:
        with open(args.prompt_file, 'r') as prompt_file:
            prompt = prompt_file.read().strip()
    except FileNotFoundError:
        print(f"Error: Prompt file {args.prompt_file} not found.")
        sys.exit(1)

    if args.context_file:
        try:
            with open(args.context_file, 'r') as context_file:
                context = context_file.read().strip()
        except FileNotFoundError:
            print(f"Error: Context file {args.context_file} not found.")
            sys.exit(1)
    else:
        context = ""

    #print("\nSummary:")
    for chunk in process_document(args.documentpath, args.model_name, prompt, context):
        print(chunk, end='', flush=True)

if __name__ == "__main__":
    main()