import hashlib
import json
import re
import time
import pprint
import openai
import groq
import os
import tiktoken
import asyncio
import google.generativeai as genai
from concurrent.futures import ThreadPoolExecutor
from dotenv import load_dotenv
from .util import debug_print, timing_decorator, DEBUG

load_dotenv()

CACHE_DIR = '.llm_cache'

openai.api_key = os.getenv("OPENAI_API_KEY")

# 2024-11-13, https://openai.com/api/pricing/
def get_price_per_token(model: str):
    if model == 'gpt-4o':
        return 0.0000025, 0.000010
    elif model == 'gpt-4o-mini':
        return 0.00000015, 0.0000006
    elif model == 'gpt-3.5-turbo':
        return 0, 0
    else:
        return 0, 0

@timing_decorator
async def parse_drugs_with_llm(text: str, model: str, prompt_file: str, system_prompt: str, response_format: object = None, temperature: float = 0.0):
    if model.startswith('gpt-'):
        return await parse_drugs_with_llm_openai(text, model, prompt_file, system_prompt, response_format, temperature)
    elif model.startswith('llama3-'):
        return await parse_drugs_with_llm_groq(text, model, prompt_file, system_prompt, response_format, temperature)
    elif model.startswith('gemini-'):
        return await parse_drugs_with_llm_gemini(text, model, prompt_file, system_prompt, response_format, temperature)
    else:
        raise Exception(f"Unsupported entity parser: {model}")

