import os
import asyncio
from openai import AsyncOpenAI
from client import SQLDebugEnv, SQLFixAction

# Environment variables
API_BASE_URL = os.environ.get("API_BASE_URL", "https://api.openai.com/v1")
MODEL_NAME = os.environ.get("MODEL_NAME", "gpt-4o-mini")
HF_TOKEN = os.environ.get("HF_TOKEN", "")
ENV_BASE_URL = os.environ.get("ENV_BASE_URL", "http://localhost:8000")

async def run_inference():
    """
    Baseline inference script. 
    Connects to the environment, resets an episode, and uses an LLM to generate a fix.
    """
    client = AsyncOpenAI(
        base_url=API_BASE_URL,
        api_key=HF_TOKEN or "dummy-key"
    )

    async with SQLDebugEnv(base_url=ENV_BASE_URL) as env:
        
        print("Starting new episode...")
        result = await env.reset(task_group="easy_syntax_fix")
        obs = result.observation
        
        print(f"Task: {obs.task_description}")
        print(f"Buggy Query: {obs.buggy_query}")
        print(f"Error: {obs.error_message}")
        
        prompt = f"""
            You are an expert SQL engineer. Your task is to fix the following buggy query.
            
            Task Description: {obs.task_description}
            Database Schema: 
            {obs.schema_sql}
            
            Sample Data:
            {obs.sample_data}
            
            Buggy Query: {obs.buggy_query}
            Error Message: {obs.error_message}
            
            Return ONLY the corrected SQL query. Do not include markdown formatting or explanations.
        """
        
        try:
            response = await client.chat.completions.create(
                model=MODEL_NAME,
                messages=[{"role": "user", "content": prompt}],
                temperature=0.0
            )
            corrected_query = response.choices[0].message.content.strip()
            
            if corrected_query.startswith("```sql"):
                corrected_query = corrected_query[6:]
            if corrected_query.endswith("```"):
                corrected_query = corrected_query[:-3]
            corrected_query = corrected_query.strip()
            
            
            step_resp = await env.step(SQLFixAction(query=corrected_query))
            
            print(f"Reward: {step_resp.reward}")
            print(f"Success: {step_resp.success}")
            print(f"Feedback: {step_resp.observation.last_attempt_result}")
            
        except Exception as e:
            print(f"LLM API Error: {e}")

if __name__ == "__main__":
    asyncio.run(run_inference())