#
# SPDX-FileCopyrightText: Hadad <hadad@linuxmail.org>
# SPDX-License-Identifier: Apache-2.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
import httpx
import asyncio

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")

INTERNAL_AI_GET_SERVER = os.getenv("INTERNAL_AI_GET_SERVER")
INTERNAL_TRAINING_DATA = os.getenv("INTERNAL_TRAINING_DATA")

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

async def fetch_response_async(host, provider_key, selected_model, messages, model_config, session_id):
    try:
        async with httpx.AsyncClient(timeout=5) as client:
            data = {"model": selected_model, "messages": messages, **model_config}
            extra = {"optillm_approach": "rto|re2|cot_reflection|self_consistency|plansearch|leap|z3|bon|moa|mcts|mcp|router|privacy|executecode|json", "session_id": session_id}
            response = await client.post(f"{host}", json={**data, "extra_body": extra, "session_id": session_id}, headers={"Authorization": f"Bearer {provider_key}"})
            response.raise_for_status()
            resp_json = response.json()
            ai_text = resp_json["choices"][0]["message"]["content"] if resp_json.get("choices") and resp_json["choices"][0].get("message") and resp_json["choices"][0]["message"].get("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)
        raise

async def chat_with_model_async(history, user_input, selected_model_display, sess):
    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]
    if INTERNAL_TRAINING_DATA:
        messages.insert(0, {"role": "system", "content": INTERNAL_TRAINING_DATA})
    messages.append({"role": "user", "content": user_input})
    global ACTIVE_CANDIDATE
    if ACTIVE_CANDIDATE is not None:
        try:
            return await fetch_response_async(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)
    tasks = [fetch_response_async(host, key, selected_model, messages, model_config, sess.session_id) for host, key in candidates]
    for task in asyncio.as_completed(tasks):
        try:
            result = await task
            ACTIVE_CANDIDATE = next(((host, key) for host, key in candidates if host and key), None)
            return result
        except Exception:
            continue
    return RESPONSES["RESPONSE_2"]

async def respond_async(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, ""])
    if len(history) > 3:
        history[:] = history[-3:]
    ai_response = await chat_with_model_async(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)
        await asyncio.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_async, inputs=[msg, user_history, selected_model, user_session], outputs=[chatbot, msg, user_session], concurrency_limit=None, api_name=INTERNAL_AI_GET_SERVER)
jarvis.launch(max_file_size="1mb")