async def parse_drugs_with_llm_openai(text: str, model: str, prompt_file: str, system_prompt: str, response_format: object, temperature: float):
    try:
        model = model
        with open(f'prompts/{prompt_file}.txt', 'r') as file:
            prompt = file.read()
            prompt = prompt.format(text=text)

        if not response_format:
            if re.match(r'prompt[1-4]', prompt_file):
                json_schema = {
                    "name": "drug_analysis",
                    "strict": True,
                    "schema": {
                        "type": "object",
                        "properties": {
                            "drugs": {
                                "type": "array",
                                "items":    {
                                    "type": "object",
                                    "properties": {
                                        "name": {"type": "string"},
                                        "dosing": {"type": "string"},
                                        "context": {"type": "string"}
                                    },
                                    "required": ["name", "dosing", "context"],
                                    "additionalProperties": False
                                }
                            }
                        },
                        "required": ["drugs"],
                        "additionalProperties": False
                    }
                }
            elif re.match(r'prompt[56]', prompt_file):
                json_schema = {
                    "name": "drug_analysis",
                    "strict": True,
                    "schema": {
                        "type": "object",
                        "properties": {
                            "drugs": {
                                "type": "array",
                                "items": {
                                    "type": "object",
                                    "properties": {
                                        "name": {"type": "string"},
                                        "dosing": {"type": "string"},
                                        "context": {"type": "string"}
                                    },
                                    "required": ["name", "dosing", "context"],
                                    "additionalProperties": False
                                }
                            },
                            "sections": {
                                "type": "object",
                                "properties": {
                                    "current_therapy": {"type": "string"},
                                    "past_therapy": {"type": "string"},
                                    "allergies": {"type": "string"}
                                },
                                "required": ["current_therapy", "past_therapy", "allergies"],
                                "additionalProperties": False
                            }
                        },
                        "required": ["drugs", "sections"],
                        "additionalProperties": False
                    }
                }
            elif re.match(r'prompt[7]', prompt_file):
                json_schema = {
                    "name": "drug_analysis",
                    "strict": True,
                    "schema": {
                        "type": "object",
                        "properties": {
                            "drugs": {
                                "type": "array",
                                "items": {
                                    "type": "object",
                                    "properties": {
                                        "name": {"type": "string"},
                                        "dosing": {"type": "string"},
                                        "context": {"type": "string"},
                                        "quote": {"type": "string"}
                                    },
                                    "required": ["name", "dosing", "context", "quote"],
                                    "additionalProperties": False
                                }
                            },
                            "sections": {
                                "type": "object",
                                "properties": {
                                    "current_therapy": {"type": "string"},
                                    "past_therapy": {"type": "string"},
                                    "allergies": {"type": "string"}
                                },
                                "required": ["current_therapy", "past_therapy", "allergies"],
                                "additionalProperties": False
                            }
                        },
                        "required": ["drugs", "sections"],
                        "additionalProperties": False
                    }
                }
            elif re.match(r'prompt[8]', prompt_file):
                json_schema = {
                    "name": "drug_analysis",
                    "strict": True,
                    "schema": {
                        "type": "object",
                        "properties": {
                            "drugs": {
                                "type": "array",
                                "items": {
                                    "type": "object",
                                    "properties": {
                                        "name": {"type": "string"},
                                        "dosing": {"type": "string"},
                                        "context": {"type": "string"},
                                        "quote": {"type": "string"},
                                        "item_type": {"type": "string"}
                                    },
                                    "required": ["name", "dosing", "context", "quote", "item_type"],
                                    "additionalProperties": False
                                }
                            },
                            "sections": {
                                "type": "object",
                                "properties": {
                                    "current_therapy": {"type": "string"},
                                    "past_therapy": {"type": "string"},
                                    "allergies": {"type": "string"}
                                },
                                "required": ["current_therapy", "past_therapy", "allergies"],
                                "additionalProperties": False
                            }
                        },
                        "required": ["drugs", "sections"],
                        "additionalProperties": False
                    }
                }

        encoding = tiktoken.encoding_for_model(model)
        prompt_tokens = len(encoding.encode(prompt))
        debug_print(f'prompt_tokens (before req): {prompt_tokens}')

        # breakpoint()
        # make async request
        loop = asyncio.get_running_loop()
        with ThreadPoolExecutor() as pool:
            response = await loop.run_in_executor(pool, completion_request_openai, model, prompt, system_prompt, temperature, json_schema)

        # response = openai.chat.completions.create(
        #     model=model,
        #     messages=[
        #         {"role": "system", "content": "You are a helpful assistant that can identify drug names in text, regardless of the language."},
        #         {"role": "user", "content": prompt}
        #     ],
        #     temperature=temperature,
        #     response_format=response_format
        # )
        # debug_print(f'response_format: {response_format}')
        debug_print(f'response: {response}')
        content = response.choices[0].message.content

        # 2024-11-13, https://openai.com/api/pricing/
        if model == 'gpt-4o':
            price_per_prompt_token = 0.0000025
            price_per_completion_token = 0.000010
        elif model == 'gpt-4o-mini':
            price_per_prompt_token = 0.00000015
            price_per_completion_token = 0.0000006
        elif model == 'gpt-3.5-turbo':
            price_per_prompt_token = 0
            price_per_completion_token = 0
        else:
            price_per_prompt_token = 0
            price_per_completion_token = 0

        cost = response.usage.prompt_tokens * price_per_prompt_token
        cost += response.usage.completion_tokens * price_per_completion_token
        debug_print(f'input_tokens (actual): {response.usage.prompt_tokens}')
        debug_print(f'output_tokens (actual): {response.usage.completion_tokens}')
        debug_print(f'Cost of OpenAI request: {cost}')

        debug_print(f'model: {model}')
        debug_print(f'prompt (md5): {hashlib.md5(prompt.encode()).hexdigest()}')
        debug_print(f'prompt: {prompt}')
        debug_print(f'content (md5): {hashlib.md5(content.encode()).hexdigest()}')
        debug_print(f'content: {content}')

        response = json.loads(content.strip().lstrip('json```').rstrip('```'))
        return response, cost

    except Exception as e:
        print(f"Error calling OpenAI API: {str(e)}")
        if DEBUG:
            breakpoint()
        return "Error analyzing text for drug names.", 0

