import gradio as gr
import pandas as pd
import os
import json
from typing import List, Union
from pydantic import BaseModel, Field
import docx2txt
import extract_msg
import base64
from openai import OpenAI
from dotenv import load_dotenv

load_dotenv()
# === Константы === кек
DATA_DIR = "./"
MODEL_NAME = os.getenv("MODEL_NAME")
# === Инициализация OpenAI клиента ===
client = OpenAI(base_url="https://jane-chat-ui.ru.tuna.am/api")

# === Pydantic-модели ===
class ProductInfo(BaseModel):
    supplier: str = Field(..., description="Наименование поставщика/ код/ ИНН")
    brand: str = Field(..., description="Бренд/ Brand")
    category: str = Field(..., description="Категория/ Category")
    name: str = Field(..., description="Наименование/ Name")
    ean: str = Field(..., description="EAN")
    article_number: str = Field(..., description="Артикул для новых/ Article number for new products")
    model_number: str = Field(..., description="Номер модели/ Model number")
    currency: str = Field(..., description="Валюта/ Currency")
    delivery_date: str = Field(..., description="Дата поставки/ Delivery date в формате ДД.ММ.ГГ")
    quantity: int = Field(..., description="Кол-во/ Quantity")
    price_10_days: float = Field(..., description="Остр. 10 дн. Цена/ Price without VAT for ext. partners")
    price_30_days: float = Field(..., description="Остр. 30 дн. Цена/ Price without VAT for ext. partners")
    price_90_days: float = Field(..., description="Остр. 90 дн. Цена/ Price without VAT for ext. partners")
    sim_type: str = Field(..., description="Тип sim, например NANO+ESIM")

class ProductTable(BaseModel):
    items: List[ProductInfo] = Field(..., description="Список всех позиций из таблицы")

# === Функции обработки файлов ===
def encode_image(image_path):
    with open(image_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode("utf-8")

def extract_docx(path):
    return docx2txt.process(path).strip()

def extract_xlsx(path):
    xls = pd.read_excel(path)
    return xls.to_json(orient="records", force_ascii=False)

def extract_msg_text(path):
    msg = extract_msg.Message(path)
    return msg.body or msg.subject or "Empty message"

# === Промпт ===
prmt = """Ты — аналитик, который извлекает структурированную информацию из документа. Тебе нужно найти ответы на следующие вопросы:

1. Наименование поставщика / код / ИНН
2. Бренд / Brand
3. Категория / Category
4. Наименование / Name
5. EAN
6. Артикул для новых / Article number for new products
7. Номер модели / Model number
8. Валюта / Currency
9. Дата поставки / Delivery date
10. Кол-во / Quantity
11. Цена без НДС при отсрочке 10 дней
12. Цена без НДС при отсрочке 30 дней
13. Цена без НДС при отсрочке 90 дней
14. Тип SIM

НЕ ВЫДУМЫВАЙ ИНФОРМАЦИЮ. PLEASE!
Если каких-то данных нет — напиши "не указано".
Старайся перечислить все указанные телефоны, даже если не вся информация по ним. 
Выведи столько же SKU сколько содержится в документе!

Сначала посмотри на структуру документа, выдели, проанализируй столбики, если они есть, а потом заполни json релевантной информацией.
"""

# === Обработка одного файла ===
def process_file_for_gpt(path):
    ext = os.path.splitext(path)[1].lower()

    if ext in ['.jpg', '.jpeg', '.png']:
        base64_image = encode_image(path)
        return [
            {"type": "text", "text": prmt + f"Проанализируй и выведи всю информацию с картинки {os.path.basename(path)}?"},
            {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"}}
        ]
    elif ext == '.docx':
        text = extract_docx(path)
    elif ext == '.xlsx':
        text = extract_xlsx(path)
    elif ext == '.msg':
        text = extract_msg_text(path)
    else:
        text = f"Формат {ext} пока не поддерживается"

    return {
        "role": "user",
        "content": prmt + f"\nФайл: {os.path.basename(path)}\n\n{text}"
    }

# === Главная функция анализа ===
# === Главная функция анализа ===
def analyze_file(selected_filename=None, uploaded_file=None):
    if uploaded_file is not None:
        path = uploaded_file
    else:
        path = os.path.join(DATA_DIR, selected_filename)

    message = process_file_for_gpt(path)

    # Обращение к GPT для парсинга таблицы
    response = client.beta.chat.completions.parse(
        model=MODEL_NAME,
        response_format=ProductTable,
        messages=[message] if isinstance(message, dict) else [{"role": "user", "content": message}]
    )

    parsed_json = response.choices[0].message.parsed.model_dump()

    # Обращение к GPT для саммари
    summary_response = client.chat.completions.create(
        model=MODEL_NAME,
        messages=[
            {"role": "system", "content": "Ты ассистент, делающий краткое саммари для аналитика на основе таблицы с товарами."},
            {"role": "user", "content": f"Вот список товаров в JSON:\n{json.dumps(parsed_json, ensure_ascii=False)}\n\nСделай краткое саммари: сколько позиций, какие бренды, какие категории и диапазон цен."}
        ]
    )

    summary_text = summary_response.choices[0].message.content.strip()
    df = pd.DataFrame(parsed_json["items"])
    return summary_text, df

# === Gradio-интерфейс ===
def get_excel_files():
    return [f for f in os.listdir(DATA_DIR) if os.path.isfile(os.path.join(DATA_DIR, f)) and f.lower().endswith(('.xlsx', '.docx', '.msg', '.jpg', '.jpeg', '.png'))]

with gr.Blocks() as demo:
    gr.Markdown("## 🤖 Сервис Raft для анализа документов и картинок")
    gr.Markdown("### Выберите документ из пула или загрузить свой")

    with gr.Row():
        file_dropdown = gr.Dropdown(label="Выберите файл из папки", choices=get_excel_files(), interactive=True)
        file_upload = gr.File(label="или загрузите свой файл", file_types=['.xlsx', '.docx', '.msg', '.jpg', '.jpeg', '.png'])

    analyze_button = gr.Button("🔍 Начать анализ")

    summary_output = gr.Textbox(label="📄 Аналитика таблицы из документа", lines=5)
    table_output = gr.Dataframe(label="📋 Таблица товаров из документа")

    analyze_button.click(fn=analyze_file, inputs=[file_dropdown, file_upload], outputs=[summary_output, table_output])

demo.launch(ssr_mode=False)