from typing import Literal
from lxml import etree
import re
from dataclasses import dataclass
import openai
import tiktoken

from tocky.extractor import AbstractExtractor, TocEntry, TocResponse
from tocky.extractor.formats import build_system_prompt, format_toc, process_extracted_output
from tocky.ocr.printer import print_ocr
from tocky.utils import avg_ocr_conf
from tocky.utils.ia import get_djvu_by_leaf_nums, get_ia_metadata, get_page_scan, ocaid_to_djvu_url


class BadOcrOnToc(Exception):
  pass


SYSTEM_PROMPT = """
You are a librarian extracting table of contents data in a structured, semantic format.

{format_instructions}

### Important Instructions:
- Fix any typos or bad characters generated by the OCR.
- Change any text in ALL UPPERCASE to be in Normal Title Case (e.g. "The Limit of a Function").
- Use the `label` field to store unimportant data like numerals (e.g. "Chapter 1", "Section 2.3").

### Examples:

{PROMPT_SAMPLES[basic]}

{PROMPT_SAMPLES[no_titles]}

{PROMPT_SAMPLES[nested]}
"""

PREV_TOC_PROMPT = """
The last table of contents was too long too fit in your context. Continue extracting the table of contents of "{book_title}".

Here is the end of what you last output. Do not output this again, though.

{prev_toc}

And here is the rest of the OCR text:
```
{ocr_text}
```
"""

TOC_PROMPT = """
Extract the table of contents from this OCR text of "{book_title}":

```
{ocr_text}
```
"""


@dataclass
class AiExtractorOptions:
  redo_ocr: bool = True
  ocr_engine: Literal['easyocr', 'tesseract', 'azure'] = 'easyocr'
  model: str = "gpt-4o-mini"
  max_sent_tokens: int = 1_000
  """
  GPT 4o mini can handle up to 128k input tokens and 16k output tokens.
  TODO: Experiment with larger input sizes.
  """
  extraction_format: Literal['json', 'markdown'] = 'json'


class AiExtractor(AbstractExtractor[AiExtractorOptions]):
  name = 'ai_extractor'

  def __init__(self):
    super().__init__()
    self.P = AiExtractorOptions()

  def extract(self, ocaid: str, detector_result: list[int]) -> list[TocEntry]:
    djvu_xml_to_fetch = set(detector_result) - set(self.S.ocr_cache.keys())
    if djvu_xml_to_fetch:
      djvu_url = ocaid_to_djvu_url(ocaid)
      start = min(djvu_xml_to_fetch)
      end = max(djvu_xml_to_fetch)
      for leaf_num, elem in get_djvu_by_leaf_nums(djvu_url, start, end):
        if leaf_num in djvu_xml_to_fetch:
          self.S.ocr_cache[leaf_num] = etree.tostring(elem, encoding='unicode')

    if self.P.redo_ocr:
      for leaf_num in detector_result:
        self.S.ocr_cache[leaf_num] = self.redo_ocr(ocaid, leaf_num, self.S.ocr_cache[leaf_num])

    self.toc_raw_ocr = [
      print_ocr(self.S.ocr_cache[leaf_num])
      for leaf_num in detector_result
    ]

    if re.search(r'([A-Za-z]{25,}|\beee+\b)', '\n'.join(self.toc_raw_ocr), flags=re.MULTILINE):
      raise BadOcrOnToc("Bad OCR on TOC")

    self.toc_response = self.extract_structured_toc(self.toc_raw_ocr, get_ia_metadata(ocaid)['metadata']['title'])

    return self.toc_response.toc

  def redo_ocr(self, ocaid: str, leaf_num: int, djvu_xml: str) -> str:
    from tocky.ocr import ocr_djvu_page

    root = etree.fromstring(djvu_xml)
    if root.xpath('.//HIDDENTEXT/@x-re-ocrd') == ['true']:
      return djvu_xml

    new_ocr = ocr_djvu_page(get_page_scan(ocaid, leaf_num), engine=self.P.ocr_engine)
    old_ocr_el = root.find('.//HIDDENTEXT')
    new_ocr_el = etree.fromstring(new_ocr).find('.//HIDDENTEXT')
    assert old_ocr_el and new_ocr_el
    if (avg_ocr_conf(new_ocr_el) or 100) > (avg_ocr_conf(old_ocr_el) or 0):
      root.replace(old_ocr_el, new_ocr_el)

    return etree.tostring(root, encoding='unicode')

  def extract_structured_toc(
    self,
    pages_ocr: list[str],
    book_title: str,
  ) -> TocResponse:
    structured_toc: list[TocEntry] = []
    prompt_tokens = 0
    completion_tokens = 0
    chunks = ['']
    for page_ocr in pages_ocr:
      extended_chunk = chunks[-1] + '\n' + page_ocr
      if len(tiktoken.encoding_for_model(self.P.model).encode(extended_chunk)) > self.P.max_sent_tokens:
        chunks.append(page_ocr)
      else:
        chunks[-1] += '\n' + page_ocr
    for chunk in chunks:
      toc_response = self.extract_structured_toc_chunk(chunk, book_title, prev_toc=structured_toc)
      structured_toc += toc_response.toc
      prompt_tokens += toc_response.prompt_tokens
      completion_tokens += toc_response.completion_tokens
    
    return TocResponse(structured_toc, prompt_tokens, completion_tokens)

  def extract_structured_toc_chunk(
    self,
    ocr_text: str,
    book_title: str,
    prev_toc: list[TocEntry] | None = None,
  ) -> TocResponse:
    if prev_toc:
      prev_toc_slice = prev_toc[-10:]
      if self.P.extraction_format == 'json':
        prev_toc_formatted = f'```json\n{format_toc(prev_toc_slice, self.P.extraction_format)}\n```'
      elif self.P.extraction_format == 'markdown':
        prev_toc_formatted = f'```\n{format_toc(prev_toc_slice, self.P.extraction_format)}\n```'
      else:
        raise ValueError(f"Unknown extraction format: {self.P.extraction_format}")

      message = PREV_TOC_PROMPT.format(
        book_title=book_title,
        prev_toc=prev_toc_formatted,
        ocr_text=ocr_text,
      )
    else:
      message = TOC_PROMPT.format(
        book_title=book_title,
        ocr_text=ocr_text,
      )
    print('openai request', message[0:500] + '...')

    system_prompt = build_system_prompt(SYSTEM_PROMPT, self.P.extraction_format)
    completion = openai.chat.completions.create(
      model=self.P.model,
      messages=[
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": message},
      ],
      # max_tokens=1024,
      n=1,
      stop=None,
      temperature=0.5,
    )

    assert completion.choices[0].message.content
    assert completion.usage

    toc = process_extracted_output(
      completion.choices[0].message.content,
      self.P.extraction_format,
    )
    return TocResponse(
        toc,
        completion.usage.prompt_tokens,
        completion.usage.completion_tokens,
    )