async def parse_drugs_with_llm_groq(text: str, model: str, prompt_file: str, system_prompt: str, response_format: object, temperature: float):
    try:
        model = model
        with open(f'prompts/{prompt_file}.txt', 'r') as file:
            prompt = file.read()
            prompt = prompt.format(text=text)

        client = groq.Groq()
        # pdb.set_trace()
        # pprint.pprint(prompt)
        t1 = time.time()
        response = client.chat.completions.create(
            model=model,
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": prompt}
            ],
            temperature=temperature,
            response_format={"type": "json_object"},
            stream=False
        )
        t2 = time.time()
        content = response.choices[0].message.content

        # 2024-11-21, https://groq.com/pricing/
        if model == 'llama3-8b-8192':
            price_per_input_token = 0.00000005
            price_per_output_token = 0.00000008
        else:
            price_per_input_token = 0
            price_per_output_token = 0

        cost = response.usage.prompt_tokens * price_per_input_token
        cost += response.usage.completion_tokens * price_per_output_token
        debug_print(f'Cost of GroqCloud request: {cost}')

        debug_print(f'model: {model}')
        debug_print(f'prompt (md5): {hashlib.md5(prompt.encode()).hexdigest()}')
        debug_print(f'prompt: {prompt}')
        debug_print(f'content (md5): {hashlib.md5(content.encode()).hexdigest()}')
        debug_print(f'content: {content}')

        response_json = content
        response_obj = json.loads(response_json, strict=False)
        t3 = time.time()
        # TODO: is this "data" default for Groq JSON mode? I didn't mention it in the prompt.
        if "data" in response_obj:
            drugs = response_obj["data"]
        elif "drugs" in response_obj:
            drugs = response_obj["drugs"]
        assert isinstance(drugs, list), f"Expected a list, got {type(drugs)}"
        assert all(isinstance(item, dict) for item in drugs), "Expected all items to be dictionaries"
        assert all(set(item.keys()) == {"name", "dosing", "context"} for item in drugs), "Expected all items to have keys 'name', 'dosing', and 'context'"

        debug_print(f'(parse_drugs_with_llm_groq) request: {round(t2 - t1, 4)}')
        debug_print(f'(parse_drugs_with_llm_groq) json parse: {round(t3 - t2, 4)}')
        return drugs, cost

    except Exception as e:
        print(f"Error calling GroqCloud API: {e.__class__.__name__}: {str(e)}")
        return "Error analyzing text for drug names.", 0

async def parse_drugs_with_llm_gemini(text: str, model: str, prompt_file: str, system_prompt: str, response_format: object, temperature: float):
    try:
        # Configure Gemini API
        genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))

        # Load prompt
        with open(f'prompts/{prompt_file}.txt', 'r') as file:
            prompt = file.read()
            prompt = prompt.format(text=text)

        # Combine system prompt and user prompt
        full_prompt = f"{system_prompt}\n\n{prompt}"

        # Configure model
        # model_name = model.replace('gemini-', '')
        model_name = model
        model = genai.GenerativeModel(model_name=model_name)

        # Make async request
        loop = asyncio.get_running_loop()
        with ThreadPoolExecutor() as pool:
            response = await loop.run_in_executor(
                pool,
                lambda: model.generate_content(
                    full_prompt,
                        generation_config=genai.types.GenerationConfig(
                        temperature=temperature
                    )
                )
            )

        content = response.text

        # Parse JSON response
        debug_print(f'model: {model_name}')
        debug_print(f'prompt (md5): {hashlib.md5(prompt.encode()).hexdigest()}')
        debug_print(f'prompt: {prompt}')
        debug_print(f'content (md5): {hashlib.md5(content.encode()).hexdigest()}')
        debug_print(f'content: {content}')

        # Clean up the response - remove markdown code blocks if present
        # content = content.strip().lstrip('```json').rstrip('```')
        content = re.sub(r'```json(.*)```.*', r'\1', content, flags=re.DOTALL)
        response_obj = json.loads(content)

        # Calculate cost (Gemini API is currently free, so cost is 0)
        cost = 0

        return response_obj, cost

    except Exception as e:
        print(f"Error calling Gemini API: {e.__class__.__name__}: {str(e)}")
        if DEBUG:
            breakpoint()
        return "Error analyzing text for drug names.", 0

