#
# SPDX-FileCopyrightText: Hadad <hadad@linuxmail.org>
# SPDX-License-Identifier: CC-BY-NC-SA-4.0
#

import gradio as gr
import requests
import json
import os
import random
import time
import pytesseract
import pdfplumber
import docx
import pandas as pd
import pptx
import fitz
import io
import uuid
import concurrent.futures
import itertools
import threading

from openai import OpenAI

from optillm.cot_reflection import cot_reflection
from optillm.leap import leap
from optillm.plansearch import plansearch
from optillm.reread import re2_approach
from optillm.rto import round_trip_optimization
from optillm.self_consistency import advanced_self_consistency_approach
from optillm.z3_solver import Z3SymPySolverSystem

from pathlib import Path
from PIL import Image
from pptx import Presentation

os.system("apt-get update -q -y && apt-get install -q -y tesseract-ocr tesseract-ocr-eng tesseract-ocr-ind libleptonica-dev libtesseract-dev")

LINUX_SERVER_HOSTS = [host for host in json.loads(os.getenv("LINUX_SERVER_HOST", "[]")) if host]
LINUX_SERVER_HOSTS_MARKED = set()
LINUX_SERVER_HOSTS_ATTEMPTS = {}

LINUX_SERVER_PROVIDER_KEYS = [key for key in json.loads(os.getenv("LINUX_SERVER_PROVIDER_KEY", "[]")) if key]
LINUX_SERVER_PROVIDER_KEYS_MARKED = set()
LINUX_SERVER_PROVIDER_KEYS_ATTEMPTS = {}

AI_TYPES = {f"AI_TYPE_{i}": os.getenv(f"AI_TYPE_{i}") for i in range(1, 7)}
RESPONSES = {f"RESPONSE_{i}": os.getenv(f"RESPONSE_{i}") for i in range(1, 10)}

MODEL_MAPPING = json.loads(os.getenv("MODEL_MAPPING", "{}"))
MODEL_CONFIG = json.loads(os.getenv("MODEL_CONFIG", "{}"))
MODEL_CHOICES = list(MODEL_MAPPING.values()) if MODEL_MAPPING else []
DEFAULT_CONFIG = json.loads(os.getenv("DEFAULT_CONFIG", "{}"))

META_TAGS = os.getenv("META_TAGS")

ALLOWED_EXTENSIONS = json.loads(os.getenv("ALLOWED_EXTENSIONS"))

ACTIVE_CANDIDATE = None

def get_available_items(items, marked):
    available = [item for item in items if item not in marked]
    random.shuffle(available)
    return available

def marked_item(item, marked, attempts):
    marked.add(item)
    attempts[item] = attempts.get(item, 0) + 1
    if attempts[item] >= 3:
        def remove_fail():
            marked.discard(item)
            if item in attempts:
                del attempts[item]
        threading.Timer(300, remove_fail).start()

class SessionWithID(requests.Session):
    def __init__(self):
        super().__init__()
        self.session_id = str(uuid.uuid4())

def create_session():
    return SessionWithID()

def get_model_key(display_name):
    return next((k for k, v in MODEL_MAPPING.items() if v == display_name), list(MODEL_MAPPING.keys())[0] if MODEL_MAPPING else MODEL_CHOICES[0])

def extract_file_content(file_path):
    ext = Path(file_path).suffix.lower()
    content = ""
    try:
        if ext == ".pdf":
            with pdfplumber.open(file_path) as pdf:
                for page in pdf.pages:
                    text = page.extract_text()
                    if text:
                        content += text + "\n"
                    tables = page.extract_tables()
                    if tables:
                        for table in tables:
                            table_str = "\n".join([", ".join(row) for row in table if row])
                            content += "\n" + table_str + "\n"
        elif ext in [".doc", ".docx"]:
            doc = docx.Document(file_path)
            for para in doc.paragraphs:
                content += para.text + "\n"
        elif ext in [".xlsx", ".xls"]:
            df = pd.read_excel(file_path)
            content += df.to_csv(index=False)
        elif ext in [".ppt", ".pptx"]:
            prs = Presentation(file_path)
            for slide in prs.slides:
                for shape in slide.shapes:
                    if hasattr(shape, "text") and shape.text:
                        content += shape.text + "\n"
        elif ext in [".png", ".jpg", ".jpeg", ".tiff", ".bmp", ".gif", ".webp"]:
            try:
                pytesseract.pytesseract.tesseract_cmd = "/usr/bin/tesseract"
                image = Image.open(file_path)
                text = pytesseract.image_to_string(image)
                content += text + "\n"
            except Exception as e:
                content += f"{e}\n"
        else:
            content = Path(file_path).read_text(encoding="utf-8")
    except Exception as e:
        content = f"{file_path}: {e}"
    return content.strip()

def process_ai_response(ai_text):
    try:
        result = round_trip_optimization(ai_text)
        result = re2_approach(result)
        result = cot_reflection(result)
        result = advanced_self_consistency_approach(result)
        result = plansearch(result)
        result = leap(result)
        solver = Z3SymPySolverSystem()
        result = solver.solve(result)
        return result
    except Exception:
        return ai_text

