import os
import re
import requests
import json
from typing import Tuple, List

from omegaconf import OmegaConf

from typing import Optional
from pydantic import Field, BaseModel

from vectara_agentic.agent import Agent
from vectara_agentic.tools import ToolsFactory, VectaraToolFactory
from vectara_agentic.tools_catalog import summarize_text

from dotenv import load_dotenv
load_dotenv(override=True)

citation_description = '''
    The citation for a particular case. 
    Citation must include the volume number, reporter, and first page. For example: 253 P.2d 136.
'''

def extract_components_from_citation(citation: str) -> dict:
    citation_components = citation.split(' ')
    volume_num = citation_components[0]
    reporter = '-'.join(citation_components[1:-1]).replace('.', '').lower()
    first_page = citation_components[-1]

    if not volume_num.isdigit():
        return {}
    if not first_page.isdigit():
        return {}

    return {'volume': int(volume_num), 'reporter': reporter, 'first_page': int(first_page)}

def create_assistant_tools(cfg):

    def get_opinion_text(
            case_citation: str = Field(description = citation_description),
            summarize: bool = Field(default=True, description="if True returns case summary, otherwise the full text of the case")
            ) -> str:
        """
        Returns the full opinion/ruling text of the case, or the summary if summarize=True.
        If there is more than one opinion for the case, the type of each opinion is returned with the text, 
        and the opinions (or their summaries) are separated by semicolons (;)

        Args
            case_citation (str): the citation for a particular case. Citation must include the volume number, reporter, and first page. For example: 253 P.2d 136.
            summarize (bool): True to return just a summary of the case, False to return full case text.
        """
        citation_dict = extract_components_from_citation(case_citation)
        if not citation_dict:
            return f"Citation is invalid: {case_citation}."
        reporter = citation_dict['reporter']
        volume_num = citation_dict['volume']
        first_page = citation_dict['first_page']
        response = requests.get(f"https://static.case.law/{reporter}/{volume_num}/cases/{first_page:04d}-01.json")
        if response.status_code != 200:
            return f"Case not found; please check the citation {case_citation}."
        res = json.loads(response.text)

        if len(res["casebody"]["opinions"]) == 1:
            text = res["casebody"]["opinions"][0]["text"]
            output = text if not summarize else summarize_text(text, "law")
        else:
            output = ""
            for opinion in res["casebody"]["opinions"]:
                text = opinion["text"] if not summarize else summarize_text(opinion["text"], "law")
                output += f"Opinion type: {opinion['type']}, text: {text};"
        
        return output

    def get_case_document_pdf(
            case_citation = Field(description = citation_description)
            ) -> str:
        """
        Given a case citation, returns a valid web url to a pdf of the case record
        """
        citation_dict = extract_components_from_citation(case_citation)
        if not citation_dict:
            return f"Citation is invalid: {case_citation}."
        reporter = citation_dict['reporter']
        volume_num = citation_dict['volume']
        first_page = citation_dict['first_page']
        response = requests.get(f"https://static.case.law/{reporter}/{volume_num}/cases/{first_page:04d}-01.json")
        if response.status_code != 200:
            return f"Case not found; please check the citation {case_citation}."
        res = json.loads(response.text)
        page_number = res["first_page_order"]
        return f"https://static.case.law/{reporter}/{volume_num}.pdf#page={page_number}"

    def get_case_document_page(
            case_citation = Field(description = citation_description)
            ) -> str:
        """
        Given a case citation, returns a valid web url to a page with information about the case.
        """
        citation_dict = extract_components_from_citation(case_citation)
        if not citation_dict:
            return f"Citation is invalid: {case_citation}."
        reporter = citation_dict['reporter']
        volume_num = citation_dict['volume']
        first_page = citation_dict['first_page']
        url = f"https://case.law/caselaw/?reporter={reporter}&volume={volume_num}&case={first_page:04d}-01"
        response = requests.get(url)
        if response.status_code != 200:
            return "Case not found; please check the citation."
        return url
        
    def get_case_name(
            case_citation = Field(description = citation_description)
            ) -> Tuple[str, str]:
        """
        Given a case citation, returns its name and name abbreviation.
        """
        citation_dict = extract_components_from_citation(case_citation)
        if not citation_dict:
            return f"Citation is invalid: {case_citation}.", f"Citation is invalid: {case_citation}."
        reporter = citation_dict['reporter']
        volume_num = citation_dict['volume']
        first_page = citation_dict['first_page']
        response = requests.get(f"https://static.case.law/{reporter}/{volume_num}/cases/{first_page:04d}-01.json")
        if response.status_code != 200:
            return "Case not found", "Case not found"
        res = json.loads(response.text)
        return res["name"], res["name_abbreviation"]

    def get_cited_cases(
            case_citation = Field(description = citation_description)
            ) -> List[dict]:
        """
        Given a case citation, returns a list of cases that are cited by the opinion of this case.
        The output is a list of cases, each a dict with the citation, name and name_abbreviation of the case.
        """
        citation_dict = extract_components_from_citation(case_citation)
        if not citation_dict:
            return [f"Citation is invalid: {case_citation}."]
        reporter = citation_dict['reporter']
        volume_num = citation_dict['volume']
        first_page = citation_dict['first_page']
        response = requests.get(f"https://static.case.law/{reporter}/{volume_num}/cases/{first_page:04d}-01.json")
        if response.status_code != 200:
            return "Case not found; please check the citation."
        res = json.loads(response.text)
        citations = res["cites_to"]
        res = []
        for citation in citations[:10]:
            name, name_abbreviation = get_case_name(citation["cite"])
            res.append({
                "citation": citation["cite"],
                "name": name,
                "name_abbreviation": name_abbreviation
            })
        return res

    def validate_url(
            url = Field(description = "A web url pointing to case-law document")
        ) -> str:
        """
        Given a link, returns whether or not the link is valid.
        If it is not valid, it should not be used in any output.
        """  
        pdf_pattern = re.compile(r'^https://static.case.law/.*')
        document_pattern = re.compile(r'^https://case.law/caselaw/?reporter=.*')
        return "URL is valid" if bool(pdf_pattern.match(url)) | bool(document_pattern.match(url)) else "URL is bad"

    class QueryCaselawArgs(BaseModel):
        query: str = Field(..., description="The user query.")

    vec_factory = VectaraToolFactory(vectara_api_key=cfg.api_key, 
                                     vectara_customer_id=cfg.customer_id, 
                                     vectara_corpus_id=cfg.corpus_id)
    summarizer = 'vectara-experimental-summary-ext-2023-12-11-med-omni'

    ask_caselaw = vec_factory.create_rag_tool(
        tool_name = "ask_caselaw",
        tool_description = "A tool for asking questions about case law in Alaska. ",
        tool_args_schema = QueryCaselawArgs,
        reranker = "chain", rerank_k = 100,
        rerank_chain = [
            {
                "type": "slingshot",
                "cutoff": 0.2
            },
            {
                "type": "mmr",
                "diversity_bias": 0.1
            },
            {
                "type": "udf",
                "user_function": "max(1000 * get('$.score') - hours(seconds(to_unix_timestamp(now()) - to_unix_timestamp(datetime_parse(get('$.document_metadata.decision_date'), 'yyyy-MM-dd')))) / 24 / 365, 0)"
            }
        ],
        n_sentences_before = 2, n_sentences_after = 2, lambda_val = 0.005,
        summary_num_results = 15,
        vectara_summarizer = summarizer,
        include_citations = False,
    )

    tools_factory = ToolsFactory()
    return (
        [ask_caselaw] +
        [tools_factory.create_tool(tool) for tool in [
            get_opinion_text,
            get_case_document_pdf,
            get_case_document_page,
            get_cited_cases,
            get_case_name,
            validate_url
        ]] 
    )