async def parse_for_entities(text: str, model: str, prompt_file: str, properties: dict, system_prompt: str, temperature: float, use_cache: bool = True):
    text_hash = hashlib.md5(text.encode()).hexdigest()[:8]
    properties_hash = hashlib.md5(json.dumps(properties).encode()).hexdigest()[:8]
    system_prompt_hash = hashlib.md5(system_prompt.encode()).hexdigest()[:8]
    cache_filename = f"{CACHE_DIR}/{text_hash}_{model}_{prompt_file}_{properties_hash}_{system_prompt_hash}_{temperature}_llm_cache.json"

    if use_cache and os.path.exists(cache_filename):
        print(f'Reading cache file: {cache_filename}')
        with open(cache_filename, "r") as cache_file:
            response = json.load(cache_file)
    else:
        if model.startswith('gpt-') or model.startswith('o1-'):
            response = await parse_for_entities_openai(text, model, prompt_file, properties, system_prompt, temperature)
        elif model.startswith('llama3-') or model.startswith('llama-3') or model.startswith('mixtral-'):
            response = await parse_for_entities_groq(text, model, prompt_file, properties, system_prompt, temperature)
        else:
            raise Exception(f"Unsupported model: {model}")

        if response != [] and use_cache:
            print(f'Writing cache file: {cache_filename}')
            with open(cache_filename, "w") as cache_file:
                json.dump(response, cache_file)

    return response

async def parse_for_multiple_entities(text: str, model: str, prompt_file: str, properties: dict, system_prompt: str, temperature: float, use_cache: bool = True):
    text_hash = hashlib.md5(text.encode()).hexdigest()[:8]
    properties_hash = hashlib.md5(json.dumps(properties).encode()).hexdigest()[:8]
    system_prompt_hash = hashlib.md5(system_prompt.encode()).hexdigest()[:8]
    cache_filename = f"{CACHE_DIR}/{text_hash}_{model}_{prompt_file}_{properties_hash}_{system_prompt_hash}_{temperature}_llm_cache.json"

    if use_cache and os.path.exists(cache_filename):
        print(f'Reading cache file: {cache_filename}')
        with open(cache_filename, "r") as cache_file:
            response = json.load(cache_file)
    else:
        if model.startswith('gpt-') or model.startswith('o1-'):
            response = await parse_for_entities_openai(text, model, prompt_file, properties, system_prompt, temperature, multiple_entities=True)
        elif model.startswith('llama3-') or model.startswith('llama-3') or model.startswith('mixtral-'):
            response = await parse_for_entities_groq(text, model, prompt_file, properties, system_prompt, temperature, multiple_entities=True)
        else:
            raise Exception(f"Unsupported model: {model}")

        if response != [] and use_cache:
            print(f'Writing cache file: {cache_filename}')
            with open(cache_filename, "w") as cache_file:
                json.dump(response, cache_file)

    return response

def completion_request_openai(model, prompt, system_prompt, temperature, json_schema):
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": prompt}
    ]

    debug_print(f'system_prompt: {system_prompt}')
    debug_print(f'prompt: {prompt}')
    debug_print(f'temperature: {temperature}')
    if json_schema:
        debug_print(f'json_schema: {json.dumps(json_schema)}')

    if json_schema:
        # breakpoint()
        return openai.chat.completions.create(
            model=model,
            temperature=temperature,
            messages=messages,
            response_format={
                "type": "json_schema",
                "json_schema": json_schema
            }
        )
    else:
        return openai.chat.completions.create(
            model=model,
            temperature=temperature,
            messages=messages,
        )