def fetch_response(host, provider_key, selected_model, messages, model_config, session_id):
    try:
        client = OpenAI(base_url=host, api_key=provider_key)
        data = {"model": selected_model, "messages": messages, **model_config}
        response = client.chat.completions.create(extra_body={"optillm_approach": "rto|re2|cot_reflection|self_consistency|plansearch|leap|z3|bon|moa|mcts|mcp|router|privacy|executecode|json", "session_id": session_id}, **data)
        ai_text = response.choices[0].message.content if response.choices and response.choices[0].message and response.choices[0].message.content else RESPONSES["RESPONSE_2"]
        return process_ai_response(ai_text)
    except Exception:
        marked_item(provider_key, LINUX_SERVER_PROVIDER_KEYS_MARKED, LINUX_SERVER_PROVIDER_KEYS_ATTEMPTS)
        marked_item(host, LINUX_SERVER_HOSTS_MARKED, LINUX_SERVER_HOSTS_ATTEMPTS)
        raise

def chat_with_model(history, user_input, selected_model_display, sess):
    global ACTIVE_CANDIDATE
    if not get_available_items(LINUX_SERVER_PROVIDER_KEYS, LINUX_SERVER_PROVIDER_KEYS_MARKED) or not get_available_items(LINUX_SERVER_HOSTS, LINUX_SERVER_HOSTS_MARKED):
        return RESPONSES["RESPONSE_3"]
    if not hasattr(sess, "session_id"):
        sess.session_id = str(uuid.uuid4())
    selected_model = get_model_key(selected_model_display)
    model_config = MODEL_CONFIG.get(selected_model, DEFAULT_CONFIG)
    messages = [{"role": "user", "content": user} for user, _ in history]
    messages += [{"role": "assistant", "content": assistant} for _, assistant in history if assistant]
    messages.append({"role": "user", "content": user_input})
    if ACTIVE_CANDIDATE is not None:
        try:
            return fetch_response(ACTIVE_CANDIDATE[0], ACTIVE_CANDIDATE[1], selected_model, messages, model_config, sess.session_id)
        except Exception:
            ACTIVE_CANDIDATE = None
    available_keys = get_available_items(LINUX_SERVER_PROVIDER_KEYS, LINUX_SERVER_PROVIDER_KEYS_MARKED)
    available_servers = get_available_items(LINUX_SERVER_HOSTS, LINUX_SERVER_HOSTS_MARKED)
    candidates = [(host, key) for host in available_servers for key in available_keys]
    random.shuffle(candidates)
    with concurrent.futures.ThreadPoolExecutor(max_workers=len(candidates)) as executor:
        futures = {executor.submit(fetch_response, host, key, selected_model, messages, model_config, sess.session_id): (host, key) for (host, key) in candidates}
        for future in concurrent.futures.as_completed(futures):
            try:
                result = future.result()
                ACTIVE_CANDIDATE = futures[future]
                for f in futures:
                    if f is not future:
                        f.cancel()
                return result
            except Exception:
                continue
    return RESPONSES["RESPONSE_2"]

def respond(multi_input, history, selected_model_display, sess):
    message = {"text": multi_input.get("text", "").strip(), "files": multi_input.get("files", [])}
    if not message["text"] and not message["files"]:
        yield history, gr.MultimodalTextbox(value=None, interactive=True), sess
        return
    combined_input = ""
    for file_item in message["files"]:
        file_path = file_item["name"] if isinstance(file_item, dict) and "name" in file_item else file_item
        file_content = extract_file_content(file_path)
        combined_input += f"{Path(file_path).name}\n\n{file_content}\n\n"
    if message["text"]:
        combined_input += message["text"]
    history.append([combined_input, ""])
    ai_response = chat_with_model(history, combined_input, selected_model_display, sess)
    history[-1][1] = ""
    def convert_to_string(data):
        if isinstance(data, (str, int, float)):
            return str(data)
        elif isinstance(data, bytes):
            return data.decode("utf-8", errors="ignore")
        elif isinstance(data, (list, tuple)):
            return "".join(map(convert_to_string, data))
        elif isinstance(data, dict):
            return json.dumps(data, ensure_ascii=False)
        else:
            return repr(data)
    for character in ai_response:
        history[-1][1] += convert_to_string(character)
        time.sleep(0.0001)
        yield history, gr.MultimodalTextbox(value=None, interactive=True), sess

def change_model(new_model_display):
    return [], create_session(), new_model_display

with gr.Blocks(fill_height=True, fill_width=True, title=AI_TYPES["AI_TYPE_4"], head=META_TAGS) as jarvis:
    user_history = gr.State([])
    user_session = gr.State(create_session())
    selected_model = gr.State(MODEL_CHOICES[0])
    chatbot = gr.Chatbot(label=AI_TYPES["AI_TYPE_1"], show_copy_button=True, scale=1, elem_id=AI_TYPES["AI_TYPE_2"])
    model_dropdown = gr.Dropdown(show_label=False, choices=MODEL_CHOICES, value=MODEL_CHOICES[0])
    with gr.Row():
        msg = gr.MultimodalTextbox(show_label=False, placeholder=RESPONSES["RESPONSE_5"], interactive=True, file_count="single", file_types=ALLOWED_EXTENSIONS)

    model_dropdown.change(fn=change_model, inputs=[model_dropdown], outputs=[user_history, user_session, selected_model])
    msg.submit(fn=respond, inputs=[msg, user_history, selected_model, user_session], outputs=[chatbot, msg, user_session])

jarvis.launch(show_api=False, max_file_size="1mb")