#!/usr/bin/env python3
import asyncio
import os
import json
import textwrap
from typing import List, Optional

from openai import OpenAI

from openenv import GenericEnvClient as EnvClient
from models import GitCommitReviewerAction

# Defaults for inference
API_BASE_URL = os.getenv("API_BASE_URL", "https://api.openai.com/v1")
MODEL_NAME = os.getenv("MODEL_NAME", "gpt-4-turbo-preview")
API_KEY = os.getenv("HF_TOKEN") or os.getenv("API_KEY")

BENCHMARK = os.getenv("GIT_COMMIT_BENCHMARK", "git-commit-reviewer")
MAX_STEPS = 5
TEMPERATURE = 0.2
MAX_TOKENS = 800

# Try to get task from env, or default to easy.
# The user's phase 1 tests test easy, medium, hard presumably by modifying script or env.
TASK_NAME = os.getenv("GIT_COMMIT_TASK", "easy")

SYSTEM_PROMPT = textwrap.dedent(
    """
    You are an expert Git commit message reviewer following the Conventional Commits specification.
    
    You will be given a git commit message and must review it.
    
    Respond STRICTLY in JSON format matching this schema:
    {
      "score": <float from 0.0 to 1.0>,
      "type_valid": <boolean>,
      "summary_length": <integer length of the first line>,
      "issues": ["list of strings detailing problems"],
      "suggestion": "fixed commit message here"
    }

    Metrics guidelines:
    - 0.0-0.3: Missing type, vague ("fix stuff", "WIP", "update")
    - 0.4-0.6: Has type but poor summary or wrong mood
    - 0.7-0.85: Good structure, minor issues
    - 0.86-1.0: Excellent - clear, scoped, imperative, explains why
    """
).strip()


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: str, reward: float, done: bool, error: Optional[str]) -> None:
    error_val = error if error else "null"
    done_val = str(done).lower()
    # Aligned with [STEP]  (two spaces)
    print(
        f"[STEP]  step={step} action={action_str} 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)
    # Aligned with [END]   (three spaces)
    print(f"[END]   success={str(success).lower()} steps={steps} score={score:.2f} rewards={rewards_str}", flush=True)


def get_model_action(client: OpenAI, commit_message: str) -> GitCommitReviewerAction:
    user_prompt = f"Review this commit message:\n\n{commit_message}"
    
    try:
        completion = client.chat.completions.create(
            model=MODEL_NAME,
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": user_prompt},
            ],
            temperature=TEMPERATURE,
            max_tokens=MAX_TOKENS,
            response_format={"type": "json_object"} if "gpt-" in MODEL_NAME.lower() else None,
        )
        
        text = (completion.choices[0].message.content or "").strip()
        # Parse text just to ensure we can build the Action object safely
        # Fallback parsing in case model adds weird ticks
        if text.startswith("```json"):
            text = text.replace("```json", "").replace("```", "").strip()
            
        data = json.loads(text)
        action = GitCommitReviewerAction(
            score=float(data.get("score", 0.0)),
            type_valid=bool(data.get("type_valid", False)),
            summary_length=int(data.get("summary_length", len(commit_message))),
            issues=list(data.get("issues", [])),
            suggestion=str(data.get("suggestion", commit_message))
        )
        
        return action
    except Exception:
        # Fallback action
        # Fallback action
        return GitCommitReviewerAction(
            score=0.1,
            type_valid=False,
            summary_length=0,
            issues=["Failed to generate response"],
            suggestion="fix: broken agent"
        )


async def run_task(task_name: str):
    client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY)

    # In standard V4 openenv, we connect to local server or container via EnvClient
    # The hackathon validator natively maps the port to 8000 when checking, 
    # but we can try to connect to localhost:8000 if not specified.
    endpoint = os.getenv("ENV_URL", "http://127.0.0.1:7860")
    env = EnvClient(base_url=endpoint)

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

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

    try:
        result = await env.reset(task=task_name)
        obs = result.observation
        commit_to_review = obs.get("commit_message", "")
        max_commits = obs.get("total_commits", MAX_STEPS)

        for step in range(1, max_commits + 1):
            if result.done or commit_to_review == "DONE":
                break

            action = get_model_action(client, commit_to_review)
            
            # Action serialization for logger
            action_str = f"review(score={action.score})"

            result = await env.step(action)
            obs = result.observation

            reward = result.reward or 0.0
            done = result.done
            error = None

            rewards.append(reward)
            steps_taken = step
            commit_to_review = obs.get("commit_message", "")

            log_step(step=step, action_str=action_str, reward=reward, done=done, error=error)

            if done:
                break
                
        # Calculate bound normalized score tightly for Phase 2 validation bound checks (must be in (0,1))
        # Total max reward is len(rewards) since max per commit is ~1.0
        max_possible = len(rewards) if len(rewards) > 0 else 1.0
        raw_score = sum(rewards) / max_possible
        score = 0.01 + raw_score * 0.98  # clamp between 0.01 and 0.99 strictly
        success = score >= 0.5

    except Exception:
        score = 0.0
        success = False
        log_end(success=success, steps=steps_taken, score=score, rewards=rewards)
        return

    # No env.close() in default http EnvClient but handled properly by exiting loop
    log_end(success=success, steps=steps_taken, score=score, rewards=rewards)


def main():
    # If the hackathon script tries to execute a full run testing all 3, we can just run the tasks.
    # The instructions say "produce a reproducible baseline score on all 3 tasks."
    asyncio.run(run_task("easy"))
    asyncio.run(run_task("medium"))
    asyncio.run(run_task("hard"))


if __name__ == "__main__":
    main()