import regex as re
from typing import Dict, List, Tuple
from langchain_core.messages import HumanMessage
from prompts import create_hypotheses_generation_prompt

from typing import Tuple


def _content_to_text(content) -> str:
    """
    Normalize LLM content payloads (string/list/dict) into plain text.
    Anthropic and Gemini models may return a list of content blocks instead of
    a plain string (same normalization as llm_content_to_text in main.py).
    """
    if content is None:
        return ""
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        parts = []
        for item in content:
            text = _content_to_text(item)
            if text:
                parts.append(text)
        return "\n".join(parts)
    if isinstance(content, dict):
        for key in ("text", "content", "output_text"):
            if key in content:
                return _content_to_text(content.get(key))
        return str(content)
    return str(content)


PATTERN = re.compile(
    r"PAPER:\s*(?P<paper>.*?)\s*[\r\n]+SYSTEM_CHART:\s*(?P<chart>.*?)(?=\s*={5,}|$)",
    re.DOTALL
)

HYPOTHESIS_HEADING = re.compile(r"^(?P<lead>[ \t#*]*Hypothesis\s*#?\s*)\d+", re.MULTILINE | re.IGNORECASE)


def number_hypotheses(iterations: List[Dict]) -> Tuple[List[str], List[Dict]]:
    """
    Renumber "Hypothesis #k" headings continuously across iterations.
    Returns the renumbered iteration texts and a flat list of
    {number, iteration, text}. An iteration with no parseable heading
    becomes a single hypothesis.
    """
    counter = 0

    def bump(match):
        nonlocal counter
        counter += 1
        return f"{match.group('lead')}{counter}"

    texts, hypotheses = [], []
    for it in iterations:
        text = HYPOTHESIS_HEADING.sub(bump, it['content'])
        starts = [m.start() for m in HYPOTHESIS_HEADING.finditer(text)]
        if not starts:
            counter += 1
            text = f"Hypothesis #{counter}:\n{text.strip()}"
            starts = [0]
        texts.append(text)
        first = counter - len(starts) + 1
        for i, (start, end) in enumerate(zip(starts, starts[1:] + [len(text)])):
            hypotheses.append({'number': first + i, 'iteration': it['iteration'], 'text': text[start:end].strip()})
    return texts, hypotheses

def parse_multiple_system_chart_files(content_list: List[str]) -> Tuple[str, str]:
    """
    Parse multiple system chart files and extract numbered tables and citations.
    Handles new format with 100 '=' delimiters and PAPER:/SYSTEM_CHART: structure.
    
    Args:
        content_list: List of raw content strings from multiple system chart files
        
    Returns:
        Tuple of (numbered_tables, numbered_citations)
    """
    papers: List[str] = []
    charts: List[str] = []

    for blob in content_list:
        if not blob or not blob.strip():
            continue
        for m in PATTERN.finditer(blob):
            papers.append(m.group("paper").strip())
            charts.append(m.group("chart").strip())

    numbered_tables = []
    numbered_citations = []
    for i, (p, c) in enumerate(zip(papers, charts), 1):
        numbered_tables.append(f"\n=== TABLE {i} ===\n{c}\n=== END OF TABLE {i} ===")
        numbered_citations.append(f"{i}. {p}")
        
    return "\n".join(numbered_tables).strip(), "\n".join(numbered_citations).strip()

def parse_system_charts_and_citations(content: str) -> Tuple[str, str]:
    """
    Parse system charts content and extract numbered tables and citations robustly.
    Updated to handle new format with 100 '=' delimiters and PAPER:/SYSTEM_CHART: structure.
    
    Args:
        content: Raw content from system_charts.txt file
        
    Returns:
        Tuple of (numbered_tables, numbered_citations)
    """
    # Use the new multi-file parser with single content
    return parse_multiple_system_chart_files([content])

