import cv2
import numpy as np
import os
import uuid
import json
import subprocess
import base64
from math import degrees
from PIL import Image
import io
from langchain_core.tools import tool
from groq import Groq
import google.generativeai as genai

from config import (
    logger, mp_pose, pose, mp_drawing,
    persistent_vars, analysis_cache
)

GROQ_API_KEY = os.environ.get("GROQ_API_KEY", "")


# ─────────────────────────────────────────────────
# ANGLE CALCULATION
# ─────────────────────────────────────────────────
def calculate_angle(p1, p2, p3):
    try:
        a  = np.array(p1)
        b  = np.array(p2)
        c  = np.array(p3)
        ab = a - b
        bc = c - b
        cos_angle = np.dot(ab, bc) / (np.linalg.norm(ab) * np.linalg.norm(bc) + 1e-6)
        return degrees(np.arccos(np.clip(cos_angle, -1.0, 1.0)))
    except Exception as e:
        logger.error(f"Angle calc error: {e}")
        return 0.0


# ─────────────────────────────────────────────────
# EXTRACT ANGLES FROM LANDMARKS
# ─────────────────────────────────────────────────
def extract_angles_from_landmarks(landmarks, w=1, h=1):
    def pt(lm):
        return [lm.x * w, lm.y * h]
    lm = landmarks
    return {
        "left_elbow": calculate_angle(
            pt(lm[mp_pose.PoseLandmark.LEFT_SHOULDER]),
            pt(lm[mp_pose.PoseLandmark.LEFT_ELBOW]),
            pt(lm[mp_pose.PoseLandmark.LEFT_WRIST])
        ),
        "right_elbow": calculate_angle(
            pt(lm[mp_pose.PoseLandmark.RIGHT_SHOULDER]),
            pt(lm[mp_pose.PoseLandmark.RIGHT_ELBOW]),
            pt(lm[mp_pose.PoseLandmark.RIGHT_WRIST])
        ),
        "left_knee": calculate_angle(
            pt(lm[mp_pose.PoseLandmark.LEFT_HIP]),
            pt(lm[mp_pose.PoseLandmark.LEFT_KNEE]),
            pt(lm[mp_pose.PoseLandmark.LEFT_ANKLE])
        ),
        "right_knee": calculate_angle(
            pt(lm[mp_pose.PoseLandmark.RIGHT_HIP]),
            pt(lm[mp_pose.PoseLandmark.RIGHT_KNEE]),
            pt(lm[mp_pose.PoseLandmark.RIGHT_ANKLE])
        ),
        "left_hip": calculate_angle(
            pt(lm[mp_pose.PoseLandmark.LEFT_SHOULDER]),
            pt(lm[mp_pose.PoseLandmark.LEFT_HIP]),
            pt(lm[mp_pose.PoseLandmark.LEFT_KNEE])
        ),
        "right_hip": calculate_angle(
            pt(lm[mp_pose.PoseLandmark.RIGHT_SHOULDER]),
            pt(lm[mp_pose.PoseLandmark.RIGHT_HIP]),
            pt(lm[mp_pose.PoseLandmark.RIGHT_KNEE])
        ),
        "left_shoulder": calculate_angle(
            pt(lm[mp_pose.PoseLandmark.LEFT_HIP]),
            pt(lm[mp_pose.PoseLandmark.LEFT_SHOULDER]),
            pt(lm[mp_pose.PoseLandmark.LEFT_ELBOW])
        ),
        "right_shoulder": calculate_angle(
            pt(lm[mp_pose.PoseLandmark.RIGHT_HIP]),
            pt(lm[mp_pose.PoseLandmark.RIGHT_SHOULDER]),
            pt(lm[mp_pose.PoseLandmark.RIGHT_ELBOW])
        ),
        "back": calculate_angle(
            pt(lm[mp_pose.PoseLandmark.LEFT_SHOULDER]),
            pt(lm[mp_pose.PoseLandmark.LEFT_HIP]),
            pt(lm[mp_pose.PoseLandmark.LEFT_ANKLE])
        ),
    }



