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

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']:
        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
        )
        with open("perplexity_messages.txt", "w",encoding='utf-8') as f:
            f.write(raw_perplex_msg)
        resp = perplexity_cli.search(raw_perplex_msg, mode='reasoning', model=model, sources=[], files={}, stream=True, language='en-US', follow_up=None, incognito=False)
        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]
                # print(i["blocks"][0]["markdown_block"]["chunks"][0],end="")
            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**