import os
from g4f.client import Client
from litellm import completion
import random
import json
import os
from g4f.Provider import DeepInfraChat,LambdaChat
from backup import Client as PerplexityClient,cookies
import re
import base64
import mimetypes
import urllib.parse

def _ext_from_url(url):
    try:
        path = urllib.parse.urlparse(url).path
        root, ext = os.path.splitext(path)
        if ext:
            return ext
    except Exception:
        pass
    return '.jpg'


gemini_api_keys=json.loads(os.environ.get("GEMINI_KEY_LIST"))
groq_api_keys=json.loads(os.environ.get("GROQ_API_KEYS"))
chutes_key=os.environ.get("CHUTES_API_KEY")
github_key=os.environ.get("GITHUB_API_KEY")

DeepInfraChat.models = ["moonshotai/Kimi-K2-Instruct","Qwen/Qwen3-235B-A22B-Thinking-2507","Qwen/Qwen3-235B-A22B","Qwen/Qwen3-30B-A3B","Qwen/Qwen3-32B","google/gemma-3-27b-it","deepseek-ai/DeepSeek-R1-Turbo","Qwen/QwQ-32B","deepseek-ai/DeepSeek-R1-0528","deepseek-ai/DeepSeek-V3-0324","meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8","meta-llama/Llama-4-Scout-17B-16E-Instruct","microsoft/Phi-4-multimodal-instruct"]

deepinframodels=["meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8","microsoft/Phi-4-multimodal-instruct","google/gemma-3-27b-it","meta-llama/Llama-4-Scout-17B-16E-Instruct"]
chutes_models={'Llama-4-Maverick-17B-128E-Instruct-FP8':'chutesai/Llama-4-Maverick-17B-128E-Instruct-FP8', "Qwen3-235B":"Qwen/Qwen3-235B-A22B","MAI-DS-R1-FP8":"microsoft/MAI-DS-R1-FP8","DeepSeek-V3-0324":"deepseek-ai/DeepSeek-V3-0324","deepseek-reasoner":"deepseek-ai/DeepSeek-R1-0528","GLM-4-32B-0414":"THUDM/GLM-4-32B-0414","GLM-Z1-32B-0414":"THUDM/GLM-Z1-32B-0414","DeepSeek-R1T-Chimera":"tngtech/DeepSeek-R1T-Chimera", "DeepSeek-R1-Zero":"deepseek-ai/DeepSeek-R1-Zero"}
github_models={"gpt4.1":"gpt-4.1","gpt-4o":"gpt-4o","o4-mini":"o4-mini"}

REASONING_CORRESPONDANCE = {"DeepSeekR1":DeepInfraChat}
os.environ["GEMINI_API_KEY"] =random.choice(gemini_api_keys)

REASONING_QWQ = {"qwq-32b":DeepInfraChat}


CHAT_CORRESPONDANCE = {"DeepSeek-V3":DeepInfraChat}

client = Client()

perplexity_cli=PerplexityClient(cookies=cookies)


def chat(messages,response_format,model="gpt-4"):
        if len(messages) ==1:
            messages[0]["role"]="user"
        response = completion(
            model="gemini/gemini-2.0-flash", 
            messages=messages,
            response_format=response_format
        )
        return str(response.choices[0].message.content)



