import json
import subprocess
from pathlib import Path
from typing import Dict, Any, Optional
from app.data_structures import SearchResult
from app.utils import *



pyright_feedback_skeleton = """After using Pyright to test your generated code:
```
GENERATED_CODE_HERE
```

Here is the static analysis feedback :
FEEDBACK_HERE

Please revise the your generated code to correct above issues.
Remember:
1. Please include a brief explanation of the changes you made.
2. Please put your revised code in a single Python code block (surrounded by ``` and ```).
3. Only generate the function body and don't generate the line number.
4. Make sure your code follows the correct indentation!
"""


class PyrightChecker:
    def __init__(self, 
                 project_path: str,
                 target_func: SearchResult,
                 pyright_path: str = "pyright",
                 container = None,
                 ):
        """PyrightChecker is used as step 1.
        Args:
            project_path (str): project folder path.
            target_file_path (str): file abs path that contains target function.
            target_func (SearchResult): target function, including: file_path, start, end, class_name, func_name, code, 
            pyright_path (str, optional): path to pyright. Defaults to "pyright".
        """
        self.project_path = project_path
        self.target_file_path = target_func.file_path
        self.target_func = target_func
        self.pyright_path = pyright_path
        self.diagnostics = None
        self.container = container
    
    def read_line_from_file(self, file_path: str, start_num: int, end_num: int) -> list:
        """
        Read the content of a specific line from a file.

        Parameters:
        - file_path: Path to the file.
        - line_number: The line number to read (1-based).

        Returns:
        - The content of the line (with trailing newline stripped), in a list.

        Raises:
        - IndexError: If the line number is out of range.
        - FileNotFoundError: If the file doesn't exist.
        """
        return_list = []
        # make sure start_num <= end_num
        if(start_num > end_num):
            start_num, end_num = end_num, start_num
        
        if self.container is None:
            with open(file_path, 'r', encoding='utf-8') as f:
                for current_line_number, line in enumerate(f, start=1):
                    if start_num <= current_line_number <= end_num:
                        return_list.append(line.rstrip('\n'))
                    if current_line_number > end_num:
                        break
                return return_list
            raise IndexError(f"Line {start_num}-{end_num} not found in file: {file_path}")

        else:
            # Docker container mode
            # Use bash command: sed -n '{start_num},{end_num}p' <file>
            cmd = f"bash -c 'sed -n \"{start_num},{end_num}p\" {file_path}'"
            try:
                exec_result = self.container.exec_run(cmd)
                output = exec_result.output.decode("utf-8") if isinstance(exec_result.output, bytes) else str(exec_result.output)
                return_list = [line.rstrip('\n') for line in output.splitlines()]
                if not return_list:
                    raise IndexError(f"Line {start_num}-{end_num} not found in file: {file_path} inside container")
                return return_list
            except Exception as e:
                raise RuntimeError(f"Failed to read file from container: {e}")
    
    
    
    def execute_pyright(self, 
                        output_json_path: Optional[str] = None
                        ) -> Optional[Dict[str, Any]]:
        """
        Run pyright on the given file and return parsed JSON output.
        """
        if self.container is None:
            try:
                cmd = f"bash -c 'cd {self.project_path} && pyright --outputjson {self.target_file_path}'"
                result = subprocess.run(
                    ["bash", "-c", cmd],
                    capture_output=True,
                    text=True,
                    check=False,
                    )
                output = json.loads(result.stdout)
                self.diagnostics = output
                
                if output_json_path:
                    with open(output_json_path, "w") as f:
                        json.dump(output, f, indent=2)
                return output
            except subprocess.CalledProcessError as e:
                print(f"Pyright failed with error:\n{e.stderr}")
                return None
        else:
            try:
                self.container.reload()
                if self.container.status != "running":
                    print(f"[i] Container '{self.container.name}' is not running. Attempting to start...")
                    self.container.start()
                    self.container.reload()
                    if self.container.status != "running":
                        raise RuntimeError(f"Failed to start container '{self.container.name}'.")

                cmd = f"bash -c 'cd {self.project_path} && pyright --outputjson {self.target_file_path}'"
                exec_result = self.container.exec_run(cmd)
                output = exec_result.output.decode("utf-8") if isinstance(exec_result.output, bytes) else str(exec_result.output)
                output = json.loads(output)
                self.diagnostics = output
                return output
            except Exception as e:
                print(f"Running Pyright in container {self.container.name} failed with error:\n{e}")
                return None
        
        
        
    def extract_diagnostics(self, output_json_path: Optional[str] = None):
        """extract related pyright feedback of the target function
        return feedback(json format or None) + bool value(False: diagnostics is empty or contain error)
        """
        if(self.diagnostics is None):
            print(f'Origin diagnostics is still empty. run execute_pyright() first.')
            return None, False
        results = []
        start_line = self.target_func.start
        end_line = self.target_func.end
        for diag in self.diagnostics.get("generalDiagnostics", []):
            if Path(diag["file"]).name == Path(self.target_file_path).name:
                
                diag["range"]["start"]["line"] += 1
                diag["range"]["end"]["line"] += 1
                
                diag_start = diag["range"]["start"]["line"]
                diag_end = diag["range"]["end"]["line"]
                if start_line <= diag_start and diag_end <= end_line:
                    code_at_line = self.read_line_from_file(self.target_file_path, diag_start, diag_end)
                    diag['src_code_list'] = code_at_line # list
                    results.append(diag)
        if output_json_path:   
            with open(output_json_path, "w", encoding="utf-8") as f:
                json.dump(results, f, indent=2)
        return results
    
    
    
    def check_func(self, output_json_path: Optional[str] = None):
        """
        Run pyright and return filtered diagnostics related to the target function.
        Optionally save pyright's full JSON output to a file.
        """
        full_diagnostics = self.execute_pyright() # json
        extracted_diagnostics = self.extract_diagnostics(output_json_path) # list
        return extracted_diagnostics
    
    
    
    def build_pyright_feedback(self, extracted_diagnostics: list):
        """based on the diagnostics related to the target function,
        build the feedback to reprompt the LLM/Agent to refine the code
        Args:
            extracted_diagnostics (list): extracted_diagnostics related to the target function
        Return:
            Feedback: str
            Bool_val: bool, True: pass pyright; False: didn't pass pyright
        """
        count = 0
        all_feedback = ''
        if (len(extracted_diagnostics) == 0):
            return all_feedback, True
        
        for ele in extracted_diagnostics:
            count += 1
            src_code_list = ele['src_code_list']
            message = ele['message']
            src_code = '/n'.join(src_code_list)
            feedback_from_ele = f'{count}. the code:\n```{src_code}```\ncaused the issue:\n{message}\n'
            all_feedback += feedback_from_ele
        
        generated_code_no_line_num = self.read_line_from_file(self.target_file_path, self.target_func.start, self.target_func.end)
        generated_code = '\n'.join(generated_code_no_line_num)
        pyright_feedback = pyright_feedback_skeleton.replace('GENERATED_CODE_HERE', generated_code)
        pyright_feedback = pyright_feedback.replace('FEEDBACK_HERE', all_feedback)
        return pyright_feedback, False
            
            