@timing_decorator
async def parse_for_entities_openai(text: str, model: str, prompt_file: str, properties: dict, system_prompt: str, temperature: float, multiple_entities: bool = False):
    try:
        # load prompt
        with open(f'prompts/{prompt_file}.txt', 'r') as file:
            prompt = file.read()
            prompt = prompt.format(text=text)

        if model.startswith('gpt-'):
            # TODO: this needs better naming.
            # TODO: multiple_entities=True currently means each property is an entity and each may or may not be repeated
            # breakpoint()

            # breakpoint()
            if multiple_entities:
                properties2 = {
                    f'{k}s': {
                        "type": "array",
                        "items": {
                            "type": "object",
                            "properties": {
                                k: v
                            },
                            "required": [k],
                            "additionalProperties": False
                        }
                    } for k, v in properties.items()
                }
                json_schema = {
                    "name": "entity_analysis",
                    "strict": True,
                    "schema": {
                        "type": "object",
                        "properties": properties2,
                        "required": list([f'{k}s' for k in properties.keys()]),
                        "additionalProperties": False
                    }
                }
            # TODO (cont'd): multiple_entities=False currently means the entire set of properties/fields is a single entity and the entire set may be repeated
            else:
                json_schema = {
                    "name": "entity_analysis",
                    "strict": True,
                    "schema": {
                        "type": "object",
                        "properties": {
                            "entities": {
                                "type": "array",
                                "items": {
                                    "type": "object",
                                    "properties": properties,
                                    "required": list(properties.keys()),
                                    "additionalProperties": False
                                }
                            }
                        },
                        "required": ["entities"],
                        "additionalProperties": False
                    }
                }

        elif model.startswith('o1-'):
            json_schema = None
        else:
            raise Exception(f"Unsupported model: {model}")

        # make async request
        loop = asyncio.get_running_loop()
        with ThreadPoolExecutor() as pool:
            response = await loop.run_in_executor(pool, completion_request_openai, model, prompt, system_prompt, temperature, json_schema)

        # calculate cost
        price_per_prompt_token, price_per_completion_token = get_price_per_token(model)
        cost = response.usage.prompt_tokens * price_per_prompt_token
        cost += response.usage.completion_tokens * price_per_completion_token
        debug_print(f'Cost of OpenAI request: {cost}')

        # parse
        try:
            content = response.choices[0].message.content
            if model.startswith('gpt-'):
                entities = json.loads(content)
            elif model.startswith('o1-'):
                content = content.replace("```json", "").replace("```", "")
                entities = json.loads(content)
            else:
                raise Exception(f"Unsupported model: {model}")
        except Exception as e:
            print(f"Error parsing OpenAI response: {e.__class__.__name__}: {str(e)}")
            entities = []
            breakpoint()

        return entities, cost

    except Exception as e:
        print(f"Error calling OpenAI API: {e.__class__.__name__}: {str(e)}")
        breakpoint()
        return "Error analyzing text for drug names.", 0

async def parse_for_entities_groq(text: str, model: str, prompt_file: str, properties: dict, system_prompt: str, temperature: float, multiple_entities: bool = False):
    try:
        with open(f'prompts/{prompt_file}.txt', 'r') as file:
            prompt = file.read()
            prompt = prompt.format(text=text)

        client = groq.Groq()
        response = client.chat.completions.create(
            model=model,
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": prompt}
            ],
            temperature=temperature,
            response_format={"type": "json_object"},
            stream=False
        )

        content = response.choices[0].message.content

        # 2024-11-21, https://groq.com/pricing/
        if model == 'llama3-8b-8192':
            price_per_input_token = 0.00000005
            price_per_output_token = 0.00000008
        else:
            price_per_input_token = 0
            price_per_output_token = 0

        cost = response.usage.prompt_tokens * price_per_input_token
        cost += response.usage.completion_tokens * price_per_output_token
        debug_print(f'Cost of GroqCloud request: {cost}')

        debug_print(f'model: {model}')
        debug_print(f'prompt (md5): {hashlib.md5(prompt.encode()).hexdigest()}')
        debug_print(f'prompt: {prompt}')
        debug_print(f'content (md5): {hashlib.md5(content.encode()).hexdigest()}')
        debug_print(f'content: {content}')

        response_json = content
        response_obj = json.loads(response_json, strict=False)
        entities = list(response_obj.values())[0] # get first and only element

        assert isinstance(entities, list), f"Expected a list, got {type(entities)}"
        assert all(isinstance(item, dict) for item in entities), "Expected all items to be dictionaries"
        assert all(set(item.keys()) == set(properties.keys()) for item in entities), "Expected all items to have keys 'name', 'context' and 'icd_codes'"

        return entities, cost

    except Exception as e:
        print(f"Error calling GroqCloud API: {e.__class__.__name__}: {str(e)}")
        return "Error analyzing text for drug names.", 0