src/chart_analyzer.py

"""
Core chart analysis engine — sends chart screenshots to AI vision models
and parses structured trading signals from the response.
"""

import json
import os
import re
from typing import Optional

from dotenv import load_dotenv

from .prompt_templates import (
    CHART_ANALYSIS_SYSTEM_PROMPT,
    CHART_ANALYSIS_USER_PROMPT,
    MULTI_SCENARIO_USER_PROMPT,
    QUICK_ANALYSIS_PROMPT,
    STAGE1_OCR_EXTRACTION_PROMPT,
    STAGE2_REASONING_PROMPT,
    STAGE2_REASONING_WITH_FUNDAMENTALS_PROMPT,
    MARKET_CONTEXT,
)
from .utils import image_to_base64, prepare_image_for_api

# Load .env file from project root
load_dotenv()


class ChartAnalyzer:
    """Analyzes stock chart screenshots using AI vision models."""

    def __init__(
        self,
        provider: str = "google",
        model: Optional[str] = None,
        api_key: Optional[str] = None,
    ):
        self.provider = provider.lower()
        self.model = model
        self.api_key = api_key
        self._client = None

        # Set defaults per provider
        if self.provider == "anthropic":
            self.model = self.model or "claude-sonnet-4-5-20250929"
            self.api_key = self.api_key or os.environ.get("ANTHROPIC_API_KEY")
        elif self.provider == "openai":
            self.model = self.model or "gpt-4o"
            self.api_key = self.api_key or os.environ.get("OPENAI_API_KEY")
        elif self.provider == "google":
            self.model = self.model or "gemini-2.5-flash"
            self.api_key = self.api_key or os.environ.get("GOOGLE_API_KEY")
        elif self.provider == "huggingface":
            self.model = self.model or "Qwen/Qwen2.5-VL-72B-Instruct"
            self.api_key = self.api_key or os.environ.get("HF_TOKEN")
        elif self.provider == "openrouter":
            self.model = self.model or "google/gemma-4-31b-it:free"
            self.api_key = self.api_key or os.environ.get("OPENROUTER_API_KEY")

    def _get_anthropic_client(self):
        """Lazy-load Anthropic client."""
        if self._client is None:
            import anthropic
            self._client = anthropic.Anthropic(api_key=self.api_key)
        return self._client

    def _get_openai_client(self):
        """Lazy-load OpenAI client."""
        if self._client is None:
            import openai
            self._client = openai.OpenAI(api_key=self.api_key)
        return self._client

    def _get_google_client(self):
        """Lazy-load Google GenAI client."""
        if self._client is None:
            from google import genai
            self._client = genai.Client(api_key=self.api_key)
        return self._client

    def _get_huggingface_client(self):
        """Lazy-load HuggingFace client (OpenAI-compatible)."""
        if self._client is None:
            import openai
            self._client = openai.OpenAI(
                base_url="https://router.huggingface.co/v1",
                api_key=self.api_key,
            )
        return self._client

    def _get_openrouter_client(self):
        """Lazy-load OpenRouter client (OpenAI-compatible)."""
        if self._client is None:
            import openai
            self._client = openai.OpenAI(
                base_url="https://openrouter.ai/api/v1",
                api_key=self.api_key,
            )
        return self._client

    def analyze(
        self,
        image_bytes: bytes,
        user_context: str = "",
        quick_mode: bool = False,
        multi_scenario: bool = False,
        market_type: str = "stocks",
    ) -> dict:
        """Analyze a chart screenshot and return structured trading signal.

        Args:
            image_bytes: Raw image bytes of the chart screenshot.
            user_context: Optional additional context from the user.
            quick_mode: If True, return a brief verdict instead of full analysis.
            multi_scenario: If True, use multi-scenario prompt (bullish/bearish/sideways).
            market_type: Market type for context (crypto/stocks/forex/gold/indices/energy).

        Returns:
            Parsed analysis dict with signal, prices, patterns, etc.
        """
        # Prepare image
        processed_bytes, media_type = prepare_image_for_api(image_bytes)
        b64_image = image_to_base64(processed_bytes)

        # Build prompt based on mode
        if quick_mode:
            prompt = QUICK_ANALYSIS_PROMPT
        elif multi_scenario:
            prompt = MULTI_SCENARIO_USER_PROMPT
        else:
            prompt = CHART_ANALYSIS_USER_PROMPT

        if user_context:
            prompt += f"\n\n**Additional context from user:** {user_context}"

        # Add market context
        market_ctx = MARKET_CONTEXT.get(market_type, "")
        if market_ctx:
            prompt += f"\n\n**Market context:** {market_ctx}"

        # Call the appropriate provider
        if self.provider == "anthropic":
            raw_response = self._call_anthropic(b64_image, media_type, prompt)
        elif self.provider == "openai":
            raw_response = self._call_openai(b64_image, media_type, prompt)
        elif self.provider == "google":
            raw_response = self._call_google(processed_bytes, media_type, prompt)
        elif self.provider == "huggingface":
            raw_response = self._call_huggingface(b64_image, media_type, prompt)
        elif self.provider == "openrouter":
            raw_response = self._call_openrouter(b64_image, media_type, prompt)
        else:
            raise ValueError(f"Unsupported provider: {self.provider}")

        # Parse response
        if quick_mode:
            return {"raw_analysis": raw_response, "mode": "quick"}

        return self._parse_analysis(raw_response)

    def _call_anthropic(self, b64_image: str, media_type: str, prompt: str) -> str:
        """Call Anthropic Claude vision API."""
        client = self._get_anthropic_client()

        response = client.messages.create(
            model=self.model,
            max_tokens=4096,
            system=CHART_ANALYSIS_SYSTEM_PROMPT,
            messages=[
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "image",
                            "source": {
                                "type": "base64",
                                "media_type": media_type,
                                "data": b64_image,
                            },
                        },
                        {
                            "type": "text",
                            "text": prompt,
                        },
                    ],
                }
            ],
        )

        return response.content[0].text

    def _call_openai(self, b64_image: str, media_type: str, prompt: str) -> str:
        """Call OpenAI GPT-4o vision API."""
        client = self._get_openai_client()

        response = client.chat.completions.create(
            model=self.model,
            max_tokens=4096,
            messages=[
                {"role": "system", "content": CHART_ANALYSIS_SYSTEM_PROMPT},
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:{media_type};base64,{b64_image}",
                                "detail": "high",
                            },
                        },
                        {
                            "type": "text",
                            "text": prompt,
                        },
                    ],
                },
            ],
        )

        return response.choices[0].message.content

    def _call_google(self, image_bytes: bytes, media_type: str, prompt: str) -> str:
        """Call Google Gemini vision API with retry on rate limit."""
        import time
        from google.genai import types

        client = self._get_google_client()

        # Build the full prompt combining system + user instructions
        full_prompt = f"{CHART_ANALYSIS_SYSTEM_PROMPT}\n\n{prompt}"

        # Create image part
        image_part = types.Part.from_bytes(data=image_bytes, mime_type=media_type)

        # Retry up to 3 times with exponential backoff for rate limits
        max_retries = 3
        for attempt in range(max_retries):
            try:
                response = client.models.generate_content(
                    model=self.model,
                    contents=[image_part, full_prompt],
                )
                return response.text
            except Exception as e:
                error_str = str(e)
                if "429" in error_str or "RESOURCE_EXHAUSTED" in error_str:
                    if attempt < max_retries - 1:
                        wait_time = (attempt + 1) * 20  # 20s, 40s, 60s
                        time.sleep(wait_time)
                        continue
                raise

    def _call_huggingface(self, b64_image: str, media_type: str, prompt: str) -> str:
        """Call HuggingFace Inference API (OpenAI-compatible)."""
        client = self._get_huggingface_client()

        full_prompt = f"{CHART_ANALYSIS_SYSTEM_PROMPT}\n\n{prompt}"

        response = client.chat.completions.create(
            model=self.model,
            max_tokens=4096,
            messages=[
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:{media_type};base64,{b64_image}",
                            },
                        },
                        {
                            "type": "text",
                            "text": full_prompt,
                        },
                    ],
                },
            ],
        )

        return response.choices[0].message.content

    def _call_openrouter(self, b64_image: str, media_type: str, prompt: str) -> str:
        """Call OpenRouter API (OpenAI-compatible, free models available)."""
        client = self._get_openrouter_client()

        response = client.chat.completions.create(
            model=self.model,
            max_tokens=4096,
            messages=[
                {"role": "system", "content": CHART_ANALYSIS_SYSTEM_PROMPT},
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:{media_type};base64,{b64_image}",
                                "detail": "high",
                            },
                        },
                        {
                            "type": "text",
                            "text": prompt,
                        },
                    ],
                },
            ],
            extra_headers={
                "HTTP-Referer": "https://github.com/openclaw-stock-analyst",
                "X-Title": "OpenClaw Stock Chart Analyst",
            },
        )

        return response.choices[0].message.content

    def analyze_two_stage(
        self,
        image_bytes: bytes,
        user_context: str = "",
        ocr_model: str = "google/gemma-4-31B-it",
        ocr_provider: str = "huggingface",
        reasoning_model: str = "Qwen/Qwen2.5-VL-72B-Instruct",
        reasoning_provider: str = "huggingface",
        progress_callback=None,
        multi_scenario: bool = False,
        market_type: str = "stocks",
    ) -> dict:
        """Two-stage pipeline: OCR extraction → AI reasoning.

        Stage 1: Extract all numbers, text, indicators from chart
        Stage 2: Reason about extracted data, generate trade signal

        Supports mixing providers — e.g., Google for OCR, OpenRouter for reasoning.
        """
        processed_bytes, media_type = prepare_image_for_api(image_bytes)
        b64_image = image_to_base64(processed_bytes)

        # ── Stage 1: OCR Extraction ──
        if progress_callback:
            progress_callback(f"🔍 Stage 1: Extracting data ({ocr_provider}/{ocr_model.split('/')[-1]})...")

        extracted_data = self._call_stage(
            provider=ocr_provider,
            model=ocr_model,
            image_bytes=processed_bytes,
            b64_image=b64_image,
            media_type=media_type,
            prompt=STAGE1_OCR_EXTRACTION_PROMPT,
            max_tokens=3000,
        )

        # ── Stage 2: Reasoning ──
        if progress_callback:
            progress_callback(f"🧠 Stage 2: Reasoning ({reasoning_provider}/{reasoning_model.split('/')[-1]})...")

        # Build user context with market type
        enriched_context = user_context or "No additional context provided."
        market_ctx = MARKET_CONTEXT.get(market_type, "")
        if market_ctx:
            enriched_context += f"\n\nMarket context: {market_ctx}"

        stage2_prompt = STAGE2_REASONING_PROMPT.format(
            extracted_data=extracted_data,
            user_context=enriched_context,
        )

        raw_response = self._call_stage(
            provider=reasoning_provider,
            model=reasoning_model,
            image_bytes=processed_bytes,
            b64_image=b64_image,
            media_type=media_type,
            prompt=stage2_prompt,
            max_tokens=4096,
        )

        # Parse and enrich with pipeline metadata
        analysis = self._parse_analysis(raw_response)
        analysis["_pipeline"] = "two-stage"
        analysis["_stage1_model"] = f"{ocr_provider}/{ocr_model}"
        analysis["_stage2_model"] = f"{reasoning_provider}/{reasoning_model}"
        analysis["_stage1_extraction"] = extracted_data

        return analysis

    def analyze_three_stage(
        self,
        image_bytes: bytes,
        ticker: str = "",
        user_context: str = "",
        ocr_model: str = "google/gemma-4-31B-it",
        ocr_provider: str = "huggingface",
        reasoning_model: str = "Qwen/Qwen2.5-VL-72B-Instruct",
        reasoning_provider: str = "huggingface",
        progress_callback=None,
        multi_scenario: bool = False,
        market_type: str = "stocks",
    ) -> dict:
        """Three-stage pipeline: Fundamentals → Chart OCR → AI Reasoning.

        Stage 0 (yfinance): Fetch news, insider activity, financials
        Stage 1 (Vision model): Extract chart data from screenshot
        Stage 2 (Reasoning model): Combine all data, generate signal
        """
        from .fundamentals import fetch_stock_fundamentals, format_fundamentals_for_prompt

        processed_bytes, media_type = prepare_image_for_api(image_bytes)
        b64_image = image_to_base64(processed_bytes)

        # ── Stage 0: Fetch Fundamentals ──
        fundamental_text = ""
        fundamental_data = {}
        if ticker:
            if progress_callback:
                progress_callback(f"📰 Stage 0: Fetching news & fundamentals for {ticker}...")
            try:
                fundamental_data = fetch_stock_fundamentals(ticker)
                fundamental_text = format_fundamentals_for_prompt(fundamental_data)
            except Exception as e:
                fundamental_text = f"(Failed to fetch fundamentals: {e})"
        else:
            fundamental_text = "(No ticker provided — skipping fundamental analysis)"

        # ── Stage 1: Chart OCR Extraction ──
        if progress_callback:
            progress_callback(f"🔍 Stage 1: Extracting data ({ocr_provider}/{ocr_model.split('/')[-1]})...")

        extracted_data = self._call_stage(
            provider=ocr_provider,
            model=ocr_model,
            image_bytes=processed_bytes,
            b64_image=b64_image,
            media_type=media_type,
            prompt=STAGE1_OCR_EXTRACTION_PROMPT,
            max_tokens=3000,
        )

        # ── Stage 2: Reasoning with Fundamentals ──
        if progress_callback:
            progress_callback(f"🧠 Stage 2: Reasoning ({reasoning_provider}/{reasoning_model.split('/')[-1]})...")

        # Build enriched context with market type
        enriched_context = user_context or "No additional context provided."
        market_ctx = MARKET_CONTEXT.get(market_type, "")
        if market_ctx:
            enriched_context += f"\n\nMarket context: {market_ctx}"

        stage2_prompt = STAGE2_REASONING_WITH_FUNDAMENTALS_PROMPT.format(
            extracted_data=extracted_data,
            fundamental_data=fundamental_text,
            user_context=enriched_context,
        )

        raw_response = self._call_stage(
            provider=reasoning_provider,
            model=reasoning_model,
            image_bytes=processed_bytes,
            b64_image=b64_image,
            media_type=media_type,
            prompt=stage2_prompt,
            max_tokens=4096,
        )

        # Parse and enrich
        analysis = self._parse_analysis(raw_response)
        analysis["_pipeline"] = "three-stage"
        analysis["_stage1_model"] = f"{ocr_provider}/{ocr_model}"
        analysis["_stage2_model"] = f"{reasoning_provider}/{reasoning_model}"
        analysis["_stage1_extraction"] = extracted_data
        analysis["_fundamental_data"] = fundamental_data
        analysis["_fundamental_text"] = fundamental_text

        return analysis

    def _call_stage(
        self,
        provider: str,
        model: str,
        image_bytes: bytes,
        b64_image: str,
        media_type: str,
        prompt: str,
        max_tokens: int = 4096,
    ) -> str:
        """Call any provider for a pipeline stage. Returns raw text response.

        This is the universal dispatcher for two/three-stage pipelines,
        allowing each stage to use a different provider and model.
        """
        import openai as _openai

        if provider == "google":
            import time
            from google import genai
            from google.genai import types

            api_key = self.api_key or os.environ.get("GOOGLE_API_KEY")
            client = genai.Client(api_key=api_key)
            image_part = types.Part.from_bytes(data=image_bytes, mime_type=media_type)

            max_retries = 3
            for attempt in range(max_retries):
                try:
                    response = client.models.generate_content(
                        model=model,
                        contents=[image_part, prompt],
                    )
                    return response.text
                except Exception as e:
                    error_str = str(e)
                    if "429" in error_str or "RESOURCE_EXHAUSTED" in error_str:
                        if attempt < max_retries - 1:
                            wait_time = (attempt + 1) * 20
                            time.sleep(wait_time)
                            continue
                    raise

        elif provider == "openrouter":
            api_key = self.api_key or os.environ.get("OPENROUTER_API_KEY")
            client = _openai.OpenAI(
                base_url="https://openrouter.ai/api/v1",
                api_key=api_key,
            )
            response = client.chat.completions.create(
                model=model,
                max_tokens=max_tokens,
                messages=[
                    {
                        "role": "user",
                        "content": [
                            {
                                "type": "image_url",
                                "image_url": {
                                    "url": f"data:{media_type};base64,{b64_image}",
                                    "detail": "high",
                                },
                            },
                            {"type": "text", "text": prompt},
                        ],
                    },
                ],
                extra_headers={
                    "HTTP-Referer": "https://github.com/openclaw-stock-analyst",
                    "X-Title": "OpenClaw Stock Chart Analyst",
                },
            )
            return response.choices[0].message.content

        elif provider == "huggingface":
            api_key = self.api_key or os.environ.get("HF_TOKEN")
            client = _openai.OpenAI(
                base_url="https://router.huggingface.co/v1",
                api_key=api_key,
            )
            response = client.chat.completions.create(
                model=model,
                max_tokens=max_tokens,
                messages=[
                    {
                        "role": "user",
                        "content": [
                            {
                                "type": "image_url",
                                "image_url": {
                                    "url": f"data:{media_type};base64,{b64_image}",
                                },
                            },
                            {"type": "text", "text": prompt},
                        ],
                    },
                ],
            )
            return response.choices[0].message.content

        else:
            raise ValueError(f"Unsupported provider for pipeline stage: {provider}")

    def _parse_analysis(self, raw_response: str) -> dict:
        """Parse the AI response into a structured analysis dict.

        Handles both clean JSON and JSON embedded in markdown code blocks.
        """
        # Try to extract JSON from the response
        json_match = re.search(r"```(?:json)?\s*\n?(.*?)\n?```", raw_response, re.DOTALL)

        if json_match:
            json_str = json_match.group(1).strip()
        else:
            # Try parsing the entire response as JSON
            json_str = raw_response.strip()

        try:
            analysis = json.loads(json_str)
            analysis["_raw_response"] = raw_response
            analysis["_parse_success"] = True
            return analysis
        except json.JSONDecodeError:
            # Return raw response with error flag
            return {
                "_raw_response": raw_response,
                "_parse_success": False,
                "_parse_error": "Could not parse JSON from AI response",
                "signal": {
                    "action": "HOLD",
                    "confidence": 0.0,
                    "entry_price": None,
                    "stop_loss": None,
                    "target_1": None,
                },
                "reasoning": raw_response,
            }