def get_agent_config() -> OmegaConf:
    cfg = OmegaConf.create({
        'customer_id': str(os.environ['VECTARA_CUSTOMER_ID']),
        'corpus_id': str(os.environ['VECTARA_CORPUS_ID']),
        'corpus_key': str(os.environ['VECTARA_CORPUS_KEY']),
        'api_key': str(os.environ['VECTARA_API_KEY']),
        'examples': os.environ.get('QUERY_EXAMPLES', None),
        'demo_name': "legal-agent",
        'demo_welcome': "Welcome to the Legal Assistant demo.",
        'demo_description': "This demo can help you prepare for a court case by providing you information about past court cases in Alaska.",
    })
    return cfg

def initialize_agent(_cfg, agent_progress_callback=None):
    
    legal_assistant_instructions = """
    - You are a helpful legal assistant, with expertise in case law for the state of Alaska.
    - The ask_caselaw tool is your primary tools for finding information about cases. 
      Do not use your own knowledge to answer questions.
    - If the ask_caselaw tool responds that it does not have enough information to answer the query,
      try to rephrase the query and call the tool again.
    - If the ask_caselaw tool returns a response along with a ist of references mentioned in its response (in [N] format), 
      Format your response to focus on the main response, and use the metadata (such as citations, decision date, or case name) 
      in each relevant reference to provide more context in your response.
    - Citations include 3 components: volume number, reporter, and first page. 
      Here are some examples: '253 P.2d 136', '10 Alaska 11', '6 C.M.A. 3'
    - Never use your internal knowledge to guess citations. Only use citations information provided by a tool or the user.
    - If two cases have conflicting rulings, assume that the case with the more current ruling date is correct.
    - If the response is based on cases that are older than 5 years, make sure to inform the user that the information may be outdated,
      since some case opinions may no longer apply in law.
    - To summarize the case, use the get_opinion_text with summarize=True.
    - Use get_opinion_text with summarize=False only when full text is needed. Consider summarizing the text when possible to make things run faster.
    - If a user wants to learn more about a case, you can call the get_case_document_pdf tool with the citation to get a valid URL.
      If this is unsuccessful, call the get_case_document_page tool instead. 
      The text displayed with this URL should be the name_abbreviation of the case (DON'T just say the info can be found here).
      Don't call the get_case_document_page tool until after you have tried the get_case_document_pdf tool.
      Don't provide URLs from any other tools. Do not generate URLs yourself.
    - When presenting a URL in your response, use the validate_url tool.
    - If a user wants to test their argument, use the ask_caselaw tool to gather information about cases related to their argument 
      and the critique_as_judge tool to determine whether their argument is sound or has issues that must be corrected.
    - Never discuss politics, and always respond politely.
    """

    agent = Agent(
        tools=create_assistant_tools(_cfg),
        topic="Case law in Alaska",
        custom_instructions=legal_assistant_instructions,
        agent_progress_callback=agent_progress_callback,
    )
    agent.report()
    return agent