def chatstream(messages,model,api_keys,tools):
    print(f"-------{model}--------")
    global llmfree
    global llmdeepseek
    global llmgroq

    cunk=""



    if model in ["o3","gpt-4.1",'grok-4','gemini-2.5-pro','claude-sonnet-4-20250514','sonar-pro','r1-1778']:
        # build the raw message text as before
        raw_perplex_msg = "".join(
            (( (f"[{message['role']}]" ) +  ("(#message)" if message['role']!="system" else "(#instructions)") ) if message['role'] != "assistant"   else "") + f"\n{message['content']}\n\n"
            for message in messages
        )

        # --- extract images / image urls into files dict ---
        files = {}
        image_counter = 1

        # regexes
        markdown_image_pattern = re.compile(r'!\[.*?\]\((.*?)\)')
        image_url_pattern = re.compile(r'https?://\S+\.(?:png|jpg|jpeg|gif|webp|svg)', re.IGNORECASE)
        data_uri_pattern = re.compile(r'data:(image/[\w.+-]+);base64,([A-Za-z0-9+/=\n\r]+)', re.IGNORECASE)

        for message in messages:
            # 1) attachments (common OpenAI-like shape)
            atts = message.get('attachments') or message.get('attachment') or []
            if isinstance(atts, dict):
                atts = [atts]
            if isinstance(atts, list):
                for att in atts:
                    if not isinstance(att, dict):
                        continue
                    # common keys for urls
                    url = att.get('url') or att.get('image_url') or att.get('src') or att.get('href')
                    if url:
                        ext = _ext_from_url(url)
                        fname = f'image_{image_counter}{ext}'
                        files[fname] = url  # pass URL as file value (per your example style)
                        image_counter += 1

            # 2) content as string -> look for data-URIs, markdown images, plain image URLs
            content = message.get('content', '') or ''
            if not isinstance(content, str):
                try:
                    content = str(content)
                except Exception:
                    content = ''

            # data URIs -> decode and attach raw bytes
            for m in data_uri_pattern.finditer(content):
                mime_type, b64 = m.groups()
                try:
                    b = base64.b64decode(b64)
                except Exception:
                    continue
                ext = mimetypes.guess_extension(mime_type) or '.bin'
                fname = f'image_{image_counter}{ext}'
                files[fname] = b  # raw bytes
                image_counter += 1

            # markdown images: ![alt](url)
            for m in markdown_image_pattern.finditer(content):
                url = m.group(1)
                ext = _ext_from_url(url)
                fname = f'image_{image_counter}{ext}'
                files[fname] = url
                image_counter += 1

            # plain image URLs in text
            for m in image_url_pattern.finditer(content):
                url = m.group(0)
                ext = _ext_from_url(url)
                fname = f'image_{image_counter}{ext}'
                files[fname] = url
                image_counter += 1

        # also include the textual messages file (so Perplexity gets the full text payload as a file)
        # files['perplexity_messages.txt'] = raw_perplex_msg

        # write the textual file to disk as before (optional, kept for compatibility/debugging)
        with open("perplexity_messages.txt", "w", encoding='utf-8') as f:
            f.write(raw_perplex_msg)

        # run the search, passing files (which may contain URL strings or raw bytes for data URIs)
        resp = perplexity_cli.search(raw_perplex_msg, mode='reasoning', model=model, sources=[], files=files, stream=True, language='en-US', follow_up=None, incognito=False)

        # stream the response, aggregate into cunk (init if needed)
        cunk = ""
        for i in resp:
            try:
                cunk = cunk + (i["blocks"][0]["markdown_block"]["chunks"][0])
                if "```json" not in cunk or "```" not in cunk:
                    yield i["blocks"][0]["markdown_block"]["chunks"][0]
            except Exception as e:
                print(e)
                pass
        yield ("RESULT: " + cunk)
        
    elif model in deepinframodels:
        try:
            response = client.chat.completions.create(
                provider=DeepInfraChat,
                model=model,
                messages=messages,
                stream=True,
                response_format="json_object"

            )
            for part in response:
                cunk=cunk+(part.choices[0].delta.content or "")
                if "```json" not in cunk or "```" not in cunk:
                    yield (part.choices[0].delta.content or "")
        except Exception as e:
            pass
        yield ("RESULT: "+cunk)

    elif model == "Qwen3-235B-A22B-Thinking-2507" :
        response = client.chat.completions.create(
            provider=DeepInfraChat,
            model=f"Qwen/Qwen3-235B-A22B-Thinking-2507",
            messages=messages,
            stream=True

        )
        for part in response:
            resp=str(part.choices[0].delta.content)
            cunk=cunk+(resp or "")

            if ("```json" not in cunk or "```" not in cunk) and resp != "None":
                yield (resp or "")
        yield ("RESULT: "+str(cunk))
    elif model == "Kimi-K2-Instruct" :
        response = client.chat.completions.create(
            provider=DeepInfraChat,
            model=f"moonshotai/Kimi-K2-Instruct",
            messages=messages,
            stream=True

        )
        for part in response:
            resp=str(part.choices[0].delta.content)
            cunk=cunk+(resp or "")

            if ("```json" not in cunk or "```" not in cunk) and resp != "None":
                yield (resp or "")
        yield ("RESULT: "+str(cunk))
    elif model == "DeepSeekR1-togetherAI":
        response = completion(model="together_ai/deepseek-ai/DeepSeek-R1", messages=messages, stream=True)

        cunk=""
        for part in response:
            cunk=cunk+(part.choices[0].delta.content or "")
            if "```json" not in cunk:
                yield(part.choices[0].delta.content or "")

        yield("RESULT: "+cunk)
    
    elif model == "DeepSeekV3-togetherAI":
        response = completion(model="together_ai/deepseek-ai/DeepSeek-V3", messages=messages, stream=True)

        cunk=""
        for part in response:
            cunk=cunk+(part.choices[0].delta.content or "")

            if "```json" not in cunk:
                yield(part.choices[0].delta.content or "")
        yield("RESULT: "+cunk)

    elif model=="groq/deepseek-r1-distill-llama-70b":
        os.environ["GROQ_API_KEY"] =random.choice(groq_api_keys)

        response = completion(model="groq/deepseek-r1-distill-llama-70b", messages=messages, stream=True)

        cunk=""
        for part in response:
            cunk=cunk+(part.choices[0].delta.content or "")
            if "```json" not in cunk:
                yield(part.choices[0].delta.content or "")
        yield("RESULT: "+cunk)
    elif model=="groq/qwq-32b":
        os.environ["GROQ_API_KEY"] =random.choice(groq_api_keys)
        response = completion(model="groq/qwen-qwq-32b", messages=messages, stream=True)

        cunk=""
        for part in response:
            cunk=cunk+(part.choices[0].delta.content or "")
            if "```json" not in cunk:
                yield(part.choices[0].delta.content or "")

        yield("RESULT: "+cunk)

    elif model=="llama-3.3-70b-versatile":
        response = completion(model="groq/llama-3.3-70b-versatile", messages=messages, stream=True)

        cunk=""
        for part in response:
            cunk=cunk+(part.choices[0].delta.content or "")
            
            if "```json" not in cunk:
                yield(part.choices[0].delta.content or "")

        yield("RESULT: "+cunk)
    elif model in chutes_models:
        response = completion(model=f"openai/{chutes_models[model]}",api_key=chutes_key,base_url="https://llm.chutes.ai/v1", messages=messages, stream=True)
        if model == "MAI-DS-R1-FP8" or model == "GLM-Z1-32B-0414" or model == "DeepSeek-R1T-Chimera" or model == "QwQ-32B-ArliAI-RpR-v1":

            yield("<think> \n")

        cunk=""
        hist = ""
        for part in response:
            x=str(part.choices[0].delta.content)
            cunk=cunk+x
            print(x, end="")

            if "```" in x and "```json"  not in cunk:
                hist = x
                continue
            if hist!="" and "json" not in cunk:
                yield(hist + x)
                hist = ""
                continue

            if ("```json" not in cunk) and (x != "None"):
                if "None" not in x:
                    yield(x)
        before, found, after = cunk.partition('</think>')
        cunk=after
        yield("RESULT: "+cunk)

    elif model in github_models:
        response = completion(model=f"github/{github_models[model]}",api_key=github_key, messages=messages,tools=tools, stream=True)

        cunk=""
        for part in response:
            chunk_dict = part.model_dump()
            cunk=cunk+(part.choices[0].delta.content or "")
            yield ("RAW: " +json.dumps(chunk_dict, indent=4))
            

        yield("RESULT: "+cunk)

    elif "gemini" in model: 
        for key in gemini_api_keys:
            try:
                os.environ["GEMINI_API_KEY"] =key

                response = completion(model=f"gemini/{model}", messages=messages, tools=tools,stream=True)

                cunk=""
                for part in response:
                    cunk=cunk+(part.choices[0].delta.content or "")
                    chunk_dict = part.model_dump()

                    yield ("RAW: " +json.dumps(chunk_dict, indent=4))

                    # if "```json" not in cunk:
                    #     yield(part.choices[0].delta.content or "")

                break         
            except Exception as e:
                print(str(e))
                pass            
        print("STOPPING")
        yield("RESULT: "+cunk)

    elif model=="deepseek.r1" or model=="deepseek-chat":

        cunk=""
        hist = ""
        if "chat" in model:
            providers = CHAT_CORRESPONDANCE
            model_name="deepseek-ai/DeepSeek-V3-0324"
        else:
            providers = REASONING_CORRESPONDANCE
            model_name="deepseek-ai/DeepSeek-R1-0528"

        for i  in range(2):
            try:
                response = client.chat.completions.create(
                    provider=DeepInfraChat,
                    model=model_name,
                    messages=messages,
                    stream=True

                    # Add any other necessary parameters
                )
                for part in response:
                    x=str(part.choices[0].delta.content)
                    cunk=cunk+x
                    print(x, end="")

                    if "```" in x and "```json"  not in cunk:
                        hist = x
                        continue
                    if hist!="" and "json" not in cunk:
                        yield(hist + x)
                        hist = ""
                        continue

                    if ("```json" not in cunk) and (x != "None"):
                        if "None" not in x:
                            yield(x)
                break
            except Exception as e:
                #yield(str(e)) 
                print(e)
                pass
        print("STOPPING")
        before, found, after = cunk.partition('</think>')
        cunk=after

        yield("RESULT: "+cunk)


    elif model=="qwq-32b" :
        yield("<think>")
        cunk=""
        providers=REASONING_QWQ
        for provider in providers:
            try:
                response = client.chat.completions.create(
                    provider=providers[provider],
                    model="Qwen/QwQ-32B",
                    messages=messages,
                    stream=True

                    # Add any other necessary parameters
                )
                for part in response:
                    cunk=cunk+(part.choices[0].delta.content or "")
                    if "```json" not in cunk or "```" not in cunk:
                        yield(part.choices[0].delta.content or "")

                break
            except Exception as e:
                
                pass
        yield("RESULT: "+cunk)


    elif "DeepSeek" in model and "dev" in model:
        cunk=""

        if "V3" in model:
            providers = CHAT_CORRESPONDANCE
        else:
            providers = REASONING_CORRESPONDANCE
        for provider in providers:
            try:
                response = client.chat.completions.create(
                    provider=providers[provider],
                    model="deepseek-r1",
                    messages=messages,
                    stream=True

                    # Add any other necessary parameters
                )
                for part in response:
                    cunk=cunk+(part.choices[0].delta.content or "")
                break
            except Exception as e:
                pass
        print("STOPPING")
        yield("RESULT: "+cunk)

# #Predict top colleges in BHU for me via CUET 2024 based on 2024 cut off scores.
# Degree:BA (Hons) Pol. Sci.
# Respond in csv format with the following headings and info: College,Cut-off Score(Out of 800),Cutoff Rank.Do not output the sources like [1],[2]...within the csv code.
# Note:
# 1)First search and research the web thoroughly end to end.
# 2)Always Use the website https://collegedunia.com to search about the cut offs.Dont use any other website but only this website.
# Mandatory:**Give Atleast 10 college names**