import json
import time
import hashlib
import base64
import asyncio
import aiohttp
from fastapi import FastAPI, Request, Response
from typing import Optional, Dict, List, Any
import uvicorn
from pydantic import BaseModel

# Configuration
API_CHALLENGE_URL = 'https://api.eqing.tech/api/altcaptcha/challenge'
NEXTWAY_CHAT_URL = 'https://origin.eqing.tech/api/openai/v1/chat/completions'
CREDENTIAL_EXPIRY_MARGIN = 60 * 1000  # 60 seconds in milliseconds
PORT = 7860
API_ENDPOINT = '/v1/chat/completions'
MODEL_NAME = "gpt-4o-free"
REQUEST_TIMEOUT = 480  # seconds
MAX_RETRIES = 3
RETRY_DELAY = 1  # seconds

# Global variables
current_credential = None
credential_expiry = None
is_refreshing_credential = False

# Model mappings
MODEL_MAPPING = {
    'gpt-4o-all-lite': 'gpt-4o-mini',
    'gpt-4o-image': 'gpt-4o-mini-image-free',
    'grok-3-re': 'grok-3-free',
    'gemini-2.0-flash': 'gemini-2.0-flash-free'
}

app = FastAPI()

class ChatRequest(BaseModel):
    messages: List[Dict[str, str]]
    model: Optional[str]
    stream: Optional[bool] = True
    temperature: Optional[float] = 0.5
    presence_penalty: Optional[float] = 0
    frequency_penalty: Optional[float] = 0
    top_p: Optional[float] = 1
    max_tokens: Optional[int] = 4000

def extract_content(text: str) -> str:
    lines = text.split('\n')
    ai_response = ''
    ignored_id = 'chatcmpl-nxaTnETiUyAVBjdRwlr38Yt3'
    created = 1687070102

    for line in lines:
        line = line.strip()
        if line.startswith('data:'):
            data_str = line[5:].strip()
            if not data_str or data_str in ['[ORIGIN]', '[DONE]']:
                continue

            try:
                json_data = json.loads(data_str)
                if json_data.get('id') == ignored_id or json_data.get('created') == created:
                    continue

                if (json_data.get('choices') and json_data['choices'][0].get('delta') and 
                    'content' in json_data['choices'][0]['delta']):
                    ai_response += json_data['choices'][0]['delta']['content']

            except json.JSONDecodeError:
                print(f'Skipping non-JSON data')

    return ai_response

async def solve_challenge(challenge: str, salt: str, algorithm: str = "SHA-512", max_number: int = 1000000):
    start_time = time.time()
    
    for number in range(max_number):
        hash_value = await verify_hash(salt, number, algorithm)
        if hash_value == challenge:
            return {"number": number, "took": int((time.time() - start_time) * 1000)}
    
    return None

async def verify_hash(salt: str, number: int, algorithm: str) -> str:
    input_str = f"{salt}{number}"
    
    if algorithm == "SHA-512":
        hash_obj = hashlib.sha512(input_str.encode())
        return hash_obj.hexdigest()
    elif algorithm == "SHA-256":
        hash_obj = hashlib.sha256(input_str.encode())
        return hash_obj.hexdigest()
    else:
        raise ValueError(f"Unsupported algorithm: {algorithm}")

async def generate_credential():
    global current_credential, credential_expiry
    
    async with aiohttp.ClientSession() as session:
        try:
            async with session.get(API_CHALLENGE_URL) as response:
                if response.status != 200:
                    print(f"Challenge request failed with status {response.status}")
                    return None
                
                data = await response.json()
                solution = await solve_challenge(
                    data['challenge'],
                    data['salt'],
                    data['algorithm'],
                    data['maxnumber']
                )

                if not solution:
                    print("Failed to solve challenge")
                    return None

                credential_obj = {
                    "algorithm": data['algorithm'],
                    "challenge": data['challenge'],
                    "number": solution['number'],
                    "salt": data['salt'],
                    "signature": data['signature'],
                    "took": solution['took']
                }

                credential = base64.b64encode(json.dumps(credential_obj).encode()).decode()
                expiry = int(data['salt'].split('?expires=')[1]) * 1000

                return {"credential": credential, "expiry": expiry}

        except Exception as e:
            print(f"Error generating credential: {e}")
            return None

async def get_credential():
    global current_credential, credential_expiry, is_refreshing_credential
    
    if (not current_credential or 
        not credential_expiry or 
        credential_expiry <= time.time() * 1000 + CREDENTIAL_EXPIRY_MARGIN):
        
        if not is_refreshing_credential:
            is_refreshing_credential = True
            try:
                cred_data = await generate_credential()
                if cred_data:
                    current_credential = cred_data['credential']
                    credential_expiry = cred_data['expiry']
            finally:
                is_refreshing_credential = False
        else:
            await asyncio.sleep(2)
    
    return current_credential

@app.post(API_ENDPOINT)
async def chat_endpoint(request: ChatRequest):
    try:
        model = MODEL_MAPPING.get(request.model, request.model or MODEL_NAME)
        response_content = await handle_chat_request(model, request.messages)
        
        if response_content is None:
            return Response(
                content="Failed to get response from API",
                status_code=500
            )

        return {
            "choices": [{
                "message": {
                    "role": "assistant",
                    "content": response_content.strip()
                },
                "finish_reason": "stop",
                "index": 0
            }],
            "model": model,
            "object": "chat.completion"
        }

    except Exception as e:
        print(f"Error handling chat request: {e}")
        return Response(content="Internal server error", status_code=500)

async def handle_chat_request(model: str, messages: List[Dict[str, str]]):
    captcha_token = await get_credential()
    if not captcha_token:
        return None

    body = {
        "messages": messages,
        "stream": True,
        "model": model,
        "temperature": 0.5,
        "presence_penalty": 0,
        "frequency_penalty": 0,
        "top_p": 1,
        "max_tokens": 4000,
        "captchaToken": captcha_token
    }

    async with aiohttp.ClientSession() as session:
        try:
            async with session.post(
                NEXTWAY_CHAT_URL,
                json=body,
                timeout=REQUEST_TIMEOUT
            ) as response:
                if response.status != 200:
                    return None

                complete_response = ""
                async for chunk in response.content:
                    chunk_text = chunk.decode()
                    content = extract_content(chunk_text)
                    if content:
                        complete_response += content

                return complete_response

        except Exception as e:
            print(f"Error in chat request: {e}")
            await get_credential()  # Refresh credential on error
            return None

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=PORT)