import re
from collections import OrderedDict
from typing import List
import json

from inspect_ai.model import (
    ChatMessageSystem,
    ChatMessageUser,
    GenerateConfig,
    ResponseSchema,
)
from inspect_ai.model import Model, get_model
from inspect_ai.util import json_schema


citations_schema = json_schema(
    {
        "type": "object",
        "properties": {
            "citations": {
                "type": "array",
                "items": {"type": "string"},
            }
        },
        "required": ["citations"],
    }
)
citations_schema = ResponseSchema(
    name="citations",
    json_schema=citations_schema,
)

sentences_schema = json_schema(
    {
        "type": "object",
        "properties": {
            "sentences": {
                "type": "array",
                "items": {"type": "string"},
            }
        },
        "required": ["sentences"],
    }
)
sentences_schema = ResponseSchema(
    name="sentences",
    json_schema=sentences_schema,
)


async def extract_citations(scorer_model, sentence: str) -> List[str]:
    model = get_model(scorer_model)
    prompt = "Extract all inline citations from the following text. Keep the inline citations wrapped in parentheses or square brackets, as they appear in the text. Return your result in JSON format with a single key `citations` which is a list of strings."
    config = GenerateConfig(response_schema=citations_schema, temperature=0.5)
    resp = await model.generate(
        prompt + "\n" + sentence,
        config=config,
    )
    o = resp.completion
    o = o[o.index("{") : o.rindex("}") + 1]
    return json.loads(o)["citations"]


async def split_sentences(scorer_model, text: str) -> List[str]:
    model = get_model(scorer_model)
    prompt = "Split the following text into sentences. Make sure to preserve the inline citations.Return your result in JSON format with a single key `sentences` which is a list of strings."
    config = GenerateConfig(response_schema=sentences_schema, temperature=0.5)
    resp = await model.generate(
        prompt + "\n" + text,
        config=config,
    )
    o = resp.completion
    o = o[o.index("{") : o.rindex("}") + 1]
    return json.loads(o)["sentences"]


async def propagate_citations(scorer_model, sentences: List[str]) -> List[str]:
    out = sentences[:]
    next_cites = None
    for i in range(len(sentences) - 1, -1, -1):
        cites = await extract_citations(scorer_model, sentences[i])
        if cites:
            next_cites = " ".join(cites)
        elif next_cites:
            out[i] = sentences[i].rstrip() + " " + next_cites
    return out


async def add_missing_citations(scorer_model, text: str):
    paragraphs = re.split(r"\n\s*\n", text.strip(), flags=re.M)
    processed = []
    for p in paragraphs:
        sents = await split_sentences(scorer_model, p)
        sents = await propagate_citations(scorer_model, sents)
        processed.extend(sents)
    return processed
