"""
Inference Script — Git Conflict Resolution OpenEnv Environment
Follows the exact OpenEnv submission format requirements.
"""
import os
import json
import sys
from typing import List, Optional

from openai import OpenAI
from environment import make_env, Action

API_BASE_URL = os.getenv("API_BASE_URL", "https://api.groq.com/openai/v1")
MODEL_NAME = os.getenv("MODEL_NAME", "llama-3.3-70b-versatile")
HF_TOKEN = os.getenv("HF_TOKEN")
LOCAL_IMAGE_NAME = os.getenv("LOCAL_IMAGE_NAME")

BENCHMARK = "git-conflict-resolver"
MAX_STEPS = 3
SUCCESS_SCORE_THRESHOLD = 0.5

API_KEY = HF_TOKEN or os.getenv("GROQ_API_KEY") or os.getenv("API_KEY") or ""

client = OpenAI(
    base_url=API_BASE_URL,
    api_key=API_KEY,
)

SYSTEM_PROMPT = """\
You are an expert software engineer specializing in resolving Git merge conflicts.
You will be given files containing Git conflict markers (<<<<<<<, =======, >>>>>>>).
Your job is to resolve each conflict by producing clean, correct code.

Rules:
1. Remove ALL conflict markers from your output
2. Merge the changes intelligently — preserve all intended features from both branches
3. The resolved code must be syntactically valid Python
4. Return ONLY a JSON object with this structure:
{
  "resolved_files": {
    "filename.py": "...full resolved content..."
  }
}
No explanation, no markdown, no backticks. Pure JSON only.
"""


# ---------------------------------------------------------------------------
# Structured log helpers — exact format required by OpenEnv
# ---------------------------------------------------------------------------

def log_start(task: str, env: str, model: str) -> None:
    print(f"[START] task={task} env={env} model={model}", flush=True)


def log_step(step: int, action: str, reward: float, done: bool, error: Optional[str]) -> None:
    error_val = error if error else "null"
    done_val = str(done).lower()
    print(
        f"[STEP] step={step} action={action} reward={reward:.2f} done={done_val} error={error_val}",
        flush=True,
    )


def log_end(success: bool, steps: int, score: float, rewards: List[float]) -> None:
    rewards_str = ",".join(f"{r:.2f}" for r in rewards)
    print(
        f"[END] success={str(success).lower()} steps={steps} score={score:.3f} rewards={rewards_str}",
        flush=True,
    )


# ---------------------------------------------------------------------------
# Prompt helpers
# ---------------------------------------------------------------------------

def format_prompt(obs) -> str:
    obs_dict = obs.model_dump()
    lines = [
        f"Task: {obs_dict['task_description']}",
        f"Hint: {obs_dict['test_cases_hint']}",
        "",
        "Conflicted files to resolve:",
    ]
    for fname, content in obs_dict["conflicted_files"].items():
        lines.append(f"\n--- {fname} ---\n{content}")
    lines.append("\nReturn the resolved files as JSON.")
    return "\n".join(lines)


def get_model_action(prompt: str) -> Optional[dict]:
    try:
        response = client.chat.completions.create(
            model=MODEL_NAME,
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": prompt},
            ],
            temperature=0.2,
            max_tokens=2000,
        )
        raw = response.choices[0].message.content.strip()

        if raw.startswith("```"):
            raw = raw.split("```")[1]
            if raw.startswith("json"):
                raw = raw[4:]
        raw = raw.strip()

        data = json.loads(raw)
        return data.get("resolved_files", {})

    except Exception as e:
        return None


# ---------------------------------------------------------------------------
# Run one task episode
# ---------------------------------------------------------------------------

def run_task(task_id: str) -> float:
    env = make_env(task_id)
    obs = env.reset()

    rewards: List[float] = []
    steps_taken = 0
    score = 0.0
    success = False

    log_start(task=task_id, env=BENCHMARK, model=MODEL_NAME)

    try:
        for step in range(1, MAX_STEPS + 1):
            prompt = format_prompt(obs)
            resolved_files = get_model_action(prompt)

            if resolved_files is None:
                log_step(step=step, action="null", reward=0.0, done=False, error="api_or_parse_error")
                rewards.append(0.0)
                steps_taken = step
                continue

            action = Action(resolved_files=resolved_files)
            obs, reward, done, info = env.step(action)

            rewards.append(reward)
            steps_taken = step
            raw_score = max(score, info["score"])
            score = min(max(raw_score, 0.01), 0.99)

            log_step(
                step=step,
                action=f"resolve_{task_id}",
                reward=reward,
                done=done,
                error=None,
            )

            if done:
                break

        success = score >= SUCCESS_SCORE_THRESHOLD

    finally:
        rewards = [min(max(r, 0.01), 0.99) for r in rewards]
        score = min(max(score, 0.01), 0.99)
        log_end(success=success, steps=steps_taken, score=score, rewards=rewards)

    return score


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main():
    tasks = ["easy", "medium", "hard"]
    scores = {}

    for task_id in tasks:
        score = run_task(task_id)
        scores[task_id] = score

    avg = sum(scores.values()) / len(scores)

    with open("baseline_scores.json", "w") as f:
        json.dump({
            "scores": scores,
            "average": avg,
            "model": MODEL_NAME,
            "api_base_url": API_BASE_URL,
        }, f, indent=2)


if __name__ == "__main__":
    main()