# ─────────────────────────────────────────────────
# EXTRACT MEDIAN ANGLES FROM VIDEO
# ─────────────────────────────────────────────────
def extract_angles_from_video(video_path: str, sample_fps: int = 2) -> dict:
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        raise RuntimeError(f"Cannot open video: {video_path}")

    fps        = cap.get(cv2.CAP_PROP_FPS) or 30
    interval   = max(1, int(fps / sample_fps))
    all_angles = {}
    frame_idx  = 0
    valid      = 0

    logger.debug(f"Extracting angles from {video_path}  fps={fps}  interval={interval}")

    while True:
        ret, frame = cap.read()
        if not ret:
            break
        if frame_idx % interval == 0:
            h, w = frame.shape[:2]
            rgb  = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            res  = pose.process(rgb)
            if res.pose_landmarks:
                angles = extract_angles_from_landmarks(res.pose_landmarks.landmark, w, h)
                for joint, angle in angles.items():
                    all_angles.setdefault(joint, []).append(angle)
                valid += 1
        frame_idx += 1

    cap.release()
    logger.debug(f"Valid frames with pose: {valid}")

    if not all_angles:
        raise RuntimeError("No pose detected in video. Check lighting/visibility.")

    return {joint: float(np.median(vals)) for joint, vals in all_angles.items()}




# ─────────────────────────────────────────────────
# DOWNLOAD YOUTUBE VIDEO
# ─────────────────────────────────────────────────
def download_youtube_video(url: str, out_path: str) -> str:
    logger.debug(f"Downloading YouTube video: {url}")
    # Hugging Face IP might be throttled; use smaller formats and strict timeouts
    format_options = ["worst", "best[height<=480]", "best[ext=mp4]"]
    last_error     = ""

    for fmt in format_options:
        cmd = [
            "yt-dlp", 
            "-f", fmt, 
            "--socket-timeout", "15", 
            "--force-ipv4",
            "--no-playlist", 
            "--no-warnings", 
            "--extractor-args", "youtube:player_client=ios,android,web",
            "-o", out_path, 
            url
        ]
        logger.debug(f"Trying yt-dlp format: {fmt}")
        try:
            result = subprocess.run(cmd, capture_output=True, text=True, timeout=25)
            if result.returncode == 0 and os.path.exists(out_path):
                logger.debug(f"Download succeeded: {fmt}")
                return out_path
            last_error = result.stderr
        except subprocess.TimeoutExpired as e:
            last_error = f"Timeout for format {fmt}: {e}"
            logger.error(last_error)

    fallback_template = out_path.replace(".mp4", ".%(ext)s")
    try:
        cmd_fallback = [
            "yt-dlp", 
            "--socket-timeout", "15", 
            "--force-ipv4",
            "--no-playlist", 
            "--no-warnings", 
            "--extractor-args", "youtube:player_client=ios,android,web",
            "-o", fallback_template, 
            url
        ]
        subprocess.run(cmd_fallback, capture_output=True, text=True, timeout=25)
    except subprocess.TimeoutExpired as e:
        logger.error(f"Fallback timeout: {e}")
        
    base     = out_path.replace(".mp4", "")
    possible = [f"{base}.{ext}" for ext in ["mp4", "webm", "mkv", "avi", "mov"]]
    for p in possible:
        if os.path.exists(p):
            if p != out_path:
                os.rename(p, out_path)
            return out_path

    raise RuntimeError(f"yt-dlp failed.\nLast error: {last_error}")


