import html
import logging
import os
import threading
from dataclasses import dataclass, field
from functools import lru_cache
from pathlib import Path

import spaces  
import gradio as gr

from lexquery.config import GEMINI_MODEL, LANGUAGES
from lexquery.models import LocalModels
from lexquery.pipeline import Engine
from lexquery.schema import Result

logging.basicConfig(level=logging.INFO, format='%(levelname)s %(name)s: %(message)s')


@lru_cache(maxsize=1)
def load_models():
    return LocalModels()


@dataclass
class Session:
    engine: Engine = field(default_factory=lambda: Engine(load_models()))
    chat: list = field(default_factory=list)
    lock: object = field(default_factory=threading.RLock)


_sessions = {}
_sessions_lock = threading.RLock()


def session(request):
    if not request.session_hash:
        raise gr.Error('Open the app in a browser session first.')
    with _sessions_lock:
        if request.session_hash not in _sessions:
            _sessions[request.session_hash] = Session()
        return _sessions[request.session_hash]


def cleanup_id(session_id):
    with _sessions_lock:
        _sessions.pop(session_id, None)


def unload(request: gr.Request):
    cleanup_id(request.session_hash)


def initialize(request: gr.Request):
    session(request)
    return request.session_hash


def safe_error(exc):
    text = str(exc)
    for name in ('GOOGLE_API_KEY', 'GEMINI_API_KEY'):
        key = os.getenv(name)
        if key:
            text = text.replace(key, '[key hidden]')
    return text


def read_uploads(files, request: gr.Request, progress=gr.Progress()):
    item = session(request)
    with item.lock:
        keep, messages = set(), []
        old_ids = set(item.engine.documents)
        for i, file in enumerate(files or []):
            path = Path(file)
            progress(i / max(1, len(files)), desc='Reading ' + path.name)
            try:
                doc = item.engine.ingest(path.name, path.read_bytes())
                keep.add(doc.id)
                messages.append(f'**{html.escape(doc.name)}** · {len(doc.pages)} pages · ready')
                uncertain = [str(p.number) for p in doc.pages if p.warnings]
                if uncertain:
                    messages.append('Extraction notes on pages ' + ', '.join(uncertain) + '; see Extracted text.')
            except Exception as exc:
                logging.exception('PDF extraction failed')
                messages.append(f'**{html.escape(path.name)}** · {html.escape(safe_error(exc))}')
        for doc_id in set(item.engine.documents) - keep:
            del item.engine.documents[doc_id]
        if old_ids != keep:
            item.engine.invalidate()
        item.engine.history.clear()
        item.chat.clear()
        choices = [(d.name, d.id) for d in item.engine.documents.values()]
        ids = [i for _, i in choices]
        progress(1, desc='Finished')
        return ('\n\n'.join(messages) or 'Upload a PDF to begin.',
                gr.Dropdown(choices=choices, value=ids, multiselect=True),
                gr.Dropdown(choices=choices, value=ids[0] if ids else None),
                [], '', '')


def extracted_text(doc_id, request: gr.Request):
    item = session(request)
    with item.lock:
        doc = item.engine.documents.get(doc_id)
        if not doc:
            return ''
        return '\n\n'.join(f'PAGE {p.number} · {p.language}\n{p.original}\n' +
                           ('\nNotes: ' + '; '.join(p.warnings) if p.warnings else '') for p in doc.pages)


def result_markdown(result):
    paragraphs = [html.escape(text) for text in result.displayed] if result.status in (
        'checked', 'conversational') else []
    for note in result.notes:
        paragraphs.append(html.escape(safe_error(note)))
    return '\n\n'.join(paragraphs) or 'I could not find a supported answer in the selected documents.'


def ask(question, ids, language, request: gr.Request):
    item = session(request)
    with item.lock:
        if not question.strip():
            return item.chat, ''
        result = item.engine.present(item.engine.ask(question, ids or []), language)
        item.chat.extend([{'role': 'user', 'content': question},
                          {'role': 'assistant', 'content': result_markdown(result)}])
        return item.chat, ''

