import httpx
from langchain.chains.summarize import load_summarize_chain
from langchain.output_parsers import PydanticOutputParser
from langchain.schema import HumanMessage
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.chat_models import ChatOpenAI
from langchain_community.document_loaders import UnstructuredURLLoader
from langchain_core.prompts import PromptTemplate
from langchain_openai import OpenAI
from pydantic import BaseModel, Field


def url2News(url):
    text_splitter = RecursiveCharacterTextSplitter(separators=["正文", "撰稿"], chunk_size=1000, chunk_overlap=20,
                                                   length_function=len)

    url_loader = UnstructuredURLLoader([url])
    data = url_loader.load_and_split(text_splitter=text_splitter)
    return data[1:2]


def news2script(news):
    llm = OpenAI(openai_api_key="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", max_tokens=1500)

    prompt_template = """总结这段新闻的内容:
    
    
    "{text}"
    
    
    总结:"""
    chinese_prompt = PromptTemplate(template=prompt_template, input_variables=["text"])
    chain = load_summarize_chain(llm, prompt=chinese_prompt)
    summary = chain.run(news)

    openai_chat = ChatOpenAI(model_name="gpt-3.5-turbo-16k",
                             openai_api_key="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")

    template = """\
        我将给你一段新闻的概括,请按照要求把这段新闻改写成郭德纲和于谦的对口相声剧本.
        新闻:"{新闻}"
        要求:"{要求}"
        {output_instructions}
    """
    parser = PydanticOutputParser(pydantic_object=XiangSheng)

    prompt = PromptTemplate(template=template, input_variables=["新闻", "要求"],
                            partial_variables={"output_instructions": parser.get_format_instructions()})

    msg = [HumanMessage(content=prompt.format(新闻=summary,
                                              要求="风趣幽默,十分讽刺,剧本对话角色为郭德纲和于谦,以他们的自我介绍为开头"))]
    res = openai_chat(msg)
    xiangsheng = parser.parse(res.content)
    return xiangsheng


class Line(BaseModel):
    character: str = Field(description="说这句台词的角色名字")
    content: str = Field(description="台词的具体内容,其中不再包括角色的名字")


class XiangSheng(BaseModel):
    script: list[Line] = Field(description="一段相声的台词剧本")


def url2xiangsheng(url):
    doc = url2News(url)
    print(doc)
    xiangsheng = news2script(doc)
    print(xiangsheng)
    return xiangsheng