# ─────────────────────────────────────────────────
# DETECT EXERCISE — Groq Llama-4 Scout vision
# ─────────────────────────────────────────────────
def detect_exercise_from_video(video_path: str) -> str:
    cap   = cv2.VideoCapture(video_path)
    total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    if total == 0:
        cap.release()
        return "Unknown Exercise"

    sample_points = np.linspace(0, total - 1, 5, dtype=int)
    b64_frames    = []

    for idx in sample_points:
        cap.set(cv2.CAP_PROP_POS_FRAMES, int(idx))
        ret, frame = cap.read()
        if not ret:
            continue
        frame_resized = cv2.resize(frame, (480, 270))
        pil_img       = Image.fromarray(cv2.cvtColor(frame_resized, cv2.COLOR_BGR2RGB))
        buffer        = io.BytesIO()
        pil_img.save(buffer, format="JPEG", quality=75)
        b64_frames.append(base64.b64encode(buffer.getvalue()).decode("utf-8"))

    cap.release()
    if not b64_frames:
        return "Unknown Exercise"

    content = [
        {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{b64}"}}
        for b64 in b64_frames
    ]
    content.append({
        "type": "text",
        "text": (
            "These are frames from a workout video. "
            "What is the single main exercise being performed? "
            "Reply with ONLY the exercise name. No explanation. "
            "Examples: Squat, Push-up, Deadlift, Lunge, Bicep Curl, Pull-up, Plank"
        )
    })

    try:
        client   = Groq(api_key=GROQ_API_KEY)
        response = client.chat.completions.create(
            model="meta-llama/llama-4-scout-17b-16e-instruct",
            messages=[{"role": "user", "content": content}],
            max_tokens=20,
            temperature=0.1
        )
        exercise = response.choices[0].message.content.strip().strip('"').strip("'")
        logger.debug(f"Detected exercise: {exercise}")
        return exercise
    except Exception as e:
        logger.error(f"Exercise detection failed: {e}")
        return "Unknown Exercise"


# ─────────────────────────────────────────────────
# COMPARE ANGLES
# ─────────────────────────────────────────────────
def compare_angles(ref_angles: dict, user_angles: dict, threshold: float = 15.0) -> dict:
    comparison = {}
    for joint in ref_angles:
        if joint not in user_angles:
            continue
        ref_val  = ref_angles[joint]
        user_val = user_angles[joint]
        dev      = user_val - ref_val
        comparison[joint] = {
            "reference": round(ref_val,  1),
            "user"     : round(user_val, 1),
            "deviation": round(dev,      1),
            "is_error" : abs(dev) > threshold,
            "direction": "higher" if dev > 0 else "lower"
        }
    return comparison


# ─────────────────────────────────────────────────
# GROQ LLM FEEDBACK
# ─────────────────────────────────────────────────
def get_llm_feedback(exercise_name: str, comparison: dict, groq_key: str) -> str:
    errors = {j: v for j, v in comparison.items() if v["is_error"]}
    good   = {j: v for j, v in comparison.items() if not v["is_error"]}

    error_lines = "\n".join([
        f"- {j.replace('_',' ').title()}: "
        f"position is {v['direction']} than ideal"
        for j, v in errors.items()
    ])
    good_lines = "\n".join([
        f"- {j.replace('_',' ').title()}: good position"
        for j, v in good.items()
    ])

    prompt = f"""You are a real gym trainer standing right next to someone while they exercise.
Speak naturally like a coach giving instant verbal cues during a workout.
Do NOT use any numbers, degrees, angles, or technical measurements.
Do NOT use bullet points or numbered lists.
Keep it short — 2 to 4 sentences max, like you're actually talking to them mid-set.
Use simple everyday language anyone can understand.

Exercise: {exercise_name}

What they're doing well:
{good_lines if good_lines else "Nothing specific detected yet"}

What needs fixing:
{error_lines if error_lines else "Nothing — their form looks great!"}

Give your quick coaching cue now. Be encouraging but direct. Sound like a real trainer."""

    try:
        client   = Groq(api_key=groq_key)
        response = client.chat.completions.create(
            model="llama-3.1-8b-instant",
            messages=[{"role": "user", "content": prompt}],
            max_tokens=300,
            temperature=0.8
        )
        return response.choices[0].message.content.strip()
    except Exception as e:
        logger.error(f"Groq feedback error: {e}")
        return f"Feedback unavailable: {e}"


# ─────────────────────────────────────────────────
# GENERATE VOICE FEEDBACK (Groq Orpheus TTS)
# ─────────────────────────────────────────────────
def generate_voice_feedback(text: str, groq_key: str) -> str:
    """
    Converts feedback text to spoken audio using Groq Orpheus TTS.
    Returns base64-encoded WAV audio string.
    """
    try:
        client = Groq(api_key=groq_key)
        response = client.audio.speech.create(
            model="canopylabs/orpheus-v1-english",
            voice="troy",
            input=text,
            response_format="wav"
        )
        # Read the audio bytes from the response
        audio_bytes = response.read()
        audio_b64 = base64.b64encode(audio_bytes).decode("utf-8")
        logger.debug(f"TTS audio generated: {len(audio_bytes)} bytes")
        return audio_b64
    except Exception as e:
        logger.error(f"TTS generation error: {e}")
        return ""


# ─────────────────────────────────────────────────
# LIVE FRAME ANALYSIS
# Called per-frame during live camera session
# ─────────────────────────────────────────────────
def analyze_live_frame(frame_b64: str, ref_angles: dict, threshold: float = 15.0) -> dict:
    """
    Decodes a base64 JPEG frame from the browser webcam.
    Runs MediaPipe pose on it.
    Returns annotated frame (base64) + angle comparison.
    """
    try:
        # Decode base64 → numpy frame
        img_bytes = base64.b64decode(frame_b64)
        np_arr    = np.frombuffer(img_bytes, np.uint8)
        frame     = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)

        if frame is None:
            return {"error": "Could not decode frame"}

        h, w = frame.shape[:2]
        rgb  = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        res  = pose.process(rgb)

        comparison   = {}
        pose_detected = False

        if res.pose_landmarks:
            pose_detected = True

            # Draw skeleton
            mp_drawing.draw_landmarks(
                frame,
                res.pose_landmarks,
                mp_pose.POSE_CONNECTIONS,
                mp_drawing.DrawingSpec(color=(0, 255, 0), thickness=2, circle_radius=3),
                mp_drawing.DrawingSpec(color=(255, 255, 255), thickness=2)
            )

            user_angles = extract_angles_from_landmarks(res.pose_landmarks.landmark, w, h)
            comparison  = compare_angles(ref_angles, user_angles, threshold)

            # Overlay joint info (no degrees — simple status)
            y = 30
            for joint, data in comparison.items():
                color = (0, 0, 255) if data["is_error"] else (0, 255, 0)
                status = "Fix" if data["is_error"] else "OK"
                label = f"{joint.replace('_',' ').title()}: {status}"
                cv2.putText(frame, label, (10, y),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.45, color, 1)
                y += 22
        else:
            cv2.putText(frame, "No pose detected — step back or improve lighting",
                        (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.55, (0, 165, 255), 2)

        # Encode annotated frame back to base64
        _, buffer   = cv2.imencode(".jpg", frame, [cv2.IMWRITE_JPEG_QUALITY, 80])
        out_b64     = base64.b64encode(buffer).decode("utf-8")

        errors     = {j: v for j, v in comparison.items() if v["is_error"]}
        good       = {j: v for j, v in comparison.items() if not v["is_error"]}
        form_score = round((len(good) / max(len(comparison), 1)) * 100, 1) if comparison else 0

        return {
            "annotated_frame": out_b64,
            "comparison"     : comparison,
            "form_score"     : form_score,
            "pose_detected"  : pose_detected,
            "errors_count"   : len(errors),
            "correct_count"  : len(good),
        }

    except Exception as e:
        logger.error(f"analyze_live_frame error: {e}")
        return {"error": str(e)}


# ─────────────────────────────────────────────────
# ANNOTATE USER VIDEO (uploaded video branch)
# ─────────────────────────────────────────────────
def annotate_user_video(user_video_path: str,
                        ref_angles: dict,
                        exercise_name: str,
                        threshold: float = 15.0) -> str:
    cap = cv2.VideoCapture(user_video_path)
    fps = cap.get(cv2.CAP_PROP_FPS) or 30
    w   = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    h   = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

    os.makedirs(os.path.join("static", "outputs"), exist_ok=True)
    uid       = uuid.uuid4()
    raw_path  = os.path.join("static", "outputs", f"raw_{uid}.mp4")
    final_path = os.path.join("static", "outputs", f"annotated_{uid}.mp4")

    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
    writer = cv2.VideoWriter(raw_path, fourcc, fps, (w, h))
    if not writer.isOpened():
        raw_path = raw_path.replace(".mp4", ".avi")
        fourcc   = cv2.VideoWriter_fourcc(*"XVID")
        writer   = cv2.VideoWriter(raw_path, fourcc, fps, (w, h))

    if not writer.isOpened():
        cap.release()
        raise RuntimeError("Cannot open VideoWriter.")

    while True:
        ret, frame = cap.read()
        if not ret:
            break

        rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        res = pose.process(rgb)

        if res.pose_landmarks:
            mp_drawing.draw_landmarks(
                frame, res.pose_landmarks, mp_pose.POSE_CONNECTIONS,
                mp_drawing.DrawingSpec(color=(0, 255, 0), thickness=2, circle_radius=3),
                mp_drawing.DrawingSpec(color=(255, 255, 255), thickness=2)
            )
            user_angles = extract_angles_from_landmarks(res.pose_landmarks.landmark, w, h)
            comparison  = compare_angles(ref_angles, user_angles, threshold)

            y = 30
            cv2.putText(frame, f"Exercise: {exercise_name}",
                        (10, y), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 0), 2)
            y += 30
            for joint, data in comparison.items():
                color = (0, 0, 255) if data["is_error"] else (0, 255, 0)
                status = "Fix" if data["is_error"] else "OK"
                cv2.putText(frame,
                    f"{joint.replace('_',' ').title()}: {status}",
                    (10, y), cv2.FONT_HERSHEY_SIMPLEX, 0.45, color, 1)
                y += 22

        writer.write(frame)

    cap.release()
    writer.release()
    logger.debug(f"Raw annotated video: {raw_path} ({os.path.getsize(raw_path)} bytes)")

    # Re-encode with ffmpeg for browser compatibility
    try:
        check = subprocess.run(["ffmpeg", "-version"], capture_output=True, text=True)
        if check.returncode == 0:
            cmd = [
                "ffmpeg", "-y", "-i", raw_path,
                "-vcodec", "libx264", "-acodec", "aac",
                "-pix_fmt", "yuv420p",
                "-movflags", "+faststart",
                "-preset", "fast",
                final_path
            ]
            result = subprocess.run(cmd, capture_output=True, text=True)
            if result.returncode == 0 and os.path.exists(final_path):
                try:
                    os.remove(raw_path)
                except Exception:
                    pass
                logger.debug(f"ffmpeg re-encode: {final_path}")
                return final_path
    except FileNotFoundError:
        pass

    logger.warning("ffmpeg not found — returning raw video")
    return raw_path