async def generate_hypotheses_from_system_charts(
    system_charts: List[Dict[str, str]], 
    llm, 
    material_goal_1: str = "Property #1 for xxx material",
    material_goal_2: str = "Property #2 for xxx material",
    iterations_count: int = 1,
    progress_callback=None
) -> Tuple[str, List[Dict]]:
    """
    Generate scientific hypotheses using the new synergistic approach.
    
    Args:
        system_charts: List of parsed system chart dictionaries with 'content' or 'raw_content' fields
        llm: Language model for hypothesis generation
        material_goal_1: First material design goal to guide hypothesis generation
        material_goal_2: Second material design goal to guide hypothesis generation
        iterations_count: Number of hypothesis iterations to generate
        progress_callback: Function to report progress updates
        
    Returns:
        (hypotheses_content, hypotheses) where hypotheses is a list of
        {number, iteration, text} numbered continuously across iterations
    """
    
    if not system_charts:
        return "No system charts provided for hypothesis generation.", []
    
    # Initial progress update
    if progress_callback:
        progress_callback({
            "status": "parsing",
            "progress": 5,
            "message": "Parsing system charts and citations..."
        })
    
    # Extract raw content from system charts
    # Support both 'raw_content' and 'content' fields for backward compatibility
    content_list = []
    for chart in system_charts:
        raw_content = chart.get('raw_content', chart.get('content', ''))
        if raw_content and raw_content.strip():
            content_list.append(raw_content)
    
    if not content_list:
        return "No valid system chart content found for hypothesis generation.", []
    
    # Parse all system chart files together using the new multi-file parser
    numbered_tables, numbered_citations = parse_multiple_system_chart_files(content_list)
    
    if progress_callback:
        progress_callback({
            "status": "generating",
            "progress": 20,
            "message": "Generating synergistic hypotheses..."
        })
    
    # Create the prompt with user-specified material goals
    hypotheses_prompt = create_hypotheses_generation_prompt(material_goal_1, material_goal_2)
    
    # Prepare the prompt in the specified format
    full_prompt = f"Tables are: {numbered_tables}\n\nCitation resources are: {numbered_citations}\n\n{hypotheses_prompt}"
    
    # Configure LLM parameters for creativity
    llm_kwargs = {}
    
    # Check if this is a reasoning model (o1 series) which doesn't support temperature
    model_name = getattr(llm, 'model_name', getattr(llm, 'model', ''))
    if not (model_name and model_name.lower().startswith('o')):
        # Set higher temperature for creativity (non-reasoning models)
        if hasattr(llm, 'temperature'):
            original_temp = llm.temperature
            llm.temperature = 1.0  # Higher temperature for creativity
        elif hasattr(llm, 'model_kwargs'):
            llm_kwargs = llm.model_kwargs.copy()
            llm_kwargs['temperature'] = 1.0
    
    try:
        # Generate hypotheses for each iteration
        all_iterations_content = []
        
        for iteration in range(iterations_count):
            # Update progress for current iteration
            if progress_callback:
                progress_start = 20 + (iteration * 60 // iterations_count)
                progress_callback({
                    "status": "llm_processing",
                    "progress": progress_start,
                    "message": f"LLM generating hypotheses (iteration {iteration + 1}/{iterations_count})..."
                })
            
            conversation_history = [HumanMessage(content=full_prompt)]
            
            # Generate with appropriate parameters
            if llm_kwargs:
                # Use model_kwargs if available
                original_kwargs = getattr(llm, 'model_kwargs', {})
                llm.model_kwargs = llm_kwargs
                response = await llm.ainvoke(conversation_history)
                llm.model_kwargs = original_kwargs
            else:
                response = await llm.ainvoke(conversation_history)
            
            iteration_content = _content_to_text(response.content)
            all_iterations_content.append({
                'iteration': iteration + 1,
                'content': iteration_content
            })
        
        if progress_callback:
            progress_callback({
                "status": "formatting",
                "progress": 80,
                "message": "Formatting results..."
            })
        
        # Number hypotheses continuously across iterations, then format
        renumbered_texts, hypotheses = number_hypotheses(all_iterations_content)
        iterations_output = ""
        for iteration_data, text in zip(all_iterations_content, renumbered_texts):
            iterations_output += f"\n\n"
            iterations_output += "=" * 80 + "\n"
            iterations_output += f"ITERATION {iteration_data['iteration']} HYPOTHESES:\n"
            iterations_output += "=" * 80 + "\n\n"
            iterations_output += text + "\n\n"
            iterations_output += "=" * 80 + "\n"
        
        # Create formatted output
        formatted_output = ""
        formatted_output += "NUMBERED SYSTEM CHARTS:\n"
        formatted_output += numbered_tables + "\n\n"
        formatted_output += "CITATION REFERENCES:\n"
        formatted_output += numbered_citations
        formatted_output += iterations_output
        
        # Restore original temperature if modified
        if hasattr(llm, 'temperature') and 'original_temp' in locals():
            llm.temperature = original_temp
            
        if progress_callback:
            progress_callback({
                "status": "completed",
                "progress": 100,
                "message": f"Generated {iterations_count} iteration{'s' if iterations_count > 1 else ''} of synergistic hypotheses from {len(system_charts)} charts"
            })
        
        return formatted_output, hypotheses
        
    except Exception as e:
        # Restore original temperature if modified
        if hasattr(llm, 'temperature') and 'original_temp' in locals():
            llm.temperature = original_temp
            
        error_msg = f"Error generating hypotheses: {str(e)}"
        return error_msg, []