@spaces.GPU
def ask(question, ids, language, request: gr.Request):
    item = session(request)
    with item.lock:
        if not question.strip():
            return item.chat, ''
        result = item.engine.present(item.engine.ask(question, ids or []), language)
        item.chat.extend([{'role': 'user', 'content': question},
                          {'role': 'assistant', 'content': result_markdown(result)}])
        return item.chat, ''


@spaces.GPU
def summarize(ids, language, request: gr.Request, progress=gr.Progress()):
    item = session(request)
    with item.lock:
        if not ids:
            return 'Upload and select a PDF first.'
        sections = []


def summarize(ids, language, request: gr.Request, progress=gr.Progress()):
    item = session(request)
    with item.lock:
        if not ids:
            return 'Upload and select a PDF first.'
        sections = []
        for i, doc_id in enumerate(ids):
            doc = item.engine.documents.get(doc_id)
            if doc is None:
                continue
            progress(i / len(ids), desc='Summarizing ' + doc.name)
            try:
                report = item.engine.report(doc_id)
                result = item.engine.present(Result.model_validate(report['summary']), language)
                section = '### ' + html.escape(doc.name) + '\n\n' + result_markdown(result)
                fields = []
                for name, facts in report['fields'].items():
                    if facts:
                        fields.append('**' + name.replace('_', ' ').title() + ':** ' +
                                      '; '.join(html.escape(f['value']) for f in facts))
                if fields:
                    section += '\n\n' + '\n\n'.join(fields)
                for note in report.get('field_notes', []):
                    section += '\n\n' + html.escape(safe_error(note))
                sections.append(section)
            except Exception as exc:
                logging.exception('Report failed')
                sections.append('### ' + html.escape(doc.name) + '\n\n' + html.escape(safe_error(exc)))
        progress(1, desc='Finished')
        return '\n\n---\n\n'.join(sections)


def clear_chat(request: gr.Request):
    item = session(request)
    with item.lock:
        item.engine.history.clear()
        item.chat.clear()
    return []


def build_app():
    with gr.Blocks(title='LexQuery', delete_cache=(3600, 3600)) as demo:
        lifecycle = gr.State(None, delete_callback=cleanup_id)
        gr.Markdown('# LexQuery\nUpload your PDFs. Ask questions. Get clear summaries.')
        with gr.Row():
            with gr.Column(scale=1, min_width=280):
                files = gr.File(label='Documents', file_types=['.pdf'], file_count='multiple', type='filepath')
                gr.Markdown('Scanned and text PDFs · English and 22 Indian languages')
                status = gr.Markdown('Upload a PDF to begin.')
                selected = gr.Dropdown(choices=[], multiselect=True, label='Use these documents')
                language = gr.Dropdown(LANGUAGES, value='English', label='Answer language')
            with gr.Column(scale=3):
                with gr.Tab('Questions'):
                    chat = gr.Chatbot(label='Conversation', height=460)
                    question = gr.Textbox(label='Your question', placeholder='What does this document say about…?')
                    with gr.Row():
                        ask_button = gr.Button('Ask', variant='primary')
                        reset = gr.Button('New conversation')
                with gr.Tab('Summaries'):
                    gr.Markdown('Summaries cover each selected document independently of your conversation.')
                    summary_button = gr.Button('Summarize documents', variant='primary')
                    summary = gr.Markdown()
                with gr.Accordion('Extracted text', open=False):
                    text_doc = gr.Dropdown(choices=[], label='Document')
                    transcript = gr.Textbox(label='Original text', lines=18, interactive=False)
        gr.Markdown('Check important findings against the original PDF.')
        files.change(read_uploads, files, [status, selected, text_doc, chat, summary, transcript])
        text_doc.change(extracted_text, text_doc, transcript)
        for trigger in (ask_button.click, question.submit):
            trigger(ask, [question, selected, language], [chat, question])
        reset.click(clear_chat, outputs=chat)
        summary_button.click(summarize, [selected, language], summary)
        demo.load(initialize, outputs=lifecycle)
        demo.unload(unload)
    return demo.queue(default_concurrency_limit=1)


if __name__ == "__main__":
    logging.info("LexQuery Gemini model: %s", GEMINI_MODEL)
    build_app().launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
    )