# ─────────────────────────────────────────────────
# MAIN TOOL: Full video analysis
# ─────────────────────────────────────────────────
@tool
def fitness_analysis_tool(youtube_url: str,
                          user_video_path: str,
                          groq_api_key: str) -> str:
    """Full fitness coach pipeline for uploaded video."""
    try:
        yt_path = os.path.join("static", "uploads", f"ref_{uuid.uuid4()}.mp4")
        os.makedirs(os.path.dirname(yt_path), exist_ok=True)

        logger.debug("Step 1: Downloading YouTube reference video...")
        download_youtube_video(youtube_url, yt_path)

        logger.debug("Step 2: Detecting exercise...")
        exercise_name = detect_exercise_from_video(yt_path)

        logger.debug("Step 3: Extracting reference angles...")
        ref_angles = extract_angles_from_video(yt_path, sample_fps=2)

        logger.debug("Step 4: Extracting user angles...")
        user_angles = extract_angles_from_video(user_video_path, sample_fps=2)

        logger.debug("Step 5: Comparing angles...")
        comparison = compare_angles(ref_angles, user_angles)

        logger.debug("Step 6: Generating feedback...")
        feedback = get_llm_feedback(exercise_name, comparison, groq_api_key)

        logger.debug("Step 7: Annotating video...")
        annotated_path = annotate_user_video(user_video_path, ref_angles, exercise_name)

        errors     = {j: v for j, v in comparison.items() if v["is_error"]}
        good       = {j: v for j, v in comparison.items() if not v["is_error"]}
        form_score = round((len(good) / max(len(comparison), 1)) * 100, 1)

        result = {
            "exercise_name"   : exercise_name,
            "form_score"      : form_score,
            "reference_angles": ref_angles,
            "user_angles"     : user_angles,
            "comparison"      : comparison,
            "errors_count"    : len(errors),
            "correct_count"   : len(good),
            "feedback"        : feedback,
            "annotated_video" : annotated_path,
        }
        analysis_cache.update(result)
        return json.dumps(result, indent=2)

    except Exception as e:
        logger.error(f"fitness_analysis_tool error: {e}")
        return json.dumps({"error": str(e)})

# ─────────────────────────────────────────────────
# YOUTUBE SEARCH TOOL
# ─────────────────────────────────────────────────
@tool
def search_youtube_tool(query: str, max_results: int = 4) -> str:
    """
    Searches YouTube for videos matching the query using yt-dlp.
    Use this to find specific workout or yoga videos for users
    based on their category or weight preferences.
    """
    logger.debug(f"Searching YouTube for: {query}")
    try:
        cmd = ["yt-dlp", f"ytsearch{max_results}:{query}", "--dump-json", "--flat-playlist", "--no-warnings"]
        result = subprocess.run(cmd, capture_output=True, text=True)
        videos = []
        if result.returncode == 0:
            for line in result.stdout.strip().split('\n'):
                if not line: continue
                try:
                    data = json.loads(line)
                    videos.append({
                        "title": data.get("title"),
                        "url": data.get("url"),
                        "id": data.get("id"),
                        "duration": data.get("duration")
                    })
                except Exception:
                    pass
        return json.dumps(videos, indent=2)
    except Exception as e:
        logger.error(f"search_youtube_tool error: {e}")
        return json.dumps({